Allow setting dataop choices with optuna

[jeromedockes]

This PR enables tuning the choices in a DataOp with optuna, either by using the usual make_randomized_search but with backend='optuna', or by passing an optuna Trial to .skb.make_learner in which case the learner is created with params chosen by the Trial

Here is a complete example with make_randomized_search:

from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
import skrub

pred = skrub.X().skb.apply(
    LogisticRegression(C=skrub.choose_float(0.01, 100, log=True, name="C")),
    y=skrub.y(),
)

env = {}
env["X"], env["y"] = make_classification()

search = pred.skb.make_randomized_search(backend="optuna").fit(env)
print(search.results_)

out:

Running optuna search for study skrub_randomized_search_62ea3375-6f40-4521-b7bc-3a8f0bfece82 in storage /var/folders/sc/zl95l3fs6lj3hngh0v3l7_v40000gn/T/tmpti1in6li_skrub_optuna_search_storage/optuna_storage
[I 2025-12-07 14:16:58,397] A new study created in Journal with name: skrub_randomized_search_62ea3375-6f40-4521-b7bc-3a8f0bfece82
[I 2025-12-07 14:16:58,445] Trial 0 finished with value: 0.8400000000000001 and parameters: {'0:C': 0.6683925536308476}. Best is trial 0 with value: 0.8400000000000001.
[...]
[I 2025-12-07 14:16:58,807] Trial 8 finished with value: 0.8300000000000001 and parameters: {'0:C': 0.2927934809630964}. Best is trial 0 with value: 0.8400000000000001.
[I 2025-12-07 14:16:58,852] Trial 9 finished with value: 0.86 and parameters: {'0:C': 0.0144131844705633}. Best is trial 9 with value: 0.86.
[I 2025-12-07 14:16:58,852] A new study created in memory with name: skrub_randomized_search_62ea3375-6f40-4521-b7bc-3a8f0bfece82
           C  mean_test_score
0   0.014413             0.86
1   0.668393             0.84
2   0.027854             0.84
3   0.495693             0.84
4   0.292793             0.83
5  15.087935             0.80
6  69.473385             0.78
7  49.885865             0.78
8  44.266861             0.78
9  88.145711             0.78

And here is a more advanced example showing direct use of optuna for more fine-grained control:

import optuna
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
import skrub

pred = skrub.X().skb.apply(
    LogisticRegression(C=skrub.choose_float(0.01, 100, log=True, name="C")),
    y=skrub.y(),
)

env = {}
env["X"], env["y"] = make_classification()

def objective(trial):
    learner = pred.skb.make_learner(choose=trial)
    return skrub.cross_validate(learner, environment=env)["test_score"].mean()

study = optuna.create_study(direction="maximize")

study.optimize(objective, n_trials=15)
print("search results:")
print(
    study.trials_dataframe()
    .sort_values("value", ascending=False)
    .filter(regex="(value|params.*)")
)

best_learner = pred.skb.make_learner(choose=study.best_trial)
print(f"best params:\n{best_learner.describe_params()}")

out:

[...]
[I 2025-10-12 08:22:49,542] Trial 14 finished with value: 0.8800000000000001 and parameters: {'0:C': 1.4194684155655257}. Best is trial 0 with value: 0.9199999999999999.
search results:
    value  params_0:C
0    0.92    0.018004
1    0.92    0.055297
5    0.92    0.039228
9    0.92    0.048863
11   0.92    0.077601
10   0.91    0.011382
12   0.91    0.010271
2    0.89    0.294633
3    0.89    0.467419
13   0.89    0.174381
4    0.88   22.024237
14   0.88    1.419468
6    0.87   59.865543
7    0.87    5.214283
8    0.87    5.283302
best params:
{'C': 0.018004223638802812}

parallelization

Optuna's n_jobs is just for thread-based parallelization. Here make_randomized_search() will use joblib for parallelization (and create an on-disk storage for the optuna study to support multiprocessing), unless n_jobs=1 or the joblib backend is threading, in which case optuna's n_jobs is used. this means skrub users get multiprocessing optuna search out of the box, which requires a complex manual setup for other optuna users

original, outdated description

setting skrub.DataOp choices with Optuna

Apparently @rcap107 was asked a few times about combining skrub dataops and optuna, which is a great idea.

This is very rough first draft of what this could look like. Still missing some functionality, tests & docs etc.

search estimator interface

The DataOp now has a .skb.make_optuna_search() method, which can be used to search for the best choice outcomes with optuna.

For example:

import skrub
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.dummy import DummyClassifier

X_a, y_a = make_classification(random_state=0)
X, y = skrub.X(X_a), skrub.y(y_a)
logistic = LogisticRegression(C=skrub.choose_float(0.1, 10.0, log=True, name="C"))
rf = RandomForestClassifier(
    n_estimators=skrub.choose_int(3, 30, name="n estimators"), random_state=0
)
classifier = skrub.choose_from(
    {"logistic": logistic, "random forest": rf, "dummy": DummyClassifier()},
    name="classifier",
)
pred = X.skb.apply(classifier, y=y)

search = pred.skb.make_optuna_search(fitted=True, random_state=0)  # <--- 🎯 optuna --------------

print("\nparam grid:")
print(pred.skb.describe_param_grid())

print("best params:")
print(search.best_learner_.describe_params())

prints

[... optuna output]
[I 2025-10-06 23:33:43,972] Trial 9 finished with value: 0.8399999999999999 and parameters: {'2:classifier': '0:logistic', '0:C': 0.10037026315054523}. Best is trial 6 with value: 0.9099999999999999.

param grid:
- classifier: 'logistic'
  C: choose_float(0.1, 10.0, log=True, name='C')
- classifier: 'random forest'
  n estimators: choose_int(3, 30, name='n estimators')
- classifier: 'dummy'

best params:
{'n estimators': 7, 'classifier': 'random forest'}

We can pass our own optuna study to control how it is created
For example above change make_optuna_search() to:

import optuna

study = optuna.create_study(
    storage="sqlite:///db.sqlite3", direction="maximize", study_name="classifier"
)

search = pred.skb.make_optuna_search(fitted=True, n_trials=200, study=study)

and we can monitor the search progress / explore results with ❯ optuna-dashboard sqlite:///db.sqlite3

low-level interface

For more control over the use of optuna, we can pass a trial to .skb.make_learner() to get a skrublearner initialized with choice outcomes suggested by the optuna trial. this is useful for more advanced optuna features (multi-objective, ask-and-tell interface, ...). Here is a simple example revisited from the optuna tutorial:

import optuna
import skrub

x = skrub.choose_float(-10, 10, name="x").as_data_op()
err = (x - 2) ** 2


def objective(trial):
    learner = err.skb.make_learner(choosing=trial)
    return learner.fit_transform({})


study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=200)
print(study.best_params)

prints something similar to

[...]
[I 2025-10-06 23:41:17,610] Trial 99 finished with value: 1.210419581181497 and parameters: {'0:x': 0.8998092978117398}. Best is trial 61 with value: 0.0005685118895572634.
{'0:x': 2.0238434873614843}

<

[rcap107]

This is really cool! Thanks a lot @jeromedockes!

[GaelVaroquaux]

Thanks a lot for the PR!!

The DataOp now has a .skb.make_optuna_search()

From an API perspective I would much prefer if we could avoid making a function for each hyper-parameter search method. It will scale very poorly as we add methods. I do realize that we might need different functions, as there will be different signatures, but can we find commonalities?

[rcap107]

Thanks a lot for the PR!!
The DataOp now has a .skb.make_optuna_search()
From an API perspective I would much prefer if we could avoid making a function for each hyper-parameter search method. It will scale very poorly as we add methods. I do realize that we might need different functions, as there will be different signatures, but can we find commonalities?

You mean something like this?

predictions.skb.make_param_search(kind="random", ...)

Is it going to be necessary? Are there so many other param search methods we can add?

[GaelVaroquaux]

You mean something like predictions.skb.make_param_search(kind="random", ...)

Yes

[jeromedockes]

The main parameters which differ between those search strategies are:

• study, timeout: only for optuna. (also n_iter is called n_trials in optuna, and optuna has a few additional less important parameters)
• n_iter/n_trials, random_state: only for optuna and randomized_search, not for grid_search

all but random_state are important/often-used parameters.

In the return value, the optuna search will have a study_ attribute which the scikit-learn ones don't. We may also later want to add capabilities like resuming the search which only apply to optuna.

I guess it's a tradeoff between having several functions or having a few "param a is ignored unless param b = x" in the docstring and different return types. Scikit-learn chose to have a different class for each search strategy. Regarding the addition of a new search strategy, if we go for multiple methods the cost is to create a new method (which is annoying and clutters the namespace), if we go for a single method, it depends on how many new parameters the new strategy needs to introduce (docstring gets more complex as the ratio shared parameters / all parameters decreases).

As strategies to add I'm guessing you have in mind the HalvingRandomSearchCV and HalvingGridSearchCV, once those stop being experimental in scikit-learn? And there are also other packages like hyperopt, not sure if we will want to support more or if optuna is enough (I would say just optuna should be good for the foreseeable future).

[jeromedockes]

In any case I think I'll split this PR. I'll remove the make_optuna_search() to keep only the possibility to pass a trial in .skb.make_learner(choose=optuna_trial) (the name "choose" is up for debate). This is the main addition which by itself enables the use of optuna with the data ops, and a minimal addition to the skrub API. The make_search method is boilerplate around that which I can add in another PR once we settle on the API question.

[GaelVaroquaux]

Thanks @jeromedockes !

After discussion with @rcap107 , here are my thoughts:

This PR caters to optuna power user. Many people are not optuna users, but would like to flip a switch to use optuna.

Looking at the example https://output.circle-artifacts.com/output/job/8860ed56-510b-47b7-b260-b71da3da42a1/artifacts/0/doc/auto_examples/data_ops/13_optuna_choices.html#sphx-glr-auto-examples-data-ops-13-optuna-choices-py with @rcap107 it seems that we could do a: make_iter_search that could use either scikit-learn's RandomSearch, or optuna (or more in the long run).

The internals of make_iter_search would really be wrapping code like the one that you wrote in the example:

import optuna


def objective(trial):
    learner = pred.skb.make_learner(choose=trial)
    cv_results = skrub.cross_validate(learner, environment=env, cv=cv)
    return cv_results["test_score"].mean()


study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=16)

best_learner = pred.skb.make_learner(choose=study.best_trial)

We would then modify the existing example to show first at the top how to easily use optuna, just flipping a switch from code using RandomSearch, and then continue with your example for advanced users.

We'll also need to do a bit of impedance matching, converting "n_iter" which is scikit-learn's convention to "n_trials", and exposing the optuna solver.

[jeromedockes]

thanks @GaelVaroquaux yes that is the plan, to have the scikit-learn estimator interface for most users, and the more bare-bones direct optuna use for optuna users. As I mentioned I opted for splitting the PR so that I could make progress on the core functionality while we decided on the make_<estimator>() specific API. Note the benefit of optuna is not only the choice of samplers but also features that make it user-friendly (like the visualizations or resuming searches), some of which will probably be more easily available through the "optuna user" interface.

[jeromedockes]

circleci failure is unrelated and will be fixed by #1794

[GaelVaroquaux]

I'm looking at this PR. It's really amazing, thank you!

I'll give some feedback as I go, as I worry that I'll be interrupt. Sorry for the resulting mess and crowding.

Feedback 1: I worry that we are adding an example with a run-time of 4mn. We have another PR doing this. it's not going to scale well, and soon it will impact negatively skrub development. Would it be possible to play a bit wit the parameters of the example (for instance number of trials) to decrease the cost of this new example a bit?

Maybe use a faster learner (I love ExtraTrees for instance, or even a Logistic?).

Feedback 2: I see that we are tuning the number of estimators of RandomForests. I don't consider this as the most important hyper parameter. I typically consider the depth as the most important feature to tune.

[GaelVaroquaux]

In the interest of compute time, could we replace one of these two by a logistic regression?

[jeromedockes]

I used a ExtraTreesRegressor and a Ridge and reduced the number of trials

[GaelVaroquaux]

A few other minor comments.

That's all I have. Thank you very much!!

[jeromedockes]

Thanks for the review @GaelVaroquaux !

• example time: I reduced the number of trials, and now use extra trees & ridge. it is still close to 2 minutes, LMK what you think. we can look for another dataset, subsample this one, or tune only the ridge's alpha instead of also including a better model.
• tuning n_estimators: yes you are right it's not a good hyperparam to tune, I think I just picked the first from the estimator's parameter list 😅 . I now use the min_samples_leaf (it also controls the depth in a way, as you suggested, but it is a bit simpler to use and visualize in the example because it is always a number whereas depth has the None option). The n_estimators is tuned in a few other places in the main branch that we can tackle in a separate documentation-only pr

I also edited a bit the user guide page (@rcap107 wrote the original version actually, thanks Riccardo!) . Part of it is very similar to the sphinx gallery example but that may not be a problem 🤔 .

[GaelVaroquaux]

Thanks heaps for your changes!! min_sample_leaf is great!

Regarding compute time:

it is still close to 2 minutes, LMK what you think.

A factor 2 is really good to get. Thanks a lot. If we can get a bit more, it's better

we can look for another dataset, subsample this one

Do you want to try this to see if you can get something that is still convincing and closer to 1mn?

[jeromedockes]

ok I just sampled half of the dataset, the examples takes just under a minute now on circleci

[GaelVaroquaux]

This is amazing!

[GaelVaroquaux]

Thank you so much @jeromedockes

[jeromedockes]

thank you both for the reviews & discussion :) looking forward to the upcoming release!