Skip to content

Implementing a port to optuna optimizers¤

This tutorial shows how to wrap an Optuna sampler as an f3dasm Block, following the ask/tell update-step pattern. An update-step block performs exactly one iteration per call: it registers the previous iteration's evaluated result with Optuna's internal study, asks for one new candidate, and appends it as an unevaluated row to data. A DataGenerator then evaluates that row, and the pair (update_step >> data_generator).loop(n) drives the optimization loop.

Use this as a template for any ask/tell third-party library.

See the Optimization Loops tutorial for the pattern itself.

from typing import Optional

import optuna

from f3dasm import Block, ExperimentData
from f3dasm._src.design.parameter import (
    ArrayParameter,
    CategoricalParameter,
    ConstantParameter,
    ContinuousParameter,
    DiscreteParameter,
)
from f3dasm._src.experimentsample import ExperimentSample
from f3dasm.design import Domain

Domain to optuna.distribution mapping¤

Optuna uses its own distribution classes to represent the parameter space. In f3dasm, the Domain class plays the same role. We need a translation:

def domain_to_optuna_distributions(domain: Domain) -> dict:
    optuna_distributions = {}
    for name, parameter in domain.input_space.items():
        if isinstance(parameter, ContinuousParameter):
            optuna_distributions[name] = (
                optuna.distributions.FloatDistribution(
                    low=parameter.lower_bound,
                    high=parameter.upper_bound,
                    log=parameter.log,
                )
            )
        elif isinstance(parameter, DiscreteParameter):
            optuna_distributions[name] = optuna.distributions.IntDistribution(
                low=parameter.lower_bound,
                high=parameter.upper_bound,
                step=parameter.step,
            )
        elif isinstance(parameter, CategoricalParameter):
            optuna_distributions[name] = (
                optuna.distributions.CategoricalDistribution(
                    parameter.categories
                )
            )
        elif isinstance(parameter, ConstantParameter):
            optuna_distributions[name] = (
                optuna.distributions.CategoricalDistribution(
                    choices=[parameter.value]
                )
            )
        elif isinstance(parameter, ArrayParameter):
            raise ValueError(
                "ArrayParameter is not supported in Optuna trials."
            )
        else:
            raise TypeError(
                f"Unsupported parameter type: {type(parameter)} for {name}"
            )
    return optuna_distributions

ExperimentSample from an optuna.Trial¤

Optuna's ask() returns a Trial that we query for suggested parameter values. We translate those values into an ExperimentSample:

def suggest_experimentsample(
    trial: optuna.Trial, domain: Domain
) -> ExperimentSample:
    optuna_dict = {}
    for name, parameter in domain.input_space.items():
        if isinstance(parameter, ContinuousParameter):
            optuna_dict[name] = trial.suggest_float(
                name=name,
                low=parameter.lower_bound,
                high=parameter.upper_bound,
                log=parameter.log,
            )
        elif isinstance(parameter, DiscreteParameter):
            optuna_dict[name] = trial.suggest_int(
                name=name,
                low=parameter.lower_bound,
                high=parameter.upper_bound,
                step=parameter.step,
            )
        elif isinstance(parameter, CategoricalParameter):
            optuna_dict[name] = trial.suggest_categorical(
                name=name, choices=parameter.categories
            )
        elif isinstance(parameter, ConstantParameter):
            optuna_dict[name] = trial.suggest_categorical(
                name=name, choices=[parameter.value]
            )
        else:
            raise TypeError(
                f"Unsupported parameter type: {type(parameter)} for {name}"
            )
    return ExperimentSample(_input_data=optuna_dict)

The update-step Block¤

Each call does three things:

  1. If the previous iteration asked for a candidate that has since been evaluated by the data generator, register it with the Optuna study as a historical trial (so TPE can learn from it).
  2. Ask the Optuna sampler for a new candidate.
  3. Append the candidate as an unevaluated row to data and return the result.

arm(data) creates the Optuna study and seeds it with whatever evaluated rows data already contains.

class OptunaUpdateStep(Block):
    """Ask/tell Optuna update-step block.

    Drive it with ``(update_step >> data_generator).loop(n)``.

    Parameters
    ----------
    optuna_sampler : optuna.samplers.BaseSampler
        The Optuna sampler used for suggestion (e.g. ``TPESampler``).
    output_name : str
        Name of the output column to minimize.
    """

    def __init__(
        self,
        optuna_sampler: optuna.samplers.BaseSampler,
        output_name: str,
    ) -> None:
        self.optuna_sampler = optuna_sampler
        self.output_name = output_name

    def arm(self, data: ExperimentData) -> None:
        self.distributions = domain_to_optuna_distributions(data.domain)
        self.study = optuna.create_study(sampler=self.optuna_sampler)

        # Seed the study with any evaluated history already in data.
        for _, es in data:
            self.study.add_trial(
                optuna.trial.create_trial(
                    params=es.input_data,
                    distributions=self.distributions,
                    value=es.output_data[self.output_name],
                )
            )
        # Tracker for the most-recently-asked trial whose result hasn't
        # been registered yet.
        self._pending_trial_index: Optional[int] = None

    def call(self, data: ExperimentData, **kwargs) -> ExperimentData:
        # 1. Register the previous iteration's evaluated result, if any.
        if self._pending_trial_index is not None:
            es = data.get_experiment_sample(self._pending_trial_index)
            self.study.add_trial(
                optuna.trial.create_trial(
                    params=es.input_data,
                    distributions=self.distributions,
                    value=es.output_data[self.output_name],
                )
            )
            self._pending_trial_index = None

        # 2. Ask for a new candidate.
        trial = self.study.ask()
        new_es = suggest_experimentsample(trial=trial, domain=data.domain)

        # 3. Append unevaluated row; remember its index so the next call can
        #    read its evaluated output and tell the study.
        new_data = ExperimentData.from_data(
            data={0: new_es},
            domain=data.domain,
            project_dir=data._project_dir,
        )
        merged = data + new_data
        self._pending_trial_index = merged.index[-1]
        return merged

Example run¤

from f3dasm import create_sampler, datagenerator


@datagenerator(output_names="y")
def f(x: float) -> float:
    return x**2 + 1


# Domain + initial design
domain = Domain()
domain.add_float(name="x", low=-10, high=10)
domain.add_output("y")

data = ExperimentData(domain=domain)
data = create_sampler("random", seed=0).call(data, n_samples=1)
f.arm(data)
data = f.call(data)

# Custom update step + canonical loop
update_step = OptunaUpdateStep(
    optuna_sampler=optuna.samplers.TPESampler(seed=42),
    output_name="y",
)

loop = (update_step >> f).loop(10)
loop.arm(data)
data = loop.call(data)

data

Comparison with the built-in tpesampler factory¤

The custom OptunaUpdateStep above is functionally equivalent to the built-in block exposed via:

from f3dasm import create_optimizer

update_step = create_optimizer("tpesampler", output_name="y")
# or, equivalently:
from f3dasm.optimization import tpesampler
update_step = tpesampler(output_name="y")

Prefer the built-in unless you need to pass a custom Optuna sampler configuration.

Next: SciPy Integration