Optimizers API

AID2E Optimizers module - Bayesian and evolutionary algorithms.

This module provides various optimization algorithms for multi-objective optimization problems, including:

• BaseOptimizer: Abstract base class defining the common optimizer interface.
• AxOptimizer: Bayesian optimization using Ax with Sobol initialization, SAASBO surrogate model, and qNEHVI acquisition function.
• PyMOOProblem: PyMOO-only public Problem wrapper for structural ask/tell integrations. Ax does not expose an equivalent public problem wrapper.
• PyMOOOptimizer: Evolutionary optimization (GA, NSGA-II, NSGA-III, MOEA/D) via PyMOO's ask/tell interface for external evaluation.
• compute_pareto_front: Backend-agnostic Pareto front utility.

Attributes:

Name	Type	Description
__version__		Version string inherited from the main aid2e package.

AID2EProblem

Bases: PyMOOProblem

Deprecated compatibility alias for PyMOOProblem.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

class AID2EProblem(PyMOOProblem):
    """Deprecated compatibility alias for ``PyMOOProblem``."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Warn and delegate to ``PyMOOProblem``."""
        warnings.warn(
            "AID2EProblem is deprecated and will be removed in the next "
            "iteration. Use PyMOOProblem instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        # TODO: Remove this compatibility alias in the next iteration.
        super().__init__(*args, **kwargs)

__init__(*args, **kwargs)

Warn and delegate to PyMOOProblem.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    """Warn and delegate to ``PyMOOProblem``."""
    warnings.warn(
        "AID2EProblem is deprecated and will be removed in the next "
        "iteration. Use PyMOOProblem instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    # TODO: Remove this compatibility alias in the next iteration.
    super().__init__(*args, **kwargs)

AxOptimizer

Bases: BaseOptimizer

Ax-based Bayesian optimization for multi-objective optimization.

This optimizer uses the Ax platform for Bayesian optimization with support for multiple objectives and native Ax Modular BoTorch generation.

Attributes:

Name	Type	Description
config		AxOptimizerConfig instance with strategy settings.
objective_names	List[str]	List of objective metric names.
experiment		Ax Experiment object managing trials.
generation_strategy		Ax GenerationStrategy for candidate generation.

Examples:

>>> from aid2e.optimizers.ax import AxOptimizer, AxOptimizerConfig
>>> search_space = SearchSpace(
...     parameters={
...         "x": {"type": "range", "bounds": [0.0, 1.0]},
...         "y": {"type": "range", "bounds": [0.0, 1.0]}
...     }
... )
>>> config = AxOptimizerConfig(
...     initialization_strategy="sobol",
...     generator="BOTORCH_MODULAR"
... )
>>> optimizer = AxOptimizer(
...     search_space=search_space,
...     config=config,
...     objective_names=["loss", "time"]
... )

Notes

This implementation defaults to a native Ax node-based generation strategy that transitions from an initializer node into Generators.BOTORCH_MODULAR.

Project: AID2E v0.0.1 - AI assisted Detector Design for EIC Homepage: https://aid2e.github.io/AID2E-framework Repository: https://github.com/aid2e/AID2E-framework.git

Source code in src/aid2e/optimizers/ax/optimizer.py

class AxOptimizer(BaseOptimizer):
    """Ax-based Bayesian optimization for multi-objective optimization.

    This optimizer uses the Ax platform for Bayesian optimization with
    support for multiple objectives and native Ax Modular BoTorch generation.

    Attributes:
        config: AxOptimizerConfig instance with strategy settings.
        objective_names: List of objective metric names.
        experiment: Ax Experiment object managing trials.
        generation_strategy: Ax GenerationStrategy for candidate generation.

    Examples:
        >>> from aid2e.optimizers.ax import AxOptimizer, AxOptimizerConfig
        >>> search_space = SearchSpace(
        ...     parameters={
        ...         "x": {"type": "range", "bounds": [0.0, 1.0]},
        ...         "y": {"type": "range", "bounds": [0.0, 1.0]}
        ...     }
        ... )
        >>> config = AxOptimizerConfig(
        ...     initialization_strategy="sobol",
        ...     generator="BOTORCH_MODULAR"
        ... )
        >>> optimizer = AxOptimizer(
        ...     search_space=search_space,
        ...     config=config,
        ...     objective_names=["loss", "time"]
        ... )

    Notes:
        This implementation defaults to a native Ax node-based generation
        strategy that transitions from an initializer node into
        ``Generators.BOTORCH_MODULAR``.

        Project: AID2E v0.0.1 - AI assisted Detector Design for EIC
        Homepage: https://aid2e.github.io/AID2E-framework
        Repository: https://github.com/aid2e/AID2E-framework.git
    """

    def __init__(
        self,
        search_space: Union[SearchSpace, DesignConfig],
        config: AxOptimizerConfig,
        objective_names: List[str],
        seed: Optional[int] = None,
        objective_directions: Optional[Dict[str, Any]] = None,
    ):
        """Initialize the Ax optimizer.

        Args:
            search_space: Parameter search space definition.
            config: AxOptimizerConfig instance with strategy settings.
            objective_names: List of objective metric names to optimize.
            seed: Random seed for reproducibility (overrides config.seed if provided).

        Raises:
            ImportError: If Ax is not installed.
            ValueError: If search_space is empty or config is invalid.

        Notes:
            The optimizer is initialized but not yet ready to suggest candidates.
            Ax Experiment and GenerationStrategy are created lazily on first use.
        """
        if not AX_AVAILABLE:
            raise ImportError(
                "Ax is required but not installed. "
                "Install with: pip install ax-platform==1.0.0"
            )
        if not AX_NODE_STRATEGY_AVAILABLE:
            raise RuntimeError(
                "The installed Ax runtime does not support the node-based "
                "generation API required by AID2E. Upgrade Ax to a version "
                "that provides CenterGenerationNode, GenerationNode, "
                "GeneratorSpec, and MinTrials."
            )

        # Initialize base class (handles DesignConfig → SearchSpace conversion)
        super().__init__(
            search_space=search_space,
            objective_names=objective_names,
            objective_directions=objective_directions,
            seed=seed if seed is not None else config.seed,
        )

        self.config = config
        # self.objective_names and self.n_objectives are inherited from BaseOptimizer

        # TODO Version check removed for now, re-enable when ready

        # Create Ax search space
        self.ax_search_space = self._create_ax_search_space()

        # Create optimization config
        self.optimization_config = self._create_optimization_config()

        # Create Ax experiment
        self.experiment = Experiment(
            name=f"aid2e_optimization",
            search_space=self.ax_search_space,
            optimization_config=self.optimization_config
        )

        # Create generation strategy
        self.generation_strategy = self._create_generation_strategy()

        # Track trials
        # self._trials and self._trial_counter are owned by BaseOptimizer

        logger.info(
            f"AxOptimizer initialized: {len(self.search_space.parameters)} params, "
            f"{len(objective_names)} objectives, strategy={config.initialization_strategy}, "
            f"generator={config.generator}"
        )

    def _parse_constraint_to_ax(
        self, constraint
    ) -> Optional[Any]:
        """Parse a ParameterConstraint rule to an Ax constraint object.

        Args:
            constraint: The ParameterConstraint from design_config.

        Returns:
            Ax constraint object (ParameterConstraint or SumConstraint), or None if parsing fails.

        Notes:
            Ax ParameterConstraint supports linear constraints of the form:
                sum(w_i * param_i) <= bound

            This method attempts to parse simple sum constraints like "x + y <= 1.5"
            into Ax's format. More complex expressions may not be supported.
        """
        import re

        rule = constraint.rule

        # Try to parse sum constraints: "param1 + param2 + ... <= bound" or "param1 + param2 + ... < bound"
        # Also handle >= and > by negating

        # Pattern: captures parameters, operator, and bound
        # Example: "DTLZ2.x1 + DTLZ2.x2 <= 1.5"
        pattern = r'^([^<>=]+)\s*([<>]=?)\s*([\d.]+)$'
        match = re.match(pattern, rule.strip())

        if not match:
            logger.warning(
                f"Constraint '{constraint.name}' has unsupported format: {rule}. "
                "Only simple sum constraints are supported (e.g., 'x + y <= 1.5')."
            )
            return None

        lhs, operator, bound_str = match.groups()
        bound = float(bound_str)

        # Parse left-hand side to extract parameters and coefficients
        # For now, only handle simple addition with coefficient 1
        # Pattern: param_name optionally preceded by + or -
        param_pattern = r'([+-]?)\s*([a-zA-Z_][a-zA-Z0-9_.]*)'
        terms = re.findall(param_pattern, lhs)

        if not terms:
            logger.warning(
                f"Constraint '{constraint.name}': Could not parse parameters from: {lhs}"
            )
            return None

        # Build constraint_dict: {param_name: coefficient}
        constraint_dict = {}
        for sign, param_name in terms:
            coeff = 1.0 if sign != '-' else -1.0
            constraint_dict[param_name.strip()] = coeff

        strict_epsilon = sys.float_info.epsilon # this is the smallest representable positive number such that 1.0 + eps != 1.0, used to convert strict inequalities to non-strict

        # Determine if upper or lower bound based on operator
        # sum <= bound or sum < bound: upper bound
        # sum >= bound or sum > bound: flip to -sum <= -bound (upper bound with negated coeffs)
        if operator in ['<=', '<']:
            is_upper_bound = True
            if operator == '<':
                bound -= strict_epsilon
        elif operator in ['>=', '>']:
            # Convert sum >= bound to -sum <= -bound
            is_upper_bound = True
            constraint_dict = {k: -v for k, v in constraint_dict.items()}
            bound = -bound
            if operator == '>':
                bound -= strict_epsilon
        else:
            logger.warning(f"Unsupported operator in constraint: {operator}")
            return None

        # Check if all coefficients are the same (typically 1.0 for sum constraints)
        coeffs = list(constraint_dict.values())
        if all(c == coeffs[0] for c in coeffs) and coeffs[0] == 1.0:
            # Use SumConstraint for simple sum constraints
            # Note: SumConstraint requires Parameter objects, not just names
            # We'll use ParameterConstraint instead which takes names
            pass

        try:
            terms_rendered = []
            for param_name, coeff in constraint_dict.items():
                if coeff == 1.0:
                    terms_rendered.append(param_name)
                elif coeff == -1.0:
                    terms_rendered.append(f"-{param_name}")
                else:
                    terms_rendered.append(f"{coeff}*{param_name}")

            inequality = " + ".join(terms_rendered).replace("+ -", "- ")
            inequality = f"{inequality} <= {bound}"
            ax_constraint = AxParameterConstraint(inequality=inequality)
            logger.debug(
                f"Converted constraint '{constraint.name}' to Ax format: "
                f"{inequality}"
            )
            return ax_constraint
        except Exception as e:
            logger.warning(
                f"Failed to create Ax constraint for '{constraint.name}': {e}"
            )
            return None

    def _create_ax_search_space(self) -> AxSearchSpace:
        """Create an Ax SearchSpace from the typed SearchSpace parameters.

        Returns:
            Ax SearchSpace describing the optimization domain with constraints.

        Raises:
            ValueError: If a parameter type is not supported by the Ax backend.

        Notes:
            Constraints from search_space.constraints are automatically converted
            to Ax ParameterConstraint objects and included in the search space.
        """

        ax_params = []
        for param_name, param in self.search_space.parameters.items():
            if isinstance(param, DesignRangeParameter):
                lower, upper = param.bounds
                ax_params.append(
                    AxRangeParameter(
                        name=param_name,
                        parameter_type=ParameterType.FLOAT,
                        lower=float(lower),
                        upper=float(upper),
                    )
                )
            elif isinstance(param, DesignChoiceParameter):
                ax_params.append(
                    AxChoiceParameter(
                        name=param_name,
                        parameter_type=ParameterType.STRING,
                        values=list(param.choices),
                        is_ordered=False,
                    )
                )
            else:
                raise ValueError(
                    f"Unsupported parameter type for Ax: {param.__class__.__name__}"
                )

        # Convert design constraints to Ax parameter constraints
        ax_constraints = []
        for constraint in self.search_space.constraints:
            ax_constraint = self._parse_constraint_to_ax(constraint)
            if ax_constraint is not None:
                ax_constraints.append(ax_constraint)
                logger.debug(f"Added constraint '{constraint.name}': {constraint.rule}")

        return AxSearchSpace(
            parameters=ax_params,
            parameter_constraints=ax_constraints if ax_constraints else None
        )

    def _create_optimization_config(self):
        """Create Ax optimization configuration for multi-objective optimization.

        Returns:
            Ax OptimizationConfig or None for single objective.
        """
        if len(self.objective_names) == 1:
            # Single objective case
            name = self.objective_names[0]
            direction = getattr(
                self.objective_directions.get(name),
                "value",
                self.objective_directions.get(name, "minimize"),
            )
            minimize = str(direction).lower() != "maximize"
            return OptimizationConfig(
                objective=Objective(
                    metric=Metric(name=name, lower_is_better=minimize),
                    minimize=minimize,
                )
            )
        else:
            # Multi-objective case
            objectives = []
            for name in self.objective_names:
                direction = getattr(
                    self.objective_directions.get(name),
                    "value",
                    self.objective_directions.get(name, "minimize"),
                )
                minimize = str(direction).lower() != "maximize"
                objectives.append(
                    Objective(
                        metric=Metric(name=name, lower_is_better=minimize),
                        minimize=minimize,
                    )
                )
            objective_thresholds = []
            if self.config.objective_thresholds:
                for name, bound in self.config.objective_thresholds.items():
                    direction = getattr(
                        self.objective_directions.get(name),
                        "value",
                        self.objective_directions.get(name, "minimize"),
                    )
                    minimize = str(direction).lower() != "maximize"
                    objective_thresholds.append(
                        ObjectiveThreshold(
                            metric=Metric(name=name, lower_is_better=minimize),
                            bound=float(bound),
                            relative=False,
                            op=ComparisonOp.LEQ if minimize else ComparisonOp.GEQ,
                        )
                    )
            return MultiObjectiveOptimizationConfig(
                objective=MultiObjective(objectives=objectives),
                objective_thresholds=objective_thresholds,
            )

    def _create_generation_strategy(self):
        """Create Ax GenerationStrategy based on config.

        Returns:
            Ax GenerationStrategy configured with chosen initialization +
            model-based optimization backend.

        Notes:
            Strategy uses the required node-based Ax API and transitions from
            initialization into the configured model-based generator.
        """
        return self._create_node_generation_strategy()

    def _create_node_generation_strategy(self):
        """Create a node-based generation strategy using the latest Ax APIs.

        Notes:
            This mirrors the modern Modular BoTorch tutorial pattern of chaining
            CenterOfSearchSpace -> initializer node -> model-based node using
            transition criteria such as MinTrials.
        """
        model_node_name = self._get_model_node_name()
        model_node = GenerationNode(
            name=model_node_name,
            generator_specs=[
                GeneratorSpec(
                    generator_enum=self._get_model_based_generator_enum(),
                    generator_kwargs=self._get_model_generator_kwargs(),
                    generator_gen_kwargs=self._get_model_generator_gen_kwargs(),
                )
            ],
        )

        nodes = []
        init_strategy = self.config.initialization_strategy.lower()
        init_trials = int(self.config.n_initial_samples)

        if init_strategy == "center":
            remaining_init_trials = max(0, init_trials - 1)
            next_node_name = model_node.name
            if remaining_init_trials > 0:
                init_node = self._build_initialization_node(
                    node_name="Sobol",
                    generator_enum=Generators.SOBOL,
                    num_trials=remaining_init_trials,
                    transition_to=model_node.name,
                )
                next_node_name = init_node.name
                nodes.append(init_node)

            nodes.insert(0, CenterGenerationNode(next_node_name=next_node_name))
            nodes.append(model_node)
            return GenerationStrategy(
                name=f"Center+{next_node_name}+{model_node.name}",
                nodes=nodes,
            )

        init_node = self._build_initialization_node(
            node_name="Random" if init_strategy == "random" else "Sobol",
            generator_enum=self._get_initialization_model_enum(),
            num_trials=init_trials,
            transition_to=model_node.name,
        )

        return GenerationStrategy(
            name=f"{init_node.name}+{model_node.name}",
            nodes=[init_node, model_node],
        )

    def _build_initialization_node(
        self,
        *,
        node_name: str,
        generator_enum: Any,
        num_trials: int,
        transition_to: str,
    ):
        """Build one initialization node for the node-based Ax API."""
        return GenerationNode(
            name=node_name,
            generator_specs=[
                GeneratorSpec(
                    generator_enum=generator_enum,
                    generator_kwargs=self._get_initialization_generator_kwargs(),
                )
            ],
            transition_criteria=[
                MinTrials(
                    threshold=num_trials,
                    transition_to=transition_to,
                    use_all_trials_in_exp=True,
                )
            ],
        )

    def _get_initialization_model_enum(self):
        """Return Ax initializer generator enum for the configured strategy."""
        init_strategy = self.config.initialization_strategy.lower()
        if init_strategy == "uniform":
            uniform = getattr(Generators, "UNIFORM", None)
            if uniform is not None:
                return uniform
            logger.warning(
                "Generators.UNIFORM is unavailable in this Ax version; "
                "falling back to Sobol for initialization."
            )
        return Generators.SOBOL

    def _get_initialization_generator_kwargs(self) -> Dict[str, Any]:
        """Return generator kwargs for initialization nodes."""
        if self.seed is None:
            return {}
        return {"seed": int(self.seed)}

    def _get_model_based_generator_enum(self):
        """Return the configured Ax model-based generator enum."""
        generator_name = self.config.generator
        generator_enum = getattr(Generators, generator_name, None)
        if generator_enum is None:
            raise ValueError(
                f"Configured Ax generator '{generator_name}' is unavailable in "
                "the installed Ax version."
            )
        return generator_enum

    def _get_model_node_name(self) -> str:
        """Return the display name for the active model-based generation node."""
        if self.config.generator == "BOTORCH_MODULAR":
            return "ModularBoTorch"
        return self.config.generator

    def _get_model_generator_kwargs(self) -> Dict[str, Any]:
        """Return resolved generator kwargs for the configured backend."""
        return resolve_generator_kwargs(
            generator_name=self.config.generator,
            generator_kwargs=self.config.generator_kwargs,
        )

    def _get_model_generator_gen_kwargs(self) -> Dict[str, Any]:
        """Return generation-time kwargs for model-based candidate generation."""
        return deepcopy(self.config.generator_gen_kwargs)

    def _split_generator_run(self, generator_run: Any) -> List[Any]:
        """Split a possibly batched Ax generator run into per-arm runs.

        Args:
            generator_run: A generator run returned by Ax.

        Returns:
            List of single-arm generator runs.
        """
        if not hasattr(generator_run, "arms"):
            raise TypeError(
                "Expected an Ax generator run with an 'arms' attribute, got "
                f"{type(generator_run).__name__}"
            )

        arms = list(generator_run.arms)
        if len(arms) <= 1:
            return [generator_run]

        from ax.core.generator_run import GeneratorRun

        weights = list(getattr(generator_run, "weights", []) or [])
        single_runs: List[Any] = []
        for index, arm in enumerate(arms):
            weight = weights[index] if index < len(weights) else 1.0
            single_run = GeneratorRun(
                arms=[arm],
                weights=[weight],
                fit_time=getattr(generator_run, "fit_time", None),
                gen_time=getattr(generator_run, "gen_time", None),
                generation_node_name=getattr(generator_run, "_generation_node_name", None),
            )
            for attr_name in ("_model_key", "_generation_node_name"):
                if hasattr(generator_run, attr_name):
                    setattr(single_run, attr_name, getattr(generator_run, attr_name))
            single_runs.append(single_run)

        return single_runs

    def _normalize_generator_runs(self, gen_result: Any) -> List[Any]:
        """Normalize Ax generation output into per-arm generator runs.

        Args:
            gen_result: Value returned by ``generation_strategy.gen(...)``.

        Returns:
            List of single-arm generator runs.

        Raises:
            TypeError: If the return shape cannot be interpreted as generator
                run output.

        Notes:
            Newer Ax internals may return wrapper/list-like structures around
            generator runs. AID2E uses one Ax generation call per requested
            batch and then splits the result into individual trial records.
        """
        if hasattr(gen_result, "arms"):
            return self._split_generator_run(gen_result)

        if isinstance(gen_result, (list, tuple)):
            generator_runs: List[Any] = []
            for item in gen_result:
                generator_runs.extend(self._normalize_generator_runs(item))
            return generator_runs

        if hasattr(gen_result, "generator_run_structs"):
            structs = getattr(gen_result, "generator_run_structs")
            generator_runs: List[Any] = []
            for struct in structs:
                generator_runs.extend(self._normalize_generator_runs(struct))
            return generator_runs

        if hasattr(gen_result, "generator_run"):
            return self._normalize_generator_runs(getattr(gen_result, "generator_run"))

        raise TypeError(
            "Unsupported Ax generation result type: "
            f"{type(gen_result).__name__}"
        )

    def _get_generation_strategy_metadata(self) -> Dict[str, Any]:
        """Return strategy progress metadata compatible across Ax APIs.

        Returns:
            Dictionary containing ``ax_step_index`` when available and
            ``ax_node_name`` for node-based strategies.
        """
        metadata: Dict[str, Any] = {"ax_step_index": -1}

        step_index = getattr(self.generation_strategy, "current_step_index", None)
        if step_index is not None:
            try:
                metadata["ax_step_index"] = int(step_index)
            except (TypeError, ValueError):
                logger.debug("Unable to coerce Ax step index '%s' to int", step_index)

        node_name = getattr(self.generation_strategy, "current_node_name", None)
        if node_name is not None:
            metadata["ax_node_name"] = str(node_name)

        return metadata

    def suggest_candidates(self, n_candidates: int = 1) -> List[Dict[str, Any]]:
        """Suggest the next batch of parameter configurations to evaluate.

        Args:
            n_candidates: Number of candidates to generate.

        Returns:
            List of parameter dictionaries ready for evaluation.

        Notes:
            Constraints are handled natively by Ax when present in search_space.
            Ax enforces constraints during candidate generation automatically.

            IMPLEMENTATION NOTE: This method generates the requested batch in a
            single Ax call whenever possible, then splits the resulting batch
            into individual AID2E trial records so the rest of the framework can
            keep using single-trial ``update_with_results`` semantics.

            The GenerationStrategy still properly tracks progress: after
            n_initial_samples individual trials complete (regardless of how
            they were generated), it automatically switches from Sobol to BO.
        """

        candidates = []
        model_keys: List[str] = []

        generator_run_result = self.generation_strategy.gen(
            experiment=self.experiment,
            n=n_candidates,
        )
        strategy_metadata = self._get_generation_strategy_metadata()
        generator_runs = self._normalize_generator_runs(generator_run_result)
        if len(generator_runs) < n_candidates:
            raise RuntimeError(
                "Ax returned fewer generator runs than requested: "
                f"requested={n_candidates}, got={len(generator_runs)}"
            )

        for generator_run in generator_runs[:n_candidates]:
            trial = self.experiment.new_trial(generator_run=generator_run)
            trial.mark_running(no_runner_required=True)
            arm = generator_run.arms[0]
            candidate_params = dict(arm.parameters)
            model_key = getattr(generator_run, "_model_key", "unknown")
            self.set_trial_status(
                trial_index=trial.index,
                status=TRIAL_STATUS_SUGGESTED,
                parameters=candidate_params,
                metrics=None,
                metadata={
                    "ax_model_key": model_key,
                    **strategy_metadata,
                },
            )
            candidates.append(candidate_params)
            model_keys.append(model_key)

        self._trial_counter = max(self._trial_counter, len(self.experiment.trials))
        strategy_metadata = self._get_generation_strategy_metadata()
        strategy_ref = strategy_metadata.get(
            "ax_node_name",
            strategy_metadata.get("ax_step_index", -1),
        )

        logger.debug(
            "Generated %d candidates using %s (trials %d-%d, strategy=%s)",
            n_candidates,
            ",".join(model_keys) if model_keys else "unknown",
            len(self.experiment.trials) - n_candidates,
            len(self.experiment.trials) - 1,
            strategy_ref,
        )
        return candidates

    def update_with_results(
        self,
        trial_index: int,
        parameters: Dict[str, Any],
        metrics: Dict[str, float]
    ) -> None:
        """Update optimizer with evaluation results from a trial.

        Args:
            trial_index: Unique identifier for the trial.
            parameters: Parameter values that were evaluated.
            metrics: Objective values obtained from evaluation.

        Examples:
            >>> optimizer.update_with_results(
            ...     trial_index=0,
            ...     parameters={'x': 0.5, 'y': 0.3},
            ...     metrics={'loss': 0.1, 'accuracy': 0.9}
            ... )

        Notes:
            Completes the trial in Ax experiment and attaches data.
            This allows the surrogate model to learn from the evaluation.
        """
        # Validate metrics
        for obj_name in self.objective_names:
            if obj_name not in metrics:
                raise ValueError(
                    f"Missing objective '{obj_name}' in metrics. "
                    f"Expected: {self.objective_names}, got: {list(metrics.keys())}"
                )

        # Get the trial from experiment
        if trial_index < len(self.experiment.trials):
            trial = self.experiment.trials[trial_index]

            # Complete the trial with data
            trial.mark_completed()

            # Attach data to experiment
            from ax.core.data import Data
            import pandas as pd

            data_rows = []
            for metric_name, metric_value in metrics.items():
                if metric_name in self.objective_names:
                    data_rows.append({
                        'trial_index': trial_index,
                        'metric_name': metric_name,
                        'metric_signature': metric_name,
                        'arm_name': trial.arm.name if trial.arm else f"arm_{trial_index}",
                        'mean': float(metric_value),
                        'sem': 0.0  # Standard error of mean (0 for deterministic)
                    })

            if data_rows:
                df = pd.DataFrame(data_rows)
                data = Data(df=df)
                self.experiment.attach_data(data)

        # Update internal trial tracking through the base API
        self.set_trial_status(
            trial_index=trial_index,
            status="completed",
            parameters=parameters,
            metrics={k: float(v) for k, v in metrics.items()},
        )

        logger.debug(
            f"Updated trial {trial_index} with {len(metrics)} metrics"
        )

    def mark_trial_failed(
        self,
        trial_index: int,
        *,
        parameters: Optional[Dict[str, Any]] = None,
        reason: Optional[str] = None,
    ) -> Trial:
        """Mark a failed evaluation in AID2E and the Ax experiment."""
        if trial_index in self.experiment.trials:
            self.experiment.trials[trial_index].mark_failed(reason=reason)
        return super().mark_trial_failed(
            trial_index,
            parameters=parameters,
            reason=reason,
        )

    def serialize_state(self) -> Dict[str, Any]:
        """Serialize optimizer state for distributed execution.

        Returns:
            Dictionary containing all necessary state to reconstruct the optimizer.

        Examples:
            >>> state = optimizer.serialize_state()
            >>> json.dumps(state)  # Should be JSON-serializable
        """
        space_payload = {
            name: param.model_dump()
            for name, param in self.search_space.parameters.items()
        }
        constraints_payload = [
            constraint.model_dump()
            for constraint in self.search_space.constraints
        ]

        return {
            "search_space": {
                "parameters": space_payload,
                "constraints": constraints_payload,
                "name": self.search_space.name,
            },
            "n_objectives": self.n_objectives,
            "seed": self.seed,
            "objective_names": self.objective_names,
            "config": {
                "initialization_strategy": self.config.initialization_strategy,
                "generator": self.config.generator,
                "generator_kwargs": deepcopy(self.config.generator_kwargs),
                "generator_gen_kwargs": deepcopy(self.config.generator_gen_kwargs),
                "objective_thresholds": (
                    deepcopy(self.config.objective_thresholds)
                    if self.config.objective_thresholds is not None
                    else None
                ),
                "n_initial_samples": self.config.n_initial_samples,
                "n_iterations": self.config.n_iterations,
                "batch_size": self.config.batch_size,
                "seed": self.config.seed,
            },
            "trials": [
                {
                    "index": t.index,
                    "parameters": t.parameters,
                    "metrics": t.metrics,
                    "status": t.status,
                    "metadata": t.metadata
                }
                for t in self._trials if t is not None
            ],
            "trial_counter": self._trial_counter
        }

    def load_state(self, state: Dict[str, Any]) -> None:
        """Load optimizer state from serialized form.

        Args:
            state: Dictionary containing serialized optimizer state.

        Raises:
            ValueError: If state is invalid or incompatible.

        Notes:
            This recreates the Ax experiment and generation strategy,
            then replays all trials to restore the optimizer state.
        """
        # Validate state
        required_keys = ["search_space", "objective_names", "config", "trials"]
        for key in required_keys:
            if key not in state:
                raise ValueError(f"Missing required key in state: {key}")

        # Restore config and objective metadata
        self.config = AxOptimizerConfig(**state["config"])
        self.objective_names = list(state["objective_names"])

        # Rebuild search space and Ax components from serialized payload
        saved_space = state["search_space"] or {}
        parameters_payload = saved_space.get("parameters", saved_space)
        constraints_payload = saved_space.get("constraints", [])

        self.search_space = SearchSpace(
            parameters=parameters_payload,
            constraints=constraints_payload,
            name=saved_space.get("name"),
        )

        self.ax_search_space = self._create_ax_search_space()
        self.optimization_config = self._create_optimization_config()
        self.experiment = Experiment(
            name="aid2e_optimization",
            search_space=self.ax_search_space,
            optimization_config=self.optimization_config,
        )
        self.generation_strategy = self._create_generation_strategy()

        # Restore trials
        self._trials = []
        for trial_data in state["trials"]:
            trial = Trial(
                index=trial_data["index"],
                parameters=trial_data["parameters"],
                metrics=trial_data.get("metrics"),
                status=trial_data.get("status", "pending"),
                metadata=trial_data.get("metadata", {})
            )

            while len(self._trials) <= trial.index:
                self._trials.append(None)
            self._trials[trial.index] = trial

            # Replay trial in Ax experiment if completed
            if trial.status == "completed" and trial.metrics:
                # Create trial in Ax
                ax_trial = self.experiment.new_trial()
                ax_trial.mark_running(no_runner_required=True)

                # Update with results
                self.update_with_results(
                    trial_index=trial.index,
                    parameters=trial.parameters,
                    metrics=trial.metrics
                )

        self._trial_counter = state.get("trial_counter", len(self._trials))
        self._trial_counter = max(self._trial_counter, len(self.experiment.trials))

        logger.info(f"Loaded optimizer state with {len(self._trials)} trials")

    def __repr__(self) -> str:
        """Return string representation of the optimizer.

        Returns:
            String describing the optimizer configuration.
        """
        return (
            f"AxOptimizer("
            f"n_params={len(self.search_space.parameters)}, "
            f"n_objectives={self.n_objectives}, "
            f"strategy={self.config.initialization_strategy}, "
            f"generator={self.config.generator}, "
            f"seed={self.seed}"
            f")"
        )

__init__(search_space, config, objective_names, seed=None, objective_directions=None)

Initialize the Ax optimizer.

Parameters:

Name	Type	Description	Default
search_space	Union[SearchSpace, DesignConfig]	Parameter search space definition.	required
config	AxOptimizerConfig	AxOptimizerConfig instance with strategy settings.	required
objective_names	List[str]	List of objective metric names to optimize.	required
seed	Optional[int]	Random seed for reproducibility (overrides config.seed if provided).	None

Raises:

Type	Description
ImportError	If Ax is not installed.
ValueError	If search_space is empty or config is invalid.

Notes

The optimizer is initialized but not yet ready to suggest candidates. Ax Experiment and GenerationStrategy are created lazily on first use.

Source code in src/aid2e/optimizers/ax/optimizer.py

def __init__(
    self,
    search_space: Union[SearchSpace, DesignConfig],
    config: AxOptimizerConfig,
    objective_names: List[str],
    seed: Optional[int] = None,
    objective_directions: Optional[Dict[str, Any]] = None,
):
    """Initialize the Ax optimizer.

    Args:
        search_space: Parameter search space definition.
        config: AxOptimizerConfig instance with strategy settings.
        objective_names: List of objective metric names to optimize.
        seed: Random seed for reproducibility (overrides config.seed if provided).

    Raises:
        ImportError: If Ax is not installed.
        ValueError: If search_space is empty or config is invalid.

    Notes:
        The optimizer is initialized but not yet ready to suggest candidates.
        Ax Experiment and GenerationStrategy are created lazily on first use.
    """
    if not AX_AVAILABLE:
        raise ImportError(
            "Ax is required but not installed. "
            "Install with: pip install ax-platform==1.0.0"
        )
    if not AX_NODE_STRATEGY_AVAILABLE:
        raise RuntimeError(
            "The installed Ax runtime does not support the node-based "
            "generation API required by AID2E. Upgrade Ax to a version "
            "that provides CenterGenerationNode, GenerationNode, "
            "GeneratorSpec, and MinTrials."
        )

    # Initialize base class (handles DesignConfig → SearchSpace conversion)
    super().__init__(
        search_space=search_space,
        objective_names=objective_names,
        objective_directions=objective_directions,
        seed=seed if seed is not None else config.seed,
    )

    self.config = config
    # self.objective_names and self.n_objectives are inherited from BaseOptimizer

    # TODO Version check removed for now, re-enable when ready

    # Create Ax search space
    self.ax_search_space = self._create_ax_search_space()

    # Create optimization config
    self.optimization_config = self._create_optimization_config()

    # Create Ax experiment
    self.experiment = Experiment(
        name=f"aid2e_optimization",
        search_space=self.ax_search_space,
        optimization_config=self.optimization_config
    )

    # Create generation strategy
    self.generation_strategy = self._create_generation_strategy()

    # Track trials
    # self._trials and self._trial_counter are owned by BaseOptimizer

    logger.info(
        f"AxOptimizer initialized: {len(self.search_space.parameters)} params, "
        f"{len(objective_names)} objectives, strategy={config.initialization_strategy}, "
        f"generator={config.generator}"
    )

__repr__()

Return string representation of the optimizer.

Returns:

Type	Description
str	String describing the optimizer configuration.

Source code in src/aid2e/optimizers/ax/optimizer.py

def __repr__(self) -> str:
    """Return string representation of the optimizer.

    Returns:
        String describing the optimizer configuration.
    """
    return (
        f"AxOptimizer("
        f"n_params={len(self.search_space.parameters)}, "
        f"n_objectives={self.n_objectives}, "
        f"strategy={self.config.initialization_strategy}, "
        f"generator={self.config.generator}, "
        f"seed={self.seed}"
        f")"
    )

load_state(state)

Load optimizer state from serialized form.

Parameters:

Name	Type	Description	Default
state	Dict[str, Any]	Dictionary containing serialized optimizer state.	required

Raises:

Type	Description
ValueError	If state is invalid or incompatible.

Notes

This recreates the Ax experiment and generation strategy, then replays all trials to restore the optimizer state.

Source code in src/aid2e/optimizers/ax/optimizer.py

def load_state(self, state: Dict[str, Any]) -> None:
    """Load optimizer state from serialized form.

    Args:
        state: Dictionary containing serialized optimizer state.

    Raises:
        ValueError: If state is invalid or incompatible.

    Notes:
        This recreates the Ax experiment and generation strategy,
        then replays all trials to restore the optimizer state.
    """
    # Validate state
    required_keys = ["search_space", "objective_names", "config", "trials"]
    for key in required_keys:
        if key not in state:
            raise ValueError(f"Missing required key in state: {key}")

    # Restore config and objective metadata
    self.config = AxOptimizerConfig(**state["config"])
    self.objective_names = list(state["objective_names"])

    # Rebuild search space and Ax components from serialized payload
    saved_space = state["search_space"] or {}
    parameters_payload = saved_space.get("parameters", saved_space)
    constraints_payload = saved_space.get("constraints", [])

    self.search_space = SearchSpace(
        parameters=parameters_payload,
        constraints=constraints_payload,
        name=saved_space.get("name"),
    )

    self.ax_search_space = self._create_ax_search_space()
    self.optimization_config = self._create_optimization_config()
    self.experiment = Experiment(
        name="aid2e_optimization",
        search_space=self.ax_search_space,
        optimization_config=self.optimization_config,
    )
    self.generation_strategy = self._create_generation_strategy()

    # Restore trials
    self._trials = []
    for trial_data in state["trials"]:
        trial = Trial(
            index=trial_data["index"],
            parameters=trial_data["parameters"],
            metrics=trial_data.get("metrics"),
            status=trial_data.get("status", "pending"),
            metadata=trial_data.get("metadata", {})
        )

        while len(self._trials) <= trial.index:
            self._trials.append(None)
        self._trials[trial.index] = trial

        # Replay trial in Ax experiment if completed
        if trial.status == "completed" and trial.metrics:
            # Create trial in Ax
            ax_trial = self.experiment.new_trial()
            ax_trial.mark_running(no_runner_required=True)

            # Update with results
            self.update_with_results(
                trial_index=trial.index,
                parameters=trial.parameters,
                metrics=trial.metrics
            )

    self._trial_counter = state.get("trial_counter", len(self._trials))
    self._trial_counter = max(self._trial_counter, len(self.experiment.trials))

    logger.info(f"Loaded optimizer state with {len(self._trials)} trials")

mark_trial_failed(trial_index, *, parameters=None, reason=None)

Mark a failed evaluation in AID2E and the Ax experiment.

Source code in src/aid2e/optimizers/ax/optimizer.py

def mark_trial_failed(
    self,
    trial_index: int,
    *,
    parameters: Optional[Dict[str, Any]] = None,
    reason: Optional[str] = None,
) -> Trial:
    """Mark a failed evaluation in AID2E and the Ax experiment."""
    if trial_index in self.experiment.trials:
        self.experiment.trials[trial_index].mark_failed(reason=reason)
    return super().mark_trial_failed(
        trial_index,
        parameters=parameters,
        reason=reason,
    )

serialize_state()

Serialize optimizer state for distributed execution.

Returns:

Type	Description
Dict[str, Any]	Dictionary containing all necessary state to reconstruct the optimizer.

Examples:

>>> state = optimizer.serialize_state()
>>> json.dumps(state)  # Should be JSON-serializable

Source code in src/aid2e/optimizers/ax/optimizer.py

def serialize_state(self) -> Dict[str, Any]:
    """Serialize optimizer state for distributed execution.

    Returns:
        Dictionary containing all necessary state to reconstruct the optimizer.

    Examples:
        >>> state = optimizer.serialize_state()
        >>> json.dumps(state)  # Should be JSON-serializable
    """
    space_payload = {
        name: param.model_dump()
        for name, param in self.search_space.parameters.items()
    }
    constraints_payload = [
        constraint.model_dump()
        for constraint in self.search_space.constraints
    ]

    return {
        "search_space": {
            "parameters": space_payload,
            "constraints": constraints_payload,
            "name": self.search_space.name,
        },
        "n_objectives": self.n_objectives,
        "seed": self.seed,
        "objective_names": self.objective_names,
        "config": {
            "initialization_strategy": self.config.initialization_strategy,
            "generator": self.config.generator,
            "generator_kwargs": deepcopy(self.config.generator_kwargs),
            "generator_gen_kwargs": deepcopy(self.config.generator_gen_kwargs),
            "objective_thresholds": (
                deepcopy(self.config.objective_thresholds)
                if self.config.objective_thresholds is not None
                else None
            ),
            "n_initial_samples": self.config.n_initial_samples,
            "n_iterations": self.config.n_iterations,
            "batch_size": self.config.batch_size,
            "seed": self.config.seed,
        },
        "trials": [
            {
                "index": t.index,
                "parameters": t.parameters,
                "metrics": t.metrics,
                "status": t.status,
                "metadata": t.metadata
            }
            for t in self._trials if t is not None
        ],
        "trial_counter": self._trial_counter
    }

suggest_candidates(n_candidates=1)

Suggest the next batch of parameter configurations to evaluate.

Parameters:

Name	Type	Description	Default
n_candidates	int	Number of candidates to generate.	1

Returns:

Type	Description
List[Dict[str, Any]]	List of parameter dictionaries ready for evaluation.

Notes

Constraints are handled natively by Ax when present in search_space. Ax enforces constraints during candidate generation automatically.

IMPLEMENTATION NOTE: This method generates the requested batch in a single Ax call whenever possible, then splits the resulting batch into individual AID2E trial records so the rest of the framework can keep using single-trial update_with_results semantics.

The GenerationStrategy still properly tracks progress: after n_initial_samples individual trials complete (regardless of how they were generated), it automatically switches from Sobol to BO.

Source code in src/aid2e/optimizers/ax/optimizer.py

def suggest_candidates(self, n_candidates: int = 1) -> List[Dict[str, Any]]:
    """Suggest the next batch of parameter configurations to evaluate.

    Args:
        n_candidates: Number of candidates to generate.

    Returns:
        List of parameter dictionaries ready for evaluation.

    Notes:
        Constraints are handled natively by Ax when present in search_space.
        Ax enforces constraints during candidate generation automatically.

        IMPLEMENTATION NOTE: This method generates the requested batch in a
        single Ax call whenever possible, then splits the resulting batch
        into individual AID2E trial records so the rest of the framework can
        keep using single-trial ``update_with_results`` semantics.

        The GenerationStrategy still properly tracks progress: after
        n_initial_samples individual trials complete (regardless of how
        they were generated), it automatically switches from Sobol to BO.
    """

    candidates = []
    model_keys: List[str] = []

    generator_run_result = self.generation_strategy.gen(
        experiment=self.experiment,
        n=n_candidates,
    )
    strategy_metadata = self._get_generation_strategy_metadata()
    generator_runs = self._normalize_generator_runs(generator_run_result)
    if len(generator_runs) < n_candidates:
        raise RuntimeError(
            "Ax returned fewer generator runs than requested: "
            f"requested={n_candidates}, got={len(generator_runs)}"
        )

    for generator_run in generator_runs[:n_candidates]:
        trial = self.experiment.new_trial(generator_run=generator_run)
        trial.mark_running(no_runner_required=True)
        arm = generator_run.arms[0]
        candidate_params = dict(arm.parameters)
        model_key = getattr(generator_run, "_model_key", "unknown")
        self.set_trial_status(
            trial_index=trial.index,
            status=TRIAL_STATUS_SUGGESTED,
            parameters=candidate_params,
            metrics=None,
            metadata={
                "ax_model_key": model_key,
                **strategy_metadata,
            },
        )
        candidates.append(candidate_params)
        model_keys.append(model_key)

    self._trial_counter = max(self._trial_counter, len(self.experiment.trials))
    strategy_metadata = self._get_generation_strategy_metadata()
    strategy_ref = strategy_metadata.get(
        "ax_node_name",
        strategy_metadata.get("ax_step_index", -1),
    )

    logger.debug(
        "Generated %d candidates using %s (trials %d-%d, strategy=%s)",
        n_candidates,
        ",".join(model_keys) if model_keys else "unknown",
        len(self.experiment.trials) - n_candidates,
        len(self.experiment.trials) - 1,
        strategy_ref,
    )
    return candidates

update_with_results(trial_index, parameters, metrics)

Update optimizer with evaluation results from a trial.

Parameters:

Name	Type	Description	Default
trial_index	int	Unique identifier for the trial.	required
parameters	Dict[str, Any]	Parameter values that were evaluated.	required
metrics	Dict[str, float]	Objective values obtained from evaluation.	required

Examples:

>>> optimizer.update_with_results(
...     trial_index=0,
...     parameters={'x': 0.5, 'y': 0.3},
...     metrics={'loss': 0.1, 'accuracy': 0.9}
... )

Notes

Completes the trial in Ax experiment and attaches data. This allows the surrogate model to learn from the evaluation.

Source code in src/aid2e/optimizers/ax/optimizer.py

def update_with_results(
    self,
    trial_index: int,
    parameters: Dict[str, Any],
    metrics: Dict[str, float]
) -> None:
    """Update optimizer with evaluation results from a trial.

    Args:
        trial_index: Unique identifier for the trial.
        parameters: Parameter values that were evaluated.
        metrics: Objective values obtained from evaluation.

    Examples:
        >>> optimizer.update_with_results(
        ...     trial_index=0,
        ...     parameters={'x': 0.5, 'y': 0.3},
        ...     metrics={'loss': 0.1, 'accuracy': 0.9}
        ... )

    Notes:
        Completes the trial in Ax experiment and attaches data.
        This allows the surrogate model to learn from the evaluation.
    """
    # Validate metrics
    for obj_name in self.objective_names:
        if obj_name not in metrics:
            raise ValueError(
                f"Missing objective '{obj_name}' in metrics. "
                f"Expected: {self.objective_names}, got: {list(metrics.keys())}"
            )

    # Get the trial from experiment
    if trial_index < len(self.experiment.trials):
        trial = self.experiment.trials[trial_index]

        # Complete the trial with data
        trial.mark_completed()

        # Attach data to experiment
        from ax.core.data import Data
        import pandas as pd

        data_rows = []
        for metric_name, metric_value in metrics.items():
            if metric_name in self.objective_names:
                data_rows.append({
                    'trial_index': trial_index,
                    'metric_name': metric_name,
                    'metric_signature': metric_name,
                    'arm_name': trial.arm.name if trial.arm else f"arm_{trial_index}",
                    'mean': float(metric_value),
                    'sem': 0.0  # Standard error of mean (0 for deterministic)
                })

        if data_rows:
            df = pd.DataFrame(data_rows)
            data = Data(df=df)
            self.experiment.attach_data(data)

    # Update internal trial tracking through the base API
    self.set_trial_status(
        trial_index=trial_index,
        status="completed",
        parameters=parameters,
        metrics={k: float(v) for k, v in metrics.items()},
    )

    logger.debug(
        f"Updated trial {trial_index} with {len(metrics)} metrics"
    )

AxOptimizerConfig

Bases: BaseModel

Configure the Ax optimizer backend for AID2E.

Define the tuning parameters used by the Ax-based Bayesian optimization workflow, including initialization behavior, model-generation settings, and core iteration controls.

Attributes:

Name	Type	Description
initialization_strategy	Literal['sobol', 'uniform', 'center']	Choose how the initial design points are drawn before model-based generation begins. Supported values are "sobol", "uniform", and "center".
generator	str	Specify the Ax generator enum name used for model-based candidate generation. The default is "BOTORCH_MODULAR".
generator_kwargs	dict[str, Any]	Provide keyword arguments for Ax GeneratorSpec setup (for example, model configuration details).
generator_gen_kwargs	dict[str, Any]	Provide generation-time keyword arguments passed into Ax candidate generation (for example, model_gen_options budgets).
objective_thresholds	Optional[dict[str, float]]	Optionally map objective metric names to threshold values for multi-objective optimization.
n_initial_samples	int	Set the number of initialization trials.
n_iterations	int	Set the total optimization iteration budget.
batch_size	int	Set the number of candidates proposed per iteration.
seed	Optional[int]	Set an optional random seed for reproducibility.

Examples:

>>> config = AxOptimizerConfig(
...     initialization_strategy="sobol",
...     generator="BOTORCH_MODULAR",
...     generator_kwargs={"fit_out_of_design": False},
...     generator_gen_kwargs={"model_gen_options": {"acqf_optimizer_kwargs": {"num_restarts": 8}}},
...     n_initial_samples=12,
...     n_iterations=60,
...     batch_size=2,
...     seed=42,
... )
>>> config.generator
'BOTORCH_MODULAR'

Notes

Legacy fields such as surrogate_model and acquisition_function are explicitly rejected.

Source code in src/aid2e/optimizers/ax/config.py

class AxOptimizerConfig(BaseModel):
    """Configure the Ax optimizer backend for AID2E.

    Define the tuning parameters used by the Ax-based Bayesian optimization
    workflow, including initialization behavior, model-generation settings, and
    core iteration controls.

    Attributes:
        initialization_strategy: Choose how the initial design points are drawn
            before model-based generation begins. Supported values are
            ``"sobol"``, ``"uniform"``, and ``"center"``.
        generator: Specify the Ax generator enum name used for model-based
            candidate generation. The default is ``"BOTORCH_MODULAR"``.
        generator_kwargs: Provide keyword arguments for Ax ``GeneratorSpec``
            setup (for example, model configuration details).
        generator_gen_kwargs: Provide generation-time keyword arguments passed
            into Ax candidate generation (for example,
            ``model_gen_options`` budgets).
        objective_thresholds: Optionally map objective metric names to
            threshold values for multi-objective optimization.
        n_initial_samples: Set the number of initialization trials.
        n_iterations: Set the total optimization iteration budget.
        batch_size: Set the number of candidates proposed per iteration.
        seed: Set an optional random seed for reproducibility.

    Examples:
        >>> config = AxOptimizerConfig(
        ...     initialization_strategy="sobol",
        ...     generator="BOTORCH_MODULAR",
        ...     generator_kwargs={"fit_out_of_design": False},
        ...     generator_gen_kwargs={"model_gen_options": {"acqf_optimizer_kwargs": {"num_restarts": 8}}},
        ...     n_initial_samples=12,
        ...     n_iterations=60,
        ...     batch_size=2,
        ...     seed=42,
        ... )
        >>> config.generator
        'BOTORCH_MODULAR'

    Notes:
        Legacy fields such as ``surrogate_model`` and
        ``acquisition_function`` are explicitly rejected.
    """
    # TODO: Remove the extra fields once the legacy config surface is fully retired.
    model_config = ConfigDict(extra="forbid")
    initialization_strategy: Literal["sobol", "uniform", "center"] = Field(
        default="sobol",
        description=(
            "Initialization strategy: 'sobol' for quasi-random initialization, "
            "'uniform' for uniform random initialization, or 'center' for one "
            "center point followed by additional initialization samples."
        ),
    )
    generator: str = Field(
        default="BOTORCH_MODULAR",
        description=(
            "Ax generator enum name. This backend currently supports "
            "'BOTORCH_MODULAR' and treats it as the default model-based backend."
        ),
    )
    generator_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Keyword arguments passed to Ax's GeneratorSpec for the configured "
            "model-based generator. YAML-friendly string values are resolved to "
            "supported Ax / BoTorch classes at runtime."
        ),
    )
    generator_gen_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Generation-time kwargs passed through Ax into candidate generation, "
            "such as optimizer budgets under 'model_gen_options'."
        ),
    )
    objective_thresholds: Optional[dict[str, float]] = Field(
        default=None,
        description=(
            "Optional objective thresholds for multi-objective optimization, "
            "keyed by metric name."
        ),
    )
    n_initial_samples: int = Field(
        default=10,
        ge=1,
        description="Number of samples in the initialization phase.",
    )
    n_iterations: int = Field(
        default=50,
        ge=1,
        description="Total number of optimization iterations.",
    )
    batch_size: int = Field(
        default=1,
        ge=1,
        description="Number of candidates to evaluate per iteration.",
    )
    seed: Optional[int] = Field(
        default=None,
        description="Random seed for reproducibility. If None, results are non-deterministic.",
    )

    @model_validator(mode="before")
    @classmethod
    def reject_legacy_fields(cls, raw_value: Any) -> Any:
        """Fail fast on the retired Ax config surface."""
        if not isinstance(raw_value, dict):
            return raw_value

        legacy_fields = [
            field_name
            for field_name in ("surrogate_model", "acquisition_function")
            if field_name in raw_value
        ]
        if legacy_fields:
            joined = ", ".join(legacy_fields)
            raise ValueError(
                f"AxOptimizerConfig no longer accepts legacy fields: {joined}. "
                "Use 'generator', 'generator_kwargs', and 'generator_gen_kwargs' "
                "instead."
            )
        return raw_value

    @field_validator("generator")
    @classmethod
    def normalize_generator(cls, value: str) -> str:
        """Normalize the configured generator to an Ax enum-style name."""
        return validate_generator_name(value)

normalize_generator(value) classmethod

Normalize the configured generator to an Ax enum-style name.

Source code in src/aid2e/optimizers/ax/config.py

@field_validator("generator")
@classmethod
def normalize_generator(cls, value: str) -> str:
    """Normalize the configured generator to an Ax enum-style name."""
    return validate_generator_name(value)

reject_legacy_fields(raw_value) classmethod

Fail fast on the retired Ax config surface.

Source code in src/aid2e/optimizers/ax/config.py

@model_validator(mode="before")
@classmethod
def reject_legacy_fields(cls, raw_value: Any) -> Any:
    """Fail fast on the retired Ax config surface."""
    if not isinstance(raw_value, dict):
        return raw_value

    legacy_fields = [
        field_name
        for field_name in ("surrogate_model", "acquisition_function")
        if field_name in raw_value
    ]
    if legacy_fields:
        joined = ", ".join(legacy_fields)
        raise ValueError(
            f"AxOptimizerConfig no longer accepts legacy fields: {joined}. "
            "Use 'generator', 'generator_kwargs', and 'generator_gen_kwargs' "
            "instead."
        )
    return raw_value

BaseOptimizer

Bases: ABC

Abstract base class for all AID2E optimizers.

Subclasses must implement the abstract methods to suggest candidates, ingest evaluation results, and surface optimizer state. The interface is intentionally minimal to support a range of backends (Ax, genetic algorithms, grid search) while keeping a consistent contract for the rest of the framework.

Concrete default implementations are provided for get_pareto_front and get_best_trial using :func:compute_pareto_front. Individual backends may override these when native support (e.g. PyMOO's built-in Pareto tools) is preferable.

Source code in src/aid2e/optimizers/base.py

class BaseOptimizer(ABC):
    """Abstract base class for all AID2E optimizers.

    Subclasses must implement the abstract methods to suggest candidates, ingest
    evaluation results, and surface optimizer state. The interface is intentionally
    minimal to support a range of backends (Ax, genetic algorithms, grid search)
    while keeping a consistent contract for the rest of the framework.

    Concrete default implementations are provided for ``get_pareto_front`` and
    ``get_best_trial`` using :func:`compute_pareto_front`.  Individual backends
    may override these when native support (e.g. PyMOO's built-in Pareto tools)
    is preferable.
    """

    def __init__(
        self,
        search_space: Union[SearchSpace, DesignConfig],
        objective_names: List[str],
        seed: Optional[int] = None,
        objective_directions: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Initialize the optimizer with a search space and objective specification.

        Args:
            search_space: Typed search space or DesignConfig to optimize over.
            objective_names: Ordered list of objective metric names.  Each name
                must match the keys returned in the ``metrics`` dict when
                ``update_with_results`` is called.
            seed: Optional integer seed for reproducibility.
            objective_directions: Optional mapping from objective name to
                ``"minimize"`` or ``"maximize"``. Missing objectives default to
                minimization.

        Raises:
            ValueError: If the search space is empty or no objective names are
                provided.
        """

        resolved_space = (
            SearchSpace.from_design_config(search_space)
            if isinstance(search_space, DesignConfig)
            else search_space
        )

        if not resolved_space.parameters:
            raise ValueError("Search space cannot be empty")
        if not objective_names:
            raise ValueError("objective_names must contain at least one name")

        self.search_space = resolved_space
        self.objective_names: List[str] = list(objective_names)
        self.objective_directions: Dict[str, Any] = dict(objective_directions or {})
        self.seed = seed

        # Shared trial history managed by the base class.
        # All backends read and write through self._trials and self._trial_counter
        # so utilities like seed_from_trials, get_trials, get_pareto_front work
        # uniformly regardless of the backend.
        self._trials: List[Optional[Trial]] = []
        self._trial_counter: int = 0

    @property
    def n_objectives(self) -> int:
        """Return the number of optimisation objectives.

        Returns:
            Integer count derived from ``objective_names``.
        """
        return len(self.objective_names)

    @abstractmethod
    def suggest_candidates(self, n_candidates: int = 1) -> List[Dict[str, Any]]:
        """Suggest next parameter configurations to evaluate.

        Args:
            n_candidates: Number of candidates to suggest.

        Returns:
            List of parameter dictionaries, where each dictionary maps
            parameter names to their suggested values.

        Raises:
            RuntimeError: If optimizer is not properly initialized.

        Examples:
            >>> candidates = optimizer.suggest_candidates(n_candidates=5)
            >>> candidates[0]
            {'x': 0.5, 'y': 0.3, 'z': 2.1}

        Notes:
            The implementation should use the configured strategy
            (e.g., Sobol sampling, Bayesian optimization, genetic algorithms).
        """
        pass

    @abstractmethod
    def update_with_results(
        self,
        trial_index: int,
        parameters: Dict[str, Any],
        metrics: Dict[str, float]
    ) -> None:
        """Update optimizer with evaluation results from a trial.

        Args:
            trial_index: Unique identifier for the trial.
            parameters: Parameter values that were evaluated.
            metrics: Objective values obtained from evaluation.
                Keys are metric names, values are metric values.

        Raises:
            ValueError: If metrics don't match expected objectives.

        Examples:
            >>> optimizer.update_with_results(
            ...     trial_index=0,
            ...     parameters={'x': 0.5, 'y': 0.3},
            ...     metrics={'loss': 0.1, 'accuracy': 0.9}
            ... )

        Notes:
            After updating, the optimizer can use this information to
            suggest better candidates in subsequent calls to suggest_candidates().
        """
        pass

    def get_trials(self) -> List[Trial]:
        """Return all recorded trials (pending, completed, and failed).

        The base implementation reads directly from ``self._trials``, which is
        owned by ``BaseOptimizer`` and kept up to date by every backend.
        Backends that maintain additional internal state may override this to
        include synthetic or reconstructed trials, but doing so is uncommon.

        Returns:
            List of non-``None`` Trial objects in creation order.

        Examples:
            >>> done = [t for t in optimizer.get_trials() if t.status == "completed"]
        """
        return [t for t in self._trials if t is not None]

    def set_trial_status(
        self,
        trial_index: int,
        status: str,
        *,
        parameters: Optional[Dict[str, Any]] = None,
        metrics: Optional[Dict[str, float]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Trial:
        """Create or update a trial entry with a new lifecycle status.

        Args:
            trial_index: Unique trial identifier.
            status: Trial lifecycle status (for example ``running``,
                ``completed``, ``aborted``).
            parameters: Optional parameter dictionary to store on the trial.
            metrics: Optional objective dictionary to store on the trial.
            metadata: Optional metadata to merge into existing metadata.

        Returns:
            The updated Trial object.

        Raises:
            ValueError: If ``trial_index`` is negative.
        """
        if trial_index < 0:
            raise ValueError("trial_index must be >= 0")

        normalized_status = str(status).strip().lower()
        while len(self._trials) <= trial_index:
            self._trials.append(None)

        existing = self._trials[trial_index]
        existing_parameters = existing.parameters if existing else {}
        existing_metrics = existing.metrics if existing else None
        existing_metadata = dict(existing.metadata) if existing and existing.metadata else {}

        trial = Trial(
            index=trial_index,
            parameters=parameters if parameters is not None else existing_parameters,
            metrics=metrics if metrics is not None else existing_metrics,
            status=normalized_status,
            metadata={**existing_metadata, **(metadata or {})},
        )
        self._trials[trial_index] = trial
        self._trial_counter = max(self._trial_counter, trial_index + 1)
        return trial

    def mark_trial_failed(
        self,
        trial_index: int,
        *,
        parameters: Optional[Dict[str, Any]] = None,
        reason: Optional[str] = None,
    ) -> Trial:
        """Record a failed evaluation without objective values."""
        return self.set_trial_status(
            trial_index,
            TRIAL_STATUS_FAILED,
            parameters=parameters,
            metadata={"reason": reason} if reason else None,
        )

    def get_optimization_results(
        self,
        errors_by_trial: Optional[Dict[int, Dict[str, float]]] = None,
    ) -> Dict[str, Any]:
        """Return a normalized optimization-results payload.

        Args:
            errors_by_trial: Optional uncertainty/error fields keyed by trial index.

        Returns:
            Dictionary containing objective names and trial records with
            parameters, metrics, and both raw and display status labels.
        """
        errors_by_trial = errors_by_trial or {}
        trials_payload: List[Dict[str, Any]] = []
        for trial in self.get_trials():
            trials_payload.append(
                {
                    "trial_index": trial.index,
                    "status": trial.status,
                    "display_status": DISPLAY_STATUS_MAP.get(
                        trial.status,
                        trial.status.title(),
                    ),
                    "design_parameters": dict(trial.parameters or {}),
                    "objectives": dict(trial.metrics or {}),
                    "objective_errors": dict(errors_by_trial.get(trial.index, {})),
                    "metadata": dict(trial.metadata or {}),
                }
            )

        return {
            "objective_names": list(self.objective_names),
            "n_objectives": self.n_objectives,
            "n_trials": len(trials_payload),
            "trials": trials_payload,
        }

    def save_optimization_results(
        self,
        output_path: Union[str, Path],
        *,
        errors_by_trial: Optional[Dict[int, Dict[str, float]]] = None,
    ) -> Dict[str, Path]:
        """Write optimization results to disk as pretty-printed JSON.

        Args:
            output_path: Directory where optimizer result files are written.
            errors_by_trial: Optional uncertainty/error fields keyed by trial index.

        Returns:
            Paths of the written result files.
        """
        output_dir = Path(output_path)
        output_dir.mkdir(parents=True, exist_ok=True)
        results_path = output_dir / "optimization_results.json"
        pareto_path = output_dir / "pareto_front.json"

        payload = self.get_optimization_results(errors_by_trial=errors_by_trial)
        with results_path.open("w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2, sort_keys=True)

        errors_by_trial = errors_by_trial or {}
        pareto_payload = [
            {
                "trial_index": trial.index,
                "design_parameters": dict(trial.parameters or {}),
                "objectives": dict(trial.metrics or {}),
                "objective_errors": dict(errors_by_trial.get(trial.index, {})),
                "metadata": dict(trial.metadata or {}),
            }
            for trial in self.get_pareto_front()
        ]
        with pareto_path.open("w", encoding="utf-8") as handle:
            json.dump(pareto_payload, handle, indent=2)
        return {
            "optimization_results": results_path,
            "pareto_front": pareto_path,
        }

    def seed_from_trials(
        self,
        trials: List[Trial],
        *,
        only_completed: bool = True,
    ) -> int:
        """Inject external trials into the optimizer history without advancing the algorithm.

        This is the *backend-switch primitive*.  It lets you:

        - Seed a new optimizer with results from a previous one (e.g. random
          init → MOEA → BO transition).
        - Inject prior knowledge before the first ``suggest_candidates`` call.
        - Resume an optimisation from a checkpoint produced by a *different*
          backend.

        Trials are appended to the internal ``_trials`` list with freshly
        assigned sequential indices (starting from the current
        ``_trial_counter``).  The original ``trial.index`` values from the
        source optimizer are preserved in each trial's ``metadata`` under the
        key ``"source_index"``.

        Args:
            trials: Iterable of Trial objects to inject.  The list may contain
                ``None`` placeholders (they are silently skipped).
            only_completed: When ``True`` (default), only trials whose
                ``status`` is ``"completed"`` are imported.  Set to ``False``
                to also import ``"pending"`` or ``"failed"`` trials.

        Returns:
            Number of trials actually injected.

        Examples:
            >>> # Transfer best results from a random-search warmup:
            >>> warmup_trials = random_opt.get_trials()
            >>> pymoo_opt.seed_from_trials(warmup_trials)
            100
            >>> # Now start MOEA generation — the history already has 100 points
            >>> candidates = pymoo_opt.suggest_candidates()

        Notes:
            - This method does NOT advance PyMOO's (or any other backend's)
              internal population.  The injected trials are purely visible in
              the history for Pareto-front and best-trial queries.
            - Backends that want to warm-start their internal state from these
              trials should override this method.
        """
        accepted = 0
        for trial in trials:
            if trial is None:
                continue
            if only_completed and trial.status != "completed":
                continue
            new_idx = self._trial_counter
            seeded_trial = Trial(
                index=new_idx,
                parameters=trial.parameters,
                metrics=trial.metrics,
                status=trial.status,
                metadata={**(trial.metadata or {}), "source_index": trial.index},
            )
            while len(self._trials) <= new_idx:
                self._trials.append(None)
            self._trials[new_idx] = seeded_trial
            self._trial_counter += 1
            accepted += 1

        if accepted:
            logger.debug(
                "seed_from_trials: injected %d trial(s) (total history: %d).",
                accepted,
                self._trial_counter,
            )
        return accepted

    def get_pareto_front(self) -> List[Trial]:
        """Retrieve the current Pareto front of non-dominated solutions.

        The default implementation delegates to :func:`compute_pareto_front`,
        which operates on the trials returned by :meth:`get_trials`.  Backends
        with native Pareto support (e.g. PyMOO) may override this method to
        expose more detailed Pareto metadata.

        Returns:
            List of Trial objects representing Pareto-optimal solutions.
            For single-objective optimisation, returns the single trial with
            the best metric value under its configured direction. Returns an
            empty list when no completed trials are available.

        Examples:
            >>> pareto_front = optimizer.get_pareto_front()
            >>> for trial in pareto_front:
            ...     print(f"Params: {trial.parameters}, Metrics: {trial.metrics}")

        Notes:
            Objective directions are read from ``self.objective_directions``.
        """
        return compute_pareto_front(
            self.get_trials(),
            self.objective_names,
            self.objective_directions,
        )

    def get_best_trial(self) -> Optional[Trial]:
        """Get the best trial found so far.

        For single-objective optimisation, returns the completed trial with the
        best metric value under its configured direction. For multi-objective,
        returns the first trial from the Pareto front (arbitrary representative;
        use :meth:`get_pareto_front` for the full front).

        Returns:
            Best Trial, or ``None`` if no completed trials exist.

        Examples:
            >>> best = optimizer.get_best_trial()
            >>> if best:
            ...     print(f"Best parameters: {best.parameters}")
            ...     print(f"Best metrics: {best.metrics}")
        """
        front = self.get_pareto_front()
        if not front:
            return None
        if self.n_objectives == 1:
            obj = self.objective_names[0]
            return min(front, key=lambda t: t.metrics[obj])
        # Multi-objective: return first element of the front
        return front[0]

    @abstractmethod
    def serialize_state(self) -> Dict[str, Any]:
        """Serialize optimizer state for distributed execution or checkpointing.

        Returns:
            Dictionary containing all necessary state to reconstruct
            the optimizer. Should be JSON-serializable.

        Examples:
            >>> state = optimizer.serialize_state()
            >>> import json
            >>> with open('optimizer_state.json', 'w') as f:
            ...     json.dump(state, f)

        Notes:
            This is crucial for distributed optimization where optimizer
            state needs to be shared across workers or checkpointed.
        """
        pass

    @abstractmethod
    def load_state(self, state: Dict[str, Any]) -> None:
        """Load optimizer state from serialized form.

        Args:
            state: Dictionary containing serialized optimizer state,
                as returned by serialize_state().

        Raises:
            ValueError: If state is invalid or incompatible.

        Examples:
            >>> import json
            >>> with open('optimizer_state.json', 'r') as f:
            ...     state = json.load(f)
            >>> optimizer.load_state(state)

        Notes:
            After loading state, the optimizer should be able to continue
            optimization as if it never stopped.
        """
        pass

    def __repr__(self) -> str:
        """Return string representation of the optimizer.

        Returns:
            String describing the optimizer configuration.
        """
        return (
            f"{self.__class__.__name__}("
            f"n_params={len(self.search_space.parameters)}, "
            f"n_objectives={self.n_objectives}, "
            f"seed={self.seed}"
            f")"
        )

n_objectives property

Return the number of optimisation objectives.

Returns:

Type	Description
int	Integer count derived from objective_names.

__init__(search_space, objective_names, seed=None, objective_directions=None)

Initialize the optimizer with a search space and objective specification.

Parameters:

Name	Type	Description	Default
search_space	Union[SearchSpace, DesignConfig]	Typed search space or DesignConfig to optimize over.	required
objective_names	List[str]	Ordered list of objective metric names. Each name must match the keys returned in the metrics dict when update_with_results is called.	required
seed	Optional[int]	Optional integer seed for reproducibility.	None
objective_directions	Optional[Dict[str, Any]]	Optional mapping from objective name to "minimize" or "maximize". Missing objectives default to minimization.	None

Raises:

Type	Description
ValueError	If the search space is empty or no objective names are provided.

Source code in src/aid2e/optimizers/base.py

def __init__(
    self,
    search_space: Union[SearchSpace, DesignConfig],
    objective_names: List[str],
    seed: Optional[int] = None,
    objective_directions: Optional[Dict[str, Any]] = None,
) -> None:
    """Initialize the optimizer with a search space and objective specification.

    Args:
        search_space: Typed search space or DesignConfig to optimize over.
        objective_names: Ordered list of objective metric names.  Each name
            must match the keys returned in the ``metrics`` dict when
            ``update_with_results`` is called.
        seed: Optional integer seed for reproducibility.
        objective_directions: Optional mapping from objective name to
            ``"minimize"`` or ``"maximize"``. Missing objectives default to
            minimization.

    Raises:
        ValueError: If the search space is empty or no objective names are
            provided.
    """

    resolved_space = (
        SearchSpace.from_design_config(search_space)
        if isinstance(search_space, DesignConfig)
        else search_space
    )

    if not resolved_space.parameters:
        raise ValueError("Search space cannot be empty")
    if not objective_names:
        raise ValueError("objective_names must contain at least one name")

    self.search_space = resolved_space
    self.objective_names: List[str] = list(objective_names)
    self.objective_directions: Dict[str, Any] = dict(objective_directions or {})
    self.seed = seed

    # Shared trial history managed by the base class.
    # All backends read and write through self._trials and self._trial_counter
    # so utilities like seed_from_trials, get_trials, get_pareto_front work
    # uniformly regardless of the backend.
    self._trials: List[Optional[Trial]] = []
    self._trial_counter: int = 0

__repr__()

Return string representation of the optimizer.

Returns:

Type	Description
str	String describing the optimizer configuration.

Source code in src/aid2e/optimizers/base.py

def __repr__(self) -> str:
    """Return string representation of the optimizer.

    Returns:
        String describing the optimizer configuration.
    """
    return (
        f"{self.__class__.__name__}("
        f"n_params={len(self.search_space.parameters)}, "
        f"n_objectives={self.n_objectives}, "
        f"seed={self.seed}"
        f")"
    )

get_best_trial()

Get the best trial found so far.

For single-objective optimisation, returns the completed trial with the best metric value under its configured direction. For multi-objective, returns the first trial from the Pareto front (arbitrary representative; use :meth:get_pareto_front for the full front).

Returns:

Type	Description
Optional[Trial]	Best Trial, or None if no completed trials exist.

Examples:

>>> best = optimizer.get_best_trial()
>>> if best:
...     print(f"Best parameters: {best.parameters}")
...     print(f"Best metrics: {best.metrics}")

Source code in src/aid2e/optimizers/base.py

def get_best_trial(self) -> Optional[Trial]:
    """Get the best trial found so far.

    For single-objective optimisation, returns the completed trial with the
    best metric value under its configured direction. For multi-objective,
    returns the first trial from the Pareto front (arbitrary representative;
    use :meth:`get_pareto_front` for the full front).

    Returns:
        Best Trial, or ``None`` if no completed trials exist.

    Examples:
        >>> best = optimizer.get_best_trial()
        >>> if best:
        ...     print(f"Best parameters: {best.parameters}")
        ...     print(f"Best metrics: {best.metrics}")
    """
    front = self.get_pareto_front()
    if not front:
        return None
    if self.n_objectives == 1:
        obj = self.objective_names[0]
        return min(front, key=lambda t: t.metrics[obj])
    # Multi-objective: return first element of the front
    return front[0]

get_optimization_results(errors_by_trial=None)

Return a normalized optimization-results payload.

Parameters:

Name	Type	Description	Default
errors_by_trial	Optional[Dict[int, Dict[str, float]]]	Optional uncertainty/error fields keyed by trial index.	None

Returns:

Type	Description
Dict[str, Any]	Dictionary containing objective names and trial records with
Dict[str, Any]	parameters, metrics, and both raw and display status labels.

Source code in src/aid2e/optimizers/base.py

def get_optimization_results(
    self,
    errors_by_trial: Optional[Dict[int, Dict[str, float]]] = None,
) -> Dict[str, Any]:
    """Return a normalized optimization-results payload.

    Args:
        errors_by_trial: Optional uncertainty/error fields keyed by trial index.

    Returns:
        Dictionary containing objective names and trial records with
        parameters, metrics, and both raw and display status labels.
    """
    errors_by_trial = errors_by_trial or {}
    trials_payload: List[Dict[str, Any]] = []
    for trial in self.get_trials():
        trials_payload.append(
            {
                "trial_index": trial.index,
                "status": trial.status,
                "display_status": DISPLAY_STATUS_MAP.get(
                    trial.status,
                    trial.status.title(),
                ),
                "design_parameters": dict(trial.parameters or {}),
                "objectives": dict(trial.metrics or {}),
                "objective_errors": dict(errors_by_trial.get(trial.index, {})),
                "metadata": dict(trial.metadata or {}),
            }
        )

    return {
        "objective_names": list(self.objective_names),
        "n_objectives": self.n_objectives,
        "n_trials": len(trials_payload),
        "trials": trials_payload,
    }

get_pareto_front()

Retrieve the current Pareto front of non-dominated solutions.

The default implementation delegates to :func:compute_pareto_front, which operates on the trials returned by :meth:get_trials. Backends with native Pareto support (e.g. PyMOO) may override this method to expose more detailed Pareto metadata.

Returns:

Type	Description
List[Trial]	List of Trial objects representing Pareto-optimal solutions.
List[Trial]	For single-objective optimisation, returns the single trial with
List[Trial]	the best metric value under its configured direction. Returns an
List[Trial]	empty list when no completed trials are available.

Examples:

>>> pareto_front = optimizer.get_pareto_front()
>>> for trial in pareto_front:
...     print(f"Params: {trial.parameters}, Metrics: {trial.metrics}")

Notes

Objective directions are read from self.objective_directions.

Source code in src/aid2e/optimizers/base.py

def get_pareto_front(self) -> List[Trial]:
    """Retrieve the current Pareto front of non-dominated solutions.

    The default implementation delegates to :func:`compute_pareto_front`,
    which operates on the trials returned by :meth:`get_trials`.  Backends
    with native Pareto support (e.g. PyMOO) may override this method to
    expose more detailed Pareto metadata.

    Returns:
        List of Trial objects representing Pareto-optimal solutions.
        For single-objective optimisation, returns the single trial with
        the best metric value under its configured direction. Returns an
        empty list when no completed trials are available.

    Examples:
        >>> pareto_front = optimizer.get_pareto_front()
        >>> for trial in pareto_front:
        ...     print(f"Params: {trial.parameters}, Metrics: {trial.metrics}")

    Notes:
        Objective directions are read from ``self.objective_directions``.
    """
    return compute_pareto_front(
        self.get_trials(),
        self.objective_names,
        self.objective_directions,
    )

get_trials()

Return all recorded trials (pending, completed, and failed).

The base implementation reads directly from self._trials, which is owned by BaseOptimizer and kept up to date by every backend. Backends that maintain additional internal state may override this to include synthetic or reconstructed trials, but doing so is uncommon.

Returns:

Type	Description
List[Trial]	List of non-None Trial objects in creation order.

Examples:

>>> done = [t for t in optimizer.get_trials() if t.status == "completed"]

Source code in src/aid2e/optimizers/base.py

def get_trials(self) -> List[Trial]:
    """Return all recorded trials (pending, completed, and failed).

    The base implementation reads directly from ``self._trials``, which is
    owned by ``BaseOptimizer`` and kept up to date by every backend.
    Backends that maintain additional internal state may override this to
    include synthetic or reconstructed trials, but doing so is uncommon.

    Returns:
        List of non-``None`` Trial objects in creation order.

    Examples:
        >>> done = [t for t in optimizer.get_trials() if t.status == "completed"]
    """
    return [t for t in self._trials if t is not None]

load_state(state) abstractmethod

Load optimizer state from serialized form.

Parameters:

Name	Type	Description	Default
state	Dict[str, Any]	Dictionary containing serialized optimizer state, as returned by serialize_state().	required

Raises:

Type	Description
ValueError	If state is invalid or incompatible.

Examples:

>>> import json
>>> with open('optimizer_state.json', 'r') as f:
...     state = json.load(f)
>>> optimizer.load_state(state)

Notes

After loading state, the optimizer should be able to continue optimization as if it never stopped.

Source code in src/aid2e/optimizers/base.py

@abstractmethod
def load_state(self, state: Dict[str, Any]) -> None:
    """Load optimizer state from serialized form.

    Args:
        state: Dictionary containing serialized optimizer state,
            as returned by serialize_state().

    Raises:
        ValueError: If state is invalid or incompatible.

    Examples:
        >>> import json
        >>> with open('optimizer_state.json', 'r') as f:
        ...     state = json.load(f)
        >>> optimizer.load_state(state)

    Notes:
        After loading state, the optimizer should be able to continue
        optimization as if it never stopped.
    """
    pass

mark_trial_failed(trial_index, *, parameters=None, reason=None)

Record a failed evaluation without objective values.

Source code in src/aid2e/optimizers/base.py

def mark_trial_failed(
    self,
    trial_index: int,
    *,
    parameters: Optional[Dict[str, Any]] = None,
    reason: Optional[str] = None,
) -> Trial:
    """Record a failed evaluation without objective values."""
    return self.set_trial_status(
        trial_index,
        TRIAL_STATUS_FAILED,
        parameters=parameters,
        metadata={"reason": reason} if reason else None,
    )

save_optimization_results(output_path, *, errors_by_trial=None)

Write optimization results to disk as pretty-printed JSON.

Parameters:

Name	Type	Description	Default
output_path	Union[str, Path]	Directory where optimizer result files are written.	required
errors_by_trial	Optional[Dict[int, Dict[str, float]]]	Optional uncertainty/error fields keyed by trial index.	None

Returns:

Type	Description
Dict[str, Path]	Paths of the written result files.

Source code in src/aid2e/optimizers/base.py

def save_optimization_results(
    self,
    output_path: Union[str, Path],
    *,
    errors_by_trial: Optional[Dict[int, Dict[str, float]]] = None,
) -> Dict[str, Path]:
    """Write optimization results to disk as pretty-printed JSON.

    Args:
        output_path: Directory where optimizer result files are written.
        errors_by_trial: Optional uncertainty/error fields keyed by trial index.

    Returns:
        Paths of the written result files.
    """
    output_dir = Path(output_path)
    output_dir.mkdir(parents=True, exist_ok=True)
    results_path = output_dir / "optimization_results.json"
    pareto_path = output_dir / "pareto_front.json"

    payload = self.get_optimization_results(errors_by_trial=errors_by_trial)
    with results_path.open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2, sort_keys=True)

    errors_by_trial = errors_by_trial or {}
    pareto_payload = [
        {
            "trial_index": trial.index,
            "design_parameters": dict(trial.parameters or {}),
            "objectives": dict(trial.metrics or {}),
            "objective_errors": dict(errors_by_trial.get(trial.index, {})),
            "metadata": dict(trial.metadata or {}),
        }
        for trial in self.get_pareto_front()
    ]
    with pareto_path.open("w", encoding="utf-8") as handle:
        json.dump(pareto_payload, handle, indent=2)
    return {
        "optimization_results": results_path,
        "pareto_front": pareto_path,
    }

seed_from_trials(trials, *, only_completed=True)

Inject external trials into the optimizer history without advancing the algorithm.

This is the backend-switch primitive. It lets you:

• Seed a new optimizer with results from a previous one (e.g. random init → MOEA → BO transition).
• Inject prior knowledge before the first suggest_candidates call.
• Resume an optimisation from a checkpoint produced by a different backend.

Trials are appended to the internal _trials list with freshly assigned sequential indices (starting from the current _trial_counter). The original trial.index values from the source optimizer are preserved in each trial's metadata under the key "source_index".

Parameters:

Name	Type	Description	Default
trials	List[Trial]	Iterable of Trial objects to inject. The list may contain None placeholders (they are silently skipped).	required
only_completed	bool	When True (default), only trials whose status is "completed" are imported. Set to False to also import "pending" or "failed" trials.	True

Returns:

Type	Description
int	Number of trials actually injected.

Examples:

>>> # Transfer best results from a random-search warmup:
>>> warmup_trials = random_opt.get_trials()
>>> pymoo_opt.seed_from_trials(warmup_trials)
100
>>> # Now start MOEA generation — the history already has 100 points
>>> candidates = pymoo_opt.suggest_candidates()

Notes

• This method does NOT advance PyMOO's (or any other backend's) internal population. The injected trials are purely visible in the history for Pareto-front and best-trial queries.
• Backends that want to warm-start their internal state from these trials should override this method.

Source code in src/aid2e/optimizers/base.py

def seed_from_trials(
    self,
    trials: List[Trial],
    *,
    only_completed: bool = True,
) -> int:
    """Inject external trials into the optimizer history without advancing the algorithm.

    This is the *backend-switch primitive*.  It lets you:

    - Seed a new optimizer with results from a previous one (e.g. random
      init → MOEA → BO transition).
    - Inject prior knowledge before the first ``suggest_candidates`` call.
    - Resume an optimisation from a checkpoint produced by a *different*
      backend.

    Trials are appended to the internal ``_trials`` list with freshly
    assigned sequential indices (starting from the current
    ``_trial_counter``).  The original ``trial.index`` values from the
    source optimizer are preserved in each trial's ``metadata`` under the
    key ``"source_index"``.

    Args:
        trials: Iterable of Trial objects to inject.  The list may contain
            ``None`` placeholders (they are silently skipped).
        only_completed: When ``True`` (default), only trials whose
            ``status`` is ``"completed"`` are imported.  Set to ``False``
            to also import ``"pending"`` or ``"failed"`` trials.

    Returns:
        Number of trials actually injected.

    Examples:
        >>> # Transfer best results from a random-search warmup:
        >>> warmup_trials = random_opt.get_trials()
        >>> pymoo_opt.seed_from_trials(warmup_trials)
        100
        >>> # Now start MOEA generation — the history already has 100 points
        >>> candidates = pymoo_opt.suggest_candidates()

    Notes:
        - This method does NOT advance PyMOO's (or any other backend's)
          internal population.  The injected trials are purely visible in
          the history for Pareto-front and best-trial queries.
        - Backends that want to warm-start their internal state from these
          trials should override this method.
    """
    accepted = 0
    for trial in trials:
        if trial is None:
            continue
        if only_completed and trial.status != "completed":
            continue
        new_idx = self._trial_counter
        seeded_trial = Trial(
            index=new_idx,
            parameters=trial.parameters,
            metrics=trial.metrics,
            status=trial.status,
            metadata={**(trial.metadata or {}), "source_index": trial.index},
        )
        while len(self._trials) <= new_idx:
            self._trials.append(None)
        self._trials[new_idx] = seeded_trial
        self._trial_counter += 1
        accepted += 1

    if accepted:
        logger.debug(
            "seed_from_trials: injected %d trial(s) (total history: %d).",
            accepted,
            self._trial_counter,
        )
    return accepted

serialize_state() abstractmethod

Serialize optimizer state for distributed execution or checkpointing.

Returns:

Type	Description
Dict[str, Any]	Dictionary containing all necessary state to reconstruct
Dict[str, Any]	the optimizer. Should be JSON-serializable.

Examples:

>>> state = optimizer.serialize_state()
>>> import json
>>> with open('optimizer_state.json', 'w') as f:
...     json.dump(state, f)

Notes

This is crucial for distributed optimization where optimizer state needs to be shared across workers or checkpointed.

Source code in src/aid2e/optimizers/base.py

@abstractmethod
def serialize_state(self) -> Dict[str, Any]:
    """Serialize optimizer state for distributed execution or checkpointing.

    Returns:
        Dictionary containing all necessary state to reconstruct
        the optimizer. Should be JSON-serializable.

    Examples:
        >>> state = optimizer.serialize_state()
        >>> import json
        >>> with open('optimizer_state.json', 'w') as f:
        ...     json.dump(state, f)

    Notes:
        This is crucial for distributed optimization where optimizer
        state needs to be shared across workers or checkpointed.
    """
    pass

set_trial_status(trial_index, status, *, parameters=None, metrics=None, metadata=None)

Create or update a trial entry with a new lifecycle status.

Parameters:

Name	Type	Description	Default
trial_index	int	Unique trial identifier.	required
status	str	Trial lifecycle status (for example running, completed, aborted).	required
parameters	Optional[Dict[str, Any]]	Optional parameter dictionary to store on the trial.	None
metrics	Optional[Dict[str, float]]	Optional objective dictionary to store on the trial.	None
metadata	Optional[Dict[str, Any]]	Optional metadata to merge into existing metadata.	None

Returns:

Type	Description
Trial	The updated Trial object.

Raises:

Type	Description
ValueError	If trial_index is negative.

Source code in src/aid2e/optimizers/base.py

def set_trial_status(
    self,
    trial_index: int,
    status: str,
    *,
    parameters: Optional[Dict[str, Any]] = None,
    metrics: Optional[Dict[str, float]] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> Trial:
    """Create or update a trial entry with a new lifecycle status.

    Args:
        trial_index: Unique trial identifier.
        status: Trial lifecycle status (for example ``running``,
            ``completed``, ``aborted``).
        parameters: Optional parameter dictionary to store on the trial.
        metrics: Optional objective dictionary to store on the trial.
        metadata: Optional metadata to merge into existing metadata.

    Returns:
        The updated Trial object.

    Raises:
        ValueError: If ``trial_index`` is negative.
    """
    if trial_index < 0:
        raise ValueError("trial_index must be >= 0")

    normalized_status = str(status).strip().lower()
    while len(self._trials) <= trial_index:
        self._trials.append(None)

    existing = self._trials[trial_index]
    existing_parameters = existing.parameters if existing else {}
    existing_metrics = existing.metrics if existing else None
    existing_metadata = dict(existing.metadata) if existing and existing.metadata else {}

    trial = Trial(
        index=trial_index,
        parameters=parameters if parameters is not None else existing_parameters,
        metrics=metrics if metrics is not None else existing_metrics,
        status=normalized_status,
        metadata={**existing_metadata, **(metadata or {})},
    )
    self._trials[trial_index] = trial
    self._trial_counter = max(self._trial_counter, trial_index + 1)
    return trial

suggest_candidates(n_candidates=1) abstractmethod

Suggest next parameter configurations to evaluate.

Parameters:

Name	Type	Description	Default
n_candidates	int	Number of candidates to suggest.	1

Returns:

Type	Description
List[Dict[str, Any]]	List of parameter dictionaries, where each dictionary maps
List[Dict[str, Any]]	parameter names to their suggested values.

Raises:

Type	Description
RuntimeError	If optimizer is not properly initialized.

Examples:

>>> candidates = optimizer.suggest_candidates(n_candidates=5)
>>> candidates[0]
{'x': 0.5, 'y': 0.3, 'z': 2.1}

Notes

The implementation should use the configured strategy (e.g., Sobol sampling, Bayesian optimization, genetic algorithms).

Source code in src/aid2e/optimizers/base.py

@abstractmethod
def suggest_candidates(self, n_candidates: int = 1) -> List[Dict[str, Any]]:
    """Suggest next parameter configurations to evaluate.

    Args:
        n_candidates: Number of candidates to suggest.

    Returns:
        List of parameter dictionaries, where each dictionary maps
        parameter names to their suggested values.

    Raises:
        RuntimeError: If optimizer is not properly initialized.

    Examples:
        >>> candidates = optimizer.suggest_candidates(n_candidates=5)
        >>> candidates[0]
        {'x': 0.5, 'y': 0.3, 'z': 2.1}

    Notes:
        The implementation should use the configured strategy
        (e.g., Sobol sampling, Bayesian optimization, genetic algorithms).
    """
    pass

update_with_results(trial_index, parameters, metrics) abstractmethod

Update optimizer with evaluation results from a trial.

Parameters:

Name	Type	Description	Default
trial_index	int	Unique identifier for the trial.	required
parameters	Dict[str, Any]	Parameter values that were evaluated.	required
metrics	Dict[str, float]	Objective values obtained from evaluation. Keys are metric names, values are metric values.	required

Raises:

Type	Description
ValueError	If metrics don't match expected objectives.

Examples:

>>> optimizer.update_with_results(
...     trial_index=0,
...     parameters={'x': 0.5, 'y': 0.3},
...     metrics={'loss': 0.1, 'accuracy': 0.9}
... )

Notes

After updating, the optimizer can use this information to suggest better candidates in subsequent calls to suggest_candidates().

Source code in src/aid2e/optimizers/base.py

@abstractmethod
def update_with_results(
    self,
    trial_index: int,
    parameters: Dict[str, Any],
    metrics: Dict[str, float]
) -> None:
    """Update optimizer with evaluation results from a trial.

    Args:
        trial_index: Unique identifier for the trial.
        parameters: Parameter values that were evaluated.
        metrics: Objective values obtained from evaluation.
            Keys are metric names, values are metric values.

    Raises:
        ValueError: If metrics don't match expected objectives.

    Examples:
        >>> optimizer.update_with_results(
        ...     trial_index=0,
        ...     parameters={'x': 0.5, 'y': 0.3},
        ...     metrics={'loss': 0.1, 'accuracy': 0.9}
        ... )

    Notes:
        After updating, the optimizer can use this information to
        suggest better candidates in subsequent calls to suggest_candidates().
    """
    pass

PyMOOOptimizer

Bases: BaseOptimizer

Evolutionary optimizer backed by PyMOO algorithms.

Implements the BaseOptimizer interface using PyMOO's ask/tell protocol so that evaluations can be performed externally (e.g. via schedulers or simulation jobs) without blocking PyMOO's internal loop.

The generation lifecycle is::

candidates = optimizer.suggest_candidates()   # calls algorithm.ask()
for i, c in enumerate(candidates):
    metrics = my_evaluate(c)
    optimizer.update_with_results(i, c, metrics)
# After the last call above, algorithm.tell() is invoked automatically.
# The next suggest_candidates() call produces the next generation.

Supported algorithms

• "ga" — Genetic Algorithm for single-objective optimisation.
• "nsga2" — NSGA-II, recommended for 2–3 objectives.
• "nsga3" — NSGA-III, recommended for 3+ objectives.
• "moead" — MOEA/D, weight-decomposition approach.

Attributes:

Name	Type	Description
config		PyMOOOptimizerConfig instance controlling algorithm behaviour.
n_gen_completed	int	Number of generations that have been fully evaluated.

Examples:

>>> from aid2e.optimizers.pymoo import PyMOOOptimizer, PyMOOOptimizerConfig
>>> from aid2e.optimizers.base import SearchSpace
>>> space = SearchSpace(
...     parameters={
...         "x": {"type": "range", "bounds": [0.0, 1.0]},
...         "y": {"type": "range", "bounds": [0.0, 1.0]},
...     }
... )
>>> config = PyMOOOptimizerConfig(pop_size=20, seed=0)
>>> opt = PyMOOOptimizer(space, config, objective_names=["loss"])
>>> candidates = opt.suggest_candidates()
>>> for trial_idx, c in enumerate(candidates):
...     opt.update_with_results(trial_idx, c, {"loss": c["x"] + c["y"]})
>>> best = opt.get_best_trial()

Notes

• Objective directions are translated to PyMOO's minimization convention.
• Linear parameter constraints are not yet forwarded to PyMOO. A warning is emitted when the search space contains constraints.
• n_candidates passed to suggest_candidates is informational only; the actual batch size is determined by the algorithm (pop_size for the first generation, n_offsprings thereafter).

Project: AID2E v0.0.0 — AI assisted Detector Design for EIC Homepage: https://aid2e.github.io/AID2E-framework Repository: https://github.com/aid2e/AID2E-framework.git

Source code in src/aid2e/optimizers/pymoo/optimizer.py

class PyMOOOptimizer(BaseOptimizer):
    """Evolutionary optimizer backed by PyMOO algorithms.

    Implements the ``BaseOptimizer`` interface using PyMOO's ask/tell protocol
    so that evaluations can be performed externally (e.g. via schedulers or
    simulation jobs) without blocking PyMOO's internal loop.

    The generation lifecycle is::

        candidates = optimizer.suggest_candidates()   # calls algorithm.ask()
        for i, c in enumerate(candidates):
            metrics = my_evaluate(c)
            optimizer.update_with_results(i, c, metrics)
        # After the last call above, algorithm.tell() is invoked automatically.
        # The next suggest_candidates() call produces the next generation.

    Supported algorithms:
        - ``"ga"`` — Genetic Algorithm for single-objective optimisation.
        - ``"nsga2"`` — NSGA-II, recommended for 2–3 objectives.
        - ``"nsga3"`` — NSGA-III, recommended for 3+ objectives.
        - ``"moead"`` — MOEA/D, weight-decomposition approach.

    Attributes:
        config: ``PyMOOOptimizerConfig`` instance controlling algorithm behaviour.
        n_gen_completed: Number of generations that have been fully evaluated.

    Examples:
        >>> from aid2e.optimizers.pymoo import PyMOOOptimizer, PyMOOOptimizerConfig
        >>> from aid2e.optimizers.base import SearchSpace
        >>> space = SearchSpace(
        ...     parameters={
        ...         "x": {"type": "range", "bounds": [0.0, 1.0]},
        ...         "y": {"type": "range", "bounds": [0.0, 1.0]},
        ...     }
        ... )
        >>> config = PyMOOOptimizerConfig(pop_size=20, seed=0)
        >>> opt = PyMOOOptimizer(space, config, objective_names=["loss"])
        >>> candidates = opt.suggest_candidates()
        >>> for trial_idx, c in enumerate(candidates):
        ...     opt.update_with_results(trial_idx, c, {"loss": c["x"] + c["y"]})
        >>> best = opt.get_best_trial()

    Notes:
        - Objective directions are translated to PyMOO's minimization convention.
        - Linear parameter constraints are not yet forwarded to PyMOO.  A
          warning is emitted when the search space contains constraints.
        - ``n_candidates`` passed to ``suggest_candidates`` is informational
          only; the actual batch size is determined by the algorithm
          (``pop_size`` for the first generation, ``n_offsprings`` thereafter).

        Project: AID2E v0.0.0 — AI assisted Detector Design for EIC
        Homepage: https://aid2e.github.io/AID2E-framework
        Repository: https://github.com/aid2e/AID2E-framework.git
    """

    def __init__(
        self,
        search_space: Union[SearchSpace, DesignConfig],
        config: PyMOOOptimizerConfig,
        objective_names: List[str],
        seed: Optional[int] = None,
        objective_directions: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Initialise the PyMOO optimizer.

        Args:
            search_space: Parameter search space or a ``DesignConfig`` instance.
                ``DesignConfig`` is automatically converted to ``SearchSpace``.
            config: ``PyMOOOptimizerConfig`` controlling algorithm selection and
                operator hyper-parameters.
            objective_names: Ordered list of objective metric names.  These must
                match keys in the ``metrics`` dict passed to
                ``update_with_results``.
            seed: Integer seed overriding ``config.seed`` when provided.
            objective_directions: Optimization direction for each objective.

        Raises:
            ImportError: If PyMOO is not installed.
            ValueError: If the search space is empty or ``objective_names`` is
                empty.

        Notes:
            The algorithm is initialised lazily; the actual PyMOO ``Problem``
            and ``Algorithm`` objects are created here but no evaluations occur
            until ``suggest_candidates`` is called.
        """
        if not PYMOO_AVAILABLE:
            raise ImportError(
                "PyMOO is required but not installed. "
                "Install with: pip install pymoo"
            )

        effective_seed = seed if seed is not None else config.seed

        super().__init__(
            search_space=search_space,
            objective_names=objective_names,
            seed=effective_seed,
            objective_directions=objective_directions,
        )

        self.config = config
        self.resolved_algorithm = self.config.resolve_algorithm(self.n_objectives)

        if self.search_space.constraints:
            logger.warning(
                "PyMOOOptimizer: %d parameter constraint(s) detected in the search "
                "space but are not yet forwarded to PyMOO. Constraint satisfaction "
                "is not guaranteed during candidate generation.",
                len(self.search_space.constraints),
            )

        # Build ordered parameter list (deterministic iteration order)
        self._param_items: List[Tuple[str, Any]] = list(
            self.search_space.parameters.items()
        )

        # Build numpy bounds
        self._xl, self._xu = self._build_bounds()
        n_var = len(self._param_items)

        # Public PyMOO problem (structural-only ask/tell mode).
        self.problem: PyMOOProblem = PyMOOProblem(
            n_var=n_var,
            n_obj=self.n_objectives,
            xl=self._xl,
            xu=self._xu,
            param_items=self._param_items,
            objective_names=self.objective_names,
        )

        # Create the PyMOO algorithm
        self._algorithm = self._create_algorithm()
        self._algorithm.setup(
            self.problem,
            seed=self.seed,
            verbose=self.config.verbose,
            termination=NoTermination(),
        )

        # _trials and _trial_counter are owned by BaseOptimizer.__init__
        self.n_gen_completed: int = 0

        # Per-generation state (cleared after each tell())
        self._generation_infills: Any = None           # pymoo Population
        self._gen_pos_to_trial_idx: Dict[int, int] = {}  # position → trial_index
        self._result_buffer: Dict[int, Optional[Dict[str, float]]] = {}

        logger.info(
            "PyMOOOptimizer initialised: algorithm=%s, pop_size=%d, "
            "n_params=%d, n_objectives=%d, seed=%s",
            self.resolved_algorithm,
            config.pop_size,
            n_var,
            self.n_objectives,
            self.seed,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_bounds(self) -> Tuple[np.ndarray, np.ndarray]:
        """Build lower/upper bound arrays from the search space parameters.

        ``RangeParameter`` contributes its ``bounds``; ``ChoiceParameter``
        contributes ``[0, n_choices − 1]`` as a continuous range to be
        rounded during decoding.

        Returns:
            Tuple of ``(xl, xu)`` each as a ``float64`` array of shape
            ``(n_var,)``.
        """
        xl, xu = [], []
        for _, param in self._param_items:
            if isinstance(param, DesignRangeParameter):
                xl.append(float(param.bounds[0]))
                xu.append(float(param.bounds[1]))
            elif isinstance(param, DesignChoiceParameter):
                xl.append(0.0)
                xu.append(float(len(param.choices) - 1))
            else:
                raise ValueError(
                    f"Unsupported parameter type for PyMOO: "
                    f"{param.__class__.__name__}"
                )
        return np.array(xl, dtype=float), np.array(xu, dtype=float)

    def _decode_x(self, x_row: np.ndarray) -> Dict[str, Any]:
        """Delegate to ``self.problem.decode_x`` for internal use.

        Args:
            x_row: 1-D float array of length ``n_var``.

        Returns:
            Parameter dictionary for the given individual.
        """
        return self.problem.decode_x(x_row)

    def _encode_params(self, params: Dict[str, Any]) -> np.ndarray:
        """Encode a parameter dictionary back to a PyMOO-compatible float vector.

        Used by ``load_state`` when re-seeding completed trials into PyMOO's
        memory structures.

        Args:
            params: Mapping of parameter names to values.

        Returns:
            1-D float array of length ``n_var``.
        """
        x_row = np.zeros(len(self._param_items), dtype=float)
        for i, (name, param) in enumerate(self._param_items):
            val = params.get(name)
            if isinstance(param, DesignRangeParameter):
                x_row[i] = float(val)
            elif isinstance(param, DesignChoiceParameter):
                if val in param.choices:
                    x_row[i] = float(param.choices.index(val))
                else:
                    x_row[i] = 0.0
        return x_row

    def _create_algorithm(self) -> Any:
        """Instantiate the PyMOO algorithm from the current config.

        Returns:
            An uninitialised PyMOO ``Algorithm`` instance.

        Raises:
            ValueError: If the resolved algorithm is not a supported identifier.

        Notes:
            Algorithm objects are created *before* ``setup()`` is called so
            that ``__init__`` can validate the config without triggering any
            sampling.
        """
        alg = self.resolved_algorithm.lower()
        n_offsprings = self.config.n_offsprings  # None → pop_size default

        crossover = SBX(
            prob=self.config.crossover_prob,
            eta=self.config.crossover_eta,
        )
        mutation = PM(eta=self.config.mutation_eta)
        sampling = FloatRandomSampling()

        if alg == "ga":
            from pymoo.algorithms.soo.nonconvex.ga import GA  # type: ignore[import]

            return GA(
                pop_size=self.config.pop_size,
                n_offsprings=n_offsprings,
                crossover=crossover,
                mutation=mutation,
                sampling=sampling,
            )

        if alg == "nsga2":
            from pymoo.algorithms.moo.nsga2 import NSGA2  # type: ignore[import]

            return NSGA2(
                pop_size=self.config.pop_size,
                n_offsprings=n_offsprings,
                crossover=crossover,
                mutation=mutation,
                sampling=sampling,
            )

        if alg == "nsga3":
            from pymoo.algorithms.moo.nsga3 import NSGA3  # type: ignore[import]
            from pymoo.util.ref_dirs import get_reference_directions  # type: ignore[import]

            ref_dirs = get_reference_directions(
                "das-dennis",
                n_dim=self.n_objectives,
                n_partitions=self.config.n_partitions,
            )
            return NSGA3(
                ref_dirs=ref_dirs,
                pop_size=self.config.pop_size,
                n_offsprings=n_offsprings,
                crossover=crossover,
                mutation=mutation,
                sampling=sampling,
            )

        if alg == "moead":
            from pymoo.algorithms.moo.moead import MOEAD  # type: ignore[import]
            from pymoo.util.ref_dirs import get_reference_directions  # type: ignore[import]

            ref_dirs = get_reference_directions(
                "uniform",
                n_dim=self.n_objectives,
                n_points=self.config.pop_size,
            )
            return MOEAD(
                ref_dirs=ref_dirs,
                n_neighbors=15,
                crossover=crossover,
                mutation=mutation,
                sampling=sampling,
            )

        raise ValueError(
            f"Unknown PyMOO algorithm '{self.resolved_algorithm}'. "
            "Supported: 'ga', 'nsga2', 'nsga3', 'moead'."
        )

    def _flush_generation(self) -> None:
        """Advance the algorithm by one generation using buffered results.

        Called when every candidate in the current generation has completed
        or failed.
        Builds the objective matrix ``F`` from the buffer and calls
        ``algorithm.tell()``.

        Notes:
            This method clears ``_generation_infills``, ``_gen_pos_to_trial_idx``,
            and ``_result_buffer`` after flushing.
        """
        n_gen = len(self._gen_pos_to_trial_idx)
        F = np.zeros((n_gen, self.n_objectives), dtype=float)

        for pos, trial_idx in self._gen_pos_to_trial_idx.items():
            metrics = self._result_buffer[trial_idx]
            if metrics is None:
                F[pos, :] = np.inf
                continue
            for j, obj in enumerate(self.objective_names):
                direction = getattr(
                    self.objective_directions.get(obj),
                    "value",
                    self.objective_directions.get(obj, "minimize"),
                )
                sign = -1.0 if str(direction).lower() == "maximize" else 1.0
                F[pos, j] = metrics[obj] * sign

        self._generation_infills.set("F", F)
        self._algorithm.tell(infills=self._generation_infills)
        self.n_gen_completed += 1

        logger.debug(
            "Generation %d completed: %d individuals evaluated.",
            self.n_gen_completed,
            n_gen,
        )

        # Reset generation state
        self._generation_infills = None
        self._gen_pos_to_trial_idx = {}
        self._result_buffer = {}

    # ------------------------------------------------------------------
    # BaseOptimizer interface overrides + extensions
    # ------------------------------------------------------------------

    def seed_from_trials(
        self,
        trials: List[Trial],
        *,
        only_completed: bool = True,
    ) -> int:
        """Inject completed trials from an external source into the history.

        Extends the base implementation with a guard that prevents seeding
        while a generation is in-flight (i.e., ``suggest_candidates`` has
        been called but not all ``update_with_results`` calls have come back).

        The injected trials are recorded in the optimizer's history and
        visible via ``get_trials()`` and ``get_pareto_front()``.  They do
        *not* advance PyMOO's internal population — the next
        ``suggest_candidates`` call still produces the next generation.

        Args:
            trials: Trials to inject, typically from a previous backend
                (e.g. random-initialisation results).
            only_completed: When ``True`` (default), non-completed trials are
                silently skipped.

        Returns:
            Number of trials actually injected.

        Raises:
            RuntimeError: If a generation is currently in-flight.

        Examples:
            >>> n = pymoo_opt.seed_from_trials(random_opt.get_trials())
            >>> print(f"Seeded {n} prior evaluations")
            >>> candidates = pymoo_opt.suggest_candidates()  # gen 1 starts
        """
        if self._generation_infills is not None:
            raise RuntimeError(
                "Cannot seed trials while a generation is in-flight. "
                "Call update_with_results() for all pending candidates first."
            )
        return super().seed_from_trials(trials, only_completed=only_completed)

    def suggest_candidates(self, n_candidates: int = 1) -> List[Dict[str, Any]]:
        """Suggest the next batch of candidates to evaluate.

        Calls ``algorithm.ask()`` to retrieve the current generation population
        from PyMOO and returns them as a list of parameter dicts.  The actual
        number of candidates is determined by the algorithm (``pop_size`` for
        the initial generation, ``n_offsprings`` thereafter), not by
        ``n_candidates``.

        Args:
            n_candidates: Advisory hint only.  A ``DEBUG``-level message is
                emitted when the hint differs from the actual batch size.

        Returns:
            List of parameter dicts, one per individual in the current
            generation.  Trial indices for these candidates begin at the
            current ``_trial_counter``.

        Raises:
            RuntimeError: If a previous generation has not yet been fully
                evaluated (i.e. some ``update_with_results`` calls are
                outstanding).

        Examples:
            >>> candidates = optimizer.suggest_candidates()
            >>> len(candidates)
            100  # pop_size
        """
        # Guard: cannot ask for a new generation while one is in flight
        if self._generation_infills is not None:
            n_pending = len(self._gen_pos_to_trial_idx) - len(self._result_buffer)
            raise RuntimeError(
                f"Cannot suggest new candidates: {n_pending} evaluation(s) from "
                "the current generation are still outstanding. Call "
                "update_with_results() for all pending candidates first."
            )

        infills = self._algorithm.ask()
        X = infills.get("X")   # shape (batch_size, n_var)
        batch_size = len(X)

        if n_candidates != 1 and n_candidates != batch_size:
            logger.debug(
                "n_candidates=%d ignored — PyMOO algorithm produces %d candidates "
                "(pop_size=%d). Use optimizer.suggest_candidates() without a hint "
                "to suppress this message.",
                n_candidates,
                batch_size,
                self.config.pop_size,
            )

        self._generation_infills = infills
        self._gen_pos_to_trial_idx = {}
        self._result_buffer = {}

        candidates: List[Dict[str, Any]] = []
        for pos, x_row in enumerate(X):
            trial_idx = self._trial_counter
            self._trial_counter += 1
            self._gen_pos_to_trial_idx[pos] = trial_idx

            params = self._decode_x(x_row)

            # Register as pending
            trial = Trial(index=trial_idx, parameters=params, status="pending")
            while len(self._trials) <= trial_idx:
                self._trials.append(None)
            self._trials[trial_idx] = trial

            candidates.append(params)

        logger.debug(
            "Generation %d: suggested %d candidates (trial indices %d–%d).",
            self.n_gen_completed + 1,
            batch_size,
            self._trial_counter - batch_size,
            self._trial_counter - 1,
        )
        return candidates

    def update_with_results(
        self,
        trial_index: int,
        parameters: Dict[str, Any],
        metrics: Dict[str, float],
    ) -> None:
        """Record evaluation results for one candidate and advance the algorithm.

        When the last outstanding candidate of the current generation is
        updated, the generation buffer is flushed automatically via
        ``algorithm.tell()``.

        Args:
            trial_index: Index as returned in ``_trial_counter`` during
                ``suggest_candidates``.  Must correspond to a pending trial.
            parameters: Parameter values that were evaluated (used for
                bookkeeping; the underlying search point is already tracked).
            metrics: Objective values keyed by objective name.  All names
                listed in ``objective_names`` must be present.

        Raises:
            ValueError: If any required objective is missing from ``metrics``.

        Examples:
            >>> optimizer.update_with_results(
            ...     trial_index=0,
            ...     parameters={"x": 0.5, "y": 0.3},
            ...     metrics={"f1": 0.1, "f2": 0.9},
            ... )
        """
        missing = [o for o in self.objective_names if o not in metrics]
        if missing:
            raise ValueError(
                f"update_with_results: missing objectives {missing}. "
                f"Expected {self.objective_names}, got {list(metrics.keys())}."
            )

        # Buffer the result
        self._result_buffer[trial_index] = {k: float(v) for k, v in metrics.items()}

        # Update trial record
        trial = Trial(
            index=trial_index,
            parameters=parameters,
            metrics={k: float(v) for k, v in metrics.items()},
            status="completed",
        )
        while len(self._trials) <= trial_index:
            self._trials.append(None)
        self._trials[trial_index] = trial

        if len(self._result_buffer) == len(self._gen_pos_to_trial_idx):
            self._flush_generation()

    def mark_trial_failed(
        self,
        trial_index: int,
        *,
        parameters: Optional[Dict[str, Any]] = None,
        reason: Optional[str] = None,
    ) -> Trial:
        """Record a failed candidate and allow its generation to finish."""
        trial = super().mark_trial_failed(
            trial_index,
            parameters=parameters,
            reason=reason,
        )
        if trial_index in self._gen_pos_to_trial_idx.values():
            self._result_buffer[trial_index] = None
            if len(self._result_buffer) == len(self._gen_pos_to_trial_idx):
                self._flush_generation()
        return trial

    def serialize_state(self) -> Dict[str, Any]:
        """Serialise optimizer state to a JSON-compatible dictionary.

        The serialised state includes the config, search space description,
        objective names, and all recorded trials.  The PyMOO algorithm's
        internal population state is serialised via ``pickle`` and encoded as
        a base-64 string to allow full resumption without re-running
        evaluations.

        Returns:
            JSON-serialisable dictionary containing all state needed to
            rebuild this optimizer via ``load_state``.

        Examples:
            >>> import json
            >>> state = optimizer.serialize_state()
            >>> with open("checkpoint.json", "w") as f:
            ...     json.dump(state, f)

        Notes:
            If the algorithm cannot be pickled (rare), the ``"algorithm_pickle"``
            key is omitted and a WARNING is logged.  ``load_state`` will then
            recreate the algorithm from the config + seed instead.
        """
        import base64
        import pickle

        space_payload = {
            name: param.model_dump()
            for name, param in self.search_space.parameters.items()
        }
        constraints_payload = [
            c.model_dump() for c in self.search_space.constraints
        ]

        algorithm_pickle: Optional[str] = None
        try:
            algorithm_pickle = base64.b64encode(
                pickle.dumps(self._algorithm)
            ).decode("ascii")
        except Exception as exc:
            logger.warning(
                "Could not pickle PyMOO algorithm state: %s. "
                "load_state will restart the algorithm from scratch.",
                exc,
            )

        return {
            "backend": "pymoo",
            "search_space": {
                "parameters": space_payload,
                "constraints": constraints_payload,
                "name": self.search_space.name,
            },
            "objective_names": self.objective_names,
            "objective_directions": {
                name: getattr(direction, "value", direction)
                for name, direction in self.objective_directions.items()
            },
            "seed": self.seed,
            "config": self.config.model_dump(),
            "resolved_algorithm": self.resolved_algorithm,
            "trials": [
                {
                    "index": t.index,
                    "parameters": t.parameters,
                    "metrics": t.metrics,
                    "status": t.status,
                    "metadata": t.metadata,
                }
                for t in self._trials
                if t is not None
            ],
            "trial_counter": self._trial_counter,
            "n_gen_completed": self.n_gen_completed,
            "algorithm_pickle": algorithm_pickle,
        }

    def load_state(self, state: Dict[str, Any]) -> None:
        """Restore optimizer state from a serialised dictionary.

        The search space, config, objective names, completed trials, and
        (when available) the pickled algorithm state are all restored.  If
        ``algorithm_pickle`` is absent or cannot be deserialised, the algorithm
        is rebuilt from the config and seed, which means the internal
        population will be reset.

        Args:
            state: Dictionary as returned by ``serialize_state``.

        Raises:
            ValueError: If required keys are missing from ``state``.
            ImportError: If PyMOO is not installed.

        Notes:
            After a successful load, ``suggest_candidates`` resumes from where
            optimisation left off (provided the algorithm pickle was valid).
            Any in-flight generation (partial results) is discarded on load.
        """
        if not PYMOO_AVAILABLE:
            raise ImportError(
                "PyMOO is required but not installed. "
                "Install with: pip install pymoo"
            )

        required = {
            "search_space", "objective_names", "objective_directions",
            "config", "trials",
        }
        missing = required - state.keys()
        if missing:
            raise ValueError(f"load_state: missing keys in state: {missing}")

        # Restore config, objectives, seed
        self.config = PyMOOOptimizerConfig(**state["config"])
        self.objective_names = list(state["objective_names"])
        self.objective_directions = dict(state["objective_directions"])
        self.seed = state.get("seed", self.config.seed)
        resolved_algorithm = self.config.resolve_algorithm(self.n_objectives)
        stored_algorithm = state.get("resolved_algorithm")
        if stored_algorithm and stored_algorithm != resolved_algorithm:
            logger.warning(
                "Stored resolved_algorithm '%s' does not match current resolved "
                "algorithm '%s'; using current value.",
                stored_algorithm,
                resolved_algorithm,
            )
        self.resolved_algorithm = resolved_algorithm

        saved_space = state["search_space"]
        self.search_space = SearchSpace(
            parameters=saved_space.get("parameters", {}),
            constraints=saved_space.get("constraints", []),
            name=saved_space.get("name"),
        )

        self._param_items = list(self.search_space.parameters.items())
        self._xl, self._xu = self._build_bounds()
        self.problem = PyMOOProblem(
            n_var=len(self._param_items),
            n_obj=self.n_objectives,
            xl=self._xl,
            xu=self._xu,
            param_items=self._param_items,
            objective_names=self.objective_names,
        )

        # Try to restore the algorithm state from pickle
        import base64
        import pickle

        algorithm_pickle = state.get("algorithm_pickle")
        if algorithm_pickle:
            try:
                self._algorithm = pickle.loads(
                    base64.b64decode(algorithm_pickle.encode("ascii"))
                )
                logger.info("PyMOO algorithm state restored from pickle.")
            except Exception as exc:
                logger.warning(
                    "Could not unpickle algorithm state (%s); "
                    "recreating from config + seed.",
                    exc,
                )
                self._algorithm = self._create_algorithm()
                self._algorithm.setup(
                    self.problem,
                    seed=self.seed,
                    verbose=self.config.verbose,
                    termination=NoTermination(),
                )
        else:
            self._algorithm = self._create_algorithm()
            self._algorithm.setup(
                self.problem,
                seed=self.seed,
                verbose=self.config.verbose,
                termination=NoTermination(),
            )

        # Restore trials
        self._trials = []
        for td in state["trials"]:
            trial = Trial(
                index=td["index"],
                parameters=td["parameters"],
                metrics=td.get("metrics"),
                status=td.get("status", "pending"),
                metadata=td.get("metadata", {}),
            )
            while len(self._trials) <= trial.index:
                self._trials.append(None)
            self._trials[trial.index] = trial

        self._trial_counter = state.get("trial_counter", len(self._trials))
        self.n_gen_completed = state.get("n_gen_completed", 0)

        # Clear any in-flight generation state
        self._generation_infills = None
        self._gen_pos_to_trial_idx = {}
        self._result_buffer = {}

        logger.info(
            "PyMOO optimizer state loaded: algorithm=%s, %d trials, %d generations completed.",
            self.resolved_algorithm,
            len(self._trials),
            self.n_gen_completed,
        )

    # ------------------------------------------------------------------
    # Dunder helpers
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        """Return a concise string representation of the optimizer.

        Returns:
            Human-readable description including algorithm, parameter count,
            objective count, and seed.
        """
        return (
            f"PyMOOOptimizer("
            f"algorithm={self.resolved_algorithm}, "
            f"pop_size={self.config.pop_size}, "
            f"n_params={len(self.search_space.parameters)}, "
            f"n_objectives={self.n_objectives}, "
            f"n_gen_completed={self.n_gen_completed}, "
            f"seed={self.seed})"
        )

__init__(search_space, config, objective_names, seed=None, objective_directions=None)

Initialise the PyMOO optimizer.

Parameters:

Name	Type	Description	Default
search_space	Union[SearchSpace, DesignConfig]	Parameter search space or a DesignConfig instance. DesignConfig is automatically converted to SearchSpace.	required
config	PyMOOOptimizerConfig	PyMOOOptimizerConfig controlling algorithm selection and operator hyper-parameters.	required
objective_names	List[str]	Ordered list of objective metric names. These must match keys in the metrics dict passed to update_with_results.	required
seed	Optional[int]	Integer seed overriding config.seed when provided.	None
objective_directions	Optional[Dict[str, Any]]	Optimization direction for each objective.	None

Raises:

Type	Description
ImportError	If PyMOO is not installed.
ValueError	If the search space is empty or objective_names is empty.

Notes

The algorithm is initialised lazily; the actual PyMOO Problem and Algorithm objects are created here but no evaluations occur until suggest_candidates is called.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def __init__(
    self,
    search_space: Union[SearchSpace, DesignConfig],
    config: PyMOOOptimizerConfig,
    objective_names: List[str],
    seed: Optional[int] = None,
    objective_directions: Optional[Dict[str, Any]] = None,
) -> None:
    """Initialise the PyMOO optimizer.

    Args:
        search_space: Parameter search space or a ``DesignConfig`` instance.
            ``DesignConfig`` is automatically converted to ``SearchSpace``.
        config: ``PyMOOOptimizerConfig`` controlling algorithm selection and
            operator hyper-parameters.
        objective_names: Ordered list of objective metric names.  These must
            match keys in the ``metrics`` dict passed to
            ``update_with_results``.
        seed: Integer seed overriding ``config.seed`` when provided.
        objective_directions: Optimization direction for each objective.

    Raises:
        ImportError: If PyMOO is not installed.
        ValueError: If the search space is empty or ``objective_names`` is
            empty.

    Notes:
        The algorithm is initialised lazily; the actual PyMOO ``Problem``
        and ``Algorithm`` objects are created here but no evaluations occur
        until ``suggest_candidates`` is called.
    """
    if not PYMOO_AVAILABLE:
        raise ImportError(
            "PyMOO is required but not installed. "
            "Install with: pip install pymoo"
        )

    effective_seed = seed if seed is not None else config.seed

    super().__init__(
        search_space=search_space,
        objective_names=objective_names,
        seed=effective_seed,
        objective_directions=objective_directions,
    )

    self.config = config
    self.resolved_algorithm = self.config.resolve_algorithm(self.n_objectives)

    if self.search_space.constraints:
        logger.warning(
            "PyMOOOptimizer: %d parameter constraint(s) detected in the search "
            "space but are not yet forwarded to PyMOO. Constraint satisfaction "
            "is not guaranteed during candidate generation.",
            len(self.search_space.constraints),
        )

    # Build ordered parameter list (deterministic iteration order)
    self._param_items: List[Tuple[str, Any]] = list(
        self.search_space.parameters.items()
    )

    # Build numpy bounds
    self._xl, self._xu = self._build_bounds()
    n_var = len(self._param_items)

    # Public PyMOO problem (structural-only ask/tell mode).
    self.problem: PyMOOProblem = PyMOOProblem(
        n_var=n_var,
        n_obj=self.n_objectives,
        xl=self._xl,
        xu=self._xu,
        param_items=self._param_items,
        objective_names=self.objective_names,
    )

    # Create the PyMOO algorithm
    self._algorithm = self._create_algorithm()
    self._algorithm.setup(
        self.problem,
        seed=self.seed,
        verbose=self.config.verbose,
        termination=NoTermination(),
    )

    # _trials and _trial_counter are owned by BaseOptimizer.__init__
    self.n_gen_completed: int = 0

    # Per-generation state (cleared after each tell())
    self._generation_infills: Any = None           # pymoo Population
    self._gen_pos_to_trial_idx: Dict[int, int] = {}  # position → trial_index
    self._result_buffer: Dict[int, Optional[Dict[str, float]]] = {}

    logger.info(
        "PyMOOOptimizer initialised: algorithm=%s, pop_size=%d, "
        "n_params=%d, n_objectives=%d, seed=%s",
        self.resolved_algorithm,
        config.pop_size,
        n_var,
        self.n_objectives,
        self.seed,
    )

__repr__()

Return a concise string representation of the optimizer.

Returns:

Type	Description
str	Human-readable description including algorithm, parameter count,
str	objective count, and seed.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def __repr__(self) -> str:
    """Return a concise string representation of the optimizer.

    Returns:
        Human-readable description including algorithm, parameter count,
        objective count, and seed.
    """
    return (
        f"PyMOOOptimizer("
        f"algorithm={self.resolved_algorithm}, "
        f"pop_size={self.config.pop_size}, "
        f"n_params={len(self.search_space.parameters)}, "
        f"n_objectives={self.n_objectives}, "
        f"n_gen_completed={self.n_gen_completed}, "
        f"seed={self.seed})"
    )

load_state(state)

Restore optimizer state from a serialised dictionary.

The search space, config, objective names, completed trials, and (when available) the pickled algorithm state are all restored. If algorithm_pickle is absent or cannot be deserialised, the algorithm is rebuilt from the config and seed, which means the internal population will be reset.

Parameters:

Name	Type	Description	Default
state	Dict[str, Any]	Dictionary as returned by serialize_state.	required

Raises:

Type	Description
ValueError	If required keys are missing from state.
ImportError	If PyMOO is not installed.

Notes

After a successful load, suggest_candidates resumes from where optimisation left off (provided the algorithm pickle was valid). Any in-flight generation (partial results) is discarded on load.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def load_state(self, state: Dict[str, Any]) -> None:
    """Restore optimizer state from a serialised dictionary.

    The search space, config, objective names, completed trials, and
    (when available) the pickled algorithm state are all restored.  If
    ``algorithm_pickle`` is absent or cannot be deserialised, the algorithm
    is rebuilt from the config and seed, which means the internal
    population will be reset.

    Args:
        state: Dictionary as returned by ``serialize_state``.

    Raises:
        ValueError: If required keys are missing from ``state``.
        ImportError: If PyMOO is not installed.

    Notes:
        After a successful load, ``suggest_candidates`` resumes from where
        optimisation left off (provided the algorithm pickle was valid).
        Any in-flight generation (partial results) is discarded on load.
    """
    if not PYMOO_AVAILABLE:
        raise ImportError(
            "PyMOO is required but not installed. "
            "Install with: pip install pymoo"
        )

    required = {
        "search_space", "objective_names", "objective_directions",
        "config", "trials",
    }
    missing = required - state.keys()
    if missing:
        raise ValueError(f"load_state: missing keys in state: {missing}")

    # Restore config, objectives, seed
    self.config = PyMOOOptimizerConfig(**state["config"])
    self.objective_names = list(state["objective_names"])
    self.objective_directions = dict(state["objective_directions"])
    self.seed = state.get("seed", self.config.seed)
    resolved_algorithm = self.config.resolve_algorithm(self.n_objectives)
    stored_algorithm = state.get("resolved_algorithm")
    if stored_algorithm and stored_algorithm != resolved_algorithm:
        logger.warning(
            "Stored resolved_algorithm '%s' does not match current resolved "
            "algorithm '%s'; using current value.",
            stored_algorithm,
            resolved_algorithm,
        )
    self.resolved_algorithm = resolved_algorithm

    saved_space = state["search_space"]
    self.search_space = SearchSpace(
        parameters=saved_space.get("parameters", {}),
        constraints=saved_space.get("constraints", []),
        name=saved_space.get("name"),
    )

    self._param_items = list(self.search_space.parameters.items())
    self._xl, self._xu = self._build_bounds()
    self.problem = PyMOOProblem(
        n_var=len(self._param_items),
        n_obj=self.n_objectives,
        xl=self._xl,
        xu=self._xu,
        param_items=self._param_items,
        objective_names=self.objective_names,
    )

    # Try to restore the algorithm state from pickle
    import base64
    import pickle

    algorithm_pickle = state.get("algorithm_pickle")
    if algorithm_pickle:
        try:
            self._algorithm = pickle.loads(
                base64.b64decode(algorithm_pickle.encode("ascii"))
            )
            logger.info("PyMOO algorithm state restored from pickle.")
        except Exception as exc:
            logger.warning(
                "Could not unpickle algorithm state (%s); "
                "recreating from config + seed.",
                exc,
            )
            self._algorithm = self._create_algorithm()
            self._algorithm.setup(
                self.problem,
                seed=self.seed,
                verbose=self.config.verbose,
                termination=NoTermination(),
            )
    else:
        self._algorithm = self._create_algorithm()
        self._algorithm.setup(
            self.problem,
            seed=self.seed,
            verbose=self.config.verbose,
            termination=NoTermination(),
        )

    # Restore trials
    self._trials = []
    for td in state["trials"]:
        trial = Trial(
            index=td["index"],
            parameters=td["parameters"],
            metrics=td.get("metrics"),
            status=td.get("status", "pending"),
            metadata=td.get("metadata", {}),
        )
        while len(self._trials) <= trial.index:
            self._trials.append(None)
        self._trials[trial.index] = trial

    self._trial_counter = state.get("trial_counter", len(self._trials))
    self.n_gen_completed = state.get("n_gen_completed", 0)

    # Clear any in-flight generation state
    self._generation_infills = None
    self._gen_pos_to_trial_idx = {}
    self._result_buffer = {}

    logger.info(
        "PyMOO optimizer state loaded: algorithm=%s, %d trials, %d generations completed.",
        self.resolved_algorithm,
        len(self._trials),
        self.n_gen_completed,
    )

mark_trial_failed(trial_index, *, parameters=None, reason=None)

Record a failed candidate and allow its generation to finish.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def mark_trial_failed(
    self,
    trial_index: int,
    *,
    parameters: Optional[Dict[str, Any]] = None,
    reason: Optional[str] = None,
) -> Trial:
    """Record a failed candidate and allow its generation to finish."""
    trial = super().mark_trial_failed(
        trial_index,
        parameters=parameters,
        reason=reason,
    )
    if trial_index in self._gen_pos_to_trial_idx.values():
        self._result_buffer[trial_index] = None
        if len(self._result_buffer) == len(self._gen_pos_to_trial_idx):
            self._flush_generation()
    return trial

seed_from_trials(trials, *, only_completed=True)

Inject completed trials from an external source into the history.

Extends the base implementation with a guard that prevents seeding while a generation is in-flight (i.e., suggest_candidates has been called but not all update_with_results calls have come back).

The injected trials are recorded in the optimizer's history and visible via get_trials() and get_pareto_front(). They do not advance PyMOO's internal population — the next suggest_candidates call still produces the next generation.

Parameters:

Name	Type	Description	Default
trials	List[Trial]	Trials to inject, typically from a previous backend (e.g. random-initialisation results).	required
only_completed	bool	When True (default), non-completed trials are silently skipped.	True

Returns:

Type	Description
int	Number of trials actually injected.

Raises:

Type	Description
RuntimeError	If a generation is currently in-flight.

Examples:

>>> n = pymoo_opt.seed_from_trials(random_opt.get_trials())
>>> print(f"Seeded {n} prior evaluations")
>>> candidates = pymoo_opt.suggest_candidates()  # gen 1 starts

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def seed_from_trials(
    self,
    trials: List[Trial],
    *,
    only_completed: bool = True,
) -> int:
    """Inject completed trials from an external source into the history.

    Extends the base implementation with a guard that prevents seeding
    while a generation is in-flight (i.e., ``suggest_candidates`` has
    been called but not all ``update_with_results`` calls have come back).

    The injected trials are recorded in the optimizer's history and
    visible via ``get_trials()`` and ``get_pareto_front()``.  They do
    *not* advance PyMOO's internal population — the next
    ``suggest_candidates`` call still produces the next generation.

    Args:
        trials: Trials to inject, typically from a previous backend
            (e.g. random-initialisation results).
        only_completed: When ``True`` (default), non-completed trials are
            silently skipped.

    Returns:
        Number of trials actually injected.

    Raises:
        RuntimeError: If a generation is currently in-flight.

    Examples:
        >>> n = pymoo_opt.seed_from_trials(random_opt.get_trials())
        >>> print(f"Seeded {n} prior evaluations")
        >>> candidates = pymoo_opt.suggest_candidates()  # gen 1 starts
    """
    if self._generation_infills is not None:
        raise RuntimeError(
            "Cannot seed trials while a generation is in-flight. "
            "Call update_with_results() for all pending candidates first."
        )
    return super().seed_from_trials(trials, only_completed=only_completed)

serialize_state()

Serialise optimizer state to a JSON-compatible dictionary.

The serialised state includes the config, search space description, objective names, and all recorded trials. The PyMOO algorithm's internal population state is serialised via pickle and encoded as a base-64 string to allow full resumption without re-running evaluations.

Returns:

Type	Description
Dict[str, Any]	JSON-serialisable dictionary containing all state needed to
Dict[str, Any]	rebuild this optimizer via load_state.

Examples:

>>> import json
>>> state = optimizer.serialize_state()
>>> with open("checkpoint.json", "w") as f:
...     json.dump(state, f)

Notes

If the algorithm cannot be pickled (rare), the "algorithm_pickle" key is omitted and a WARNING is logged. load_state will then recreate the algorithm from the config + seed instead.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def serialize_state(self) -> Dict[str, Any]:
    """Serialise optimizer state to a JSON-compatible dictionary.

    The serialised state includes the config, search space description,
    objective names, and all recorded trials.  The PyMOO algorithm's
    internal population state is serialised via ``pickle`` and encoded as
    a base-64 string to allow full resumption without re-running
    evaluations.

    Returns:
        JSON-serialisable dictionary containing all state needed to
        rebuild this optimizer via ``load_state``.

    Examples:
        >>> import json
        >>> state = optimizer.serialize_state()
        >>> with open("checkpoint.json", "w") as f:
        ...     json.dump(state, f)

    Notes:
        If the algorithm cannot be pickled (rare), the ``"algorithm_pickle"``
        key is omitted and a WARNING is logged.  ``load_state`` will then
        recreate the algorithm from the config + seed instead.
    """
    import base64
    import pickle

    space_payload = {
        name: param.model_dump()
        for name, param in self.search_space.parameters.items()
    }
    constraints_payload = [
        c.model_dump() for c in self.search_space.constraints
    ]

    algorithm_pickle: Optional[str] = None
    try:
        algorithm_pickle = base64.b64encode(
            pickle.dumps(self._algorithm)
        ).decode("ascii")
    except Exception as exc:
        logger.warning(
            "Could not pickle PyMOO algorithm state: %s. "
            "load_state will restart the algorithm from scratch.",
            exc,
        )

    return {
        "backend": "pymoo",
        "search_space": {
            "parameters": space_payload,
            "constraints": constraints_payload,
            "name": self.search_space.name,
        },
        "objective_names": self.objective_names,
        "objective_directions": {
            name: getattr(direction, "value", direction)
            for name, direction in self.objective_directions.items()
        },
        "seed": self.seed,
        "config": self.config.model_dump(),
        "resolved_algorithm": self.resolved_algorithm,
        "trials": [
            {
                "index": t.index,
                "parameters": t.parameters,
                "metrics": t.metrics,
                "status": t.status,
                "metadata": t.metadata,
            }
            for t in self._trials
            if t is not None
        ],
        "trial_counter": self._trial_counter,
        "n_gen_completed": self.n_gen_completed,
        "algorithm_pickle": algorithm_pickle,
    }

suggest_candidates(n_candidates=1)

Suggest the next batch of candidates to evaluate.

Calls algorithm.ask() to retrieve the current generation population from PyMOO and returns them as a list of parameter dicts. The actual number of candidates is determined by the algorithm (pop_size for the initial generation, n_offsprings thereafter), not by n_candidates.

Parameters:

Name	Type	Description	Default
n_candidates	int	Advisory hint only. A DEBUG-level message is emitted when the hint differs from the actual batch size.	1

Returns:

Type	Description
List[Dict[str, Any]]	List of parameter dicts, one per individual in the current
List[Dict[str, Any]]	generation. Trial indices for these candidates begin at the
List[Dict[str, Any]]	current _trial_counter.

Raises:

Type	Description
RuntimeError	If a previous generation has not yet been fully evaluated (i.e. some update_with_results calls are outstanding).

Examples:

>>> candidates = optimizer.suggest_candidates()
>>> len(candidates)
100  # pop_size

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def suggest_candidates(self, n_candidates: int = 1) -> List[Dict[str, Any]]:
    """Suggest the next batch of candidates to evaluate.

    Calls ``algorithm.ask()`` to retrieve the current generation population
    from PyMOO and returns them as a list of parameter dicts.  The actual
    number of candidates is determined by the algorithm (``pop_size`` for
    the initial generation, ``n_offsprings`` thereafter), not by
    ``n_candidates``.

    Args:
        n_candidates: Advisory hint only.  A ``DEBUG``-level message is
            emitted when the hint differs from the actual batch size.

    Returns:
        List of parameter dicts, one per individual in the current
        generation.  Trial indices for these candidates begin at the
        current ``_trial_counter``.

    Raises:
        RuntimeError: If a previous generation has not yet been fully
            evaluated (i.e. some ``update_with_results`` calls are
            outstanding).

    Examples:
        >>> candidates = optimizer.suggest_candidates()
        >>> len(candidates)
        100  # pop_size
    """
    # Guard: cannot ask for a new generation while one is in flight
    if self._generation_infills is not None:
        n_pending = len(self._gen_pos_to_trial_idx) - len(self._result_buffer)
        raise RuntimeError(
            f"Cannot suggest new candidates: {n_pending} evaluation(s) from "
            "the current generation are still outstanding. Call "
            "update_with_results() for all pending candidates first."
        )

    infills = self._algorithm.ask()
    X = infills.get("X")   # shape (batch_size, n_var)
    batch_size = len(X)

    if n_candidates != 1 and n_candidates != batch_size:
        logger.debug(
            "n_candidates=%d ignored — PyMOO algorithm produces %d candidates "
            "(pop_size=%d). Use optimizer.suggest_candidates() without a hint "
            "to suppress this message.",
            n_candidates,
            batch_size,
            self.config.pop_size,
        )

    self._generation_infills = infills
    self._gen_pos_to_trial_idx = {}
    self._result_buffer = {}

    candidates: List[Dict[str, Any]] = []
    for pos, x_row in enumerate(X):
        trial_idx = self._trial_counter
        self._trial_counter += 1
        self._gen_pos_to_trial_idx[pos] = trial_idx

        params = self._decode_x(x_row)

        # Register as pending
        trial = Trial(index=trial_idx, parameters=params, status="pending")
        while len(self._trials) <= trial_idx:
            self._trials.append(None)
        self._trials[trial_idx] = trial

        candidates.append(params)

    logger.debug(
        "Generation %d: suggested %d candidates (trial indices %d–%d).",
        self.n_gen_completed + 1,
        batch_size,
        self._trial_counter - batch_size,
        self._trial_counter - 1,
    )
    return candidates

update_with_results(trial_index, parameters, metrics)

Record evaluation results for one candidate and advance the algorithm.

When the last outstanding candidate of the current generation is updated, the generation buffer is flushed automatically via algorithm.tell().

Parameters:

Name	Type	Description	Default
trial_index	int	Index as returned in _trial_counter during suggest_candidates. Must correspond to a pending trial.	required
parameters	Dict[str, Any]	Parameter values that were evaluated (used for bookkeeping; the underlying search point is already tracked).	required
metrics	Dict[str, float]	Objective values keyed by objective name. All names listed in objective_names must be present.	required

Raises:

Type	Description
ValueError	If any required objective is missing from metrics.

Examples:

>>> optimizer.update_with_results(
...     trial_index=0,
...     parameters={"x": 0.5, "y": 0.3},
...     metrics={"f1": 0.1, "f2": 0.9},
... )

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def update_with_results(
    self,
    trial_index: int,
    parameters: Dict[str, Any],
    metrics: Dict[str, float],
) -> None:
    """Record evaluation results for one candidate and advance the algorithm.

    When the last outstanding candidate of the current generation is
    updated, the generation buffer is flushed automatically via
    ``algorithm.tell()``.

    Args:
        trial_index: Index as returned in ``_trial_counter`` during
            ``suggest_candidates``.  Must correspond to a pending trial.
        parameters: Parameter values that were evaluated (used for
            bookkeeping; the underlying search point is already tracked).
        metrics: Objective values keyed by objective name.  All names
            listed in ``objective_names`` must be present.

    Raises:
        ValueError: If any required objective is missing from ``metrics``.

    Examples:
        >>> optimizer.update_with_results(
        ...     trial_index=0,
        ...     parameters={"x": 0.5, "y": 0.3},
        ...     metrics={"f1": 0.1, "f2": 0.9},
        ... )
    """
    missing = [o for o in self.objective_names if o not in metrics]
    if missing:
        raise ValueError(
            f"update_with_results: missing objectives {missing}. "
            f"Expected {self.objective_names}, got {list(metrics.keys())}."
        )

    # Buffer the result
    self._result_buffer[trial_index] = {k: float(v) for k, v in metrics.items()}

    # Update trial record
    trial = Trial(
        index=trial_index,
        parameters=parameters,
        metrics={k: float(v) for k, v in metrics.items()},
        status="completed",
    )
    while len(self._trials) <= trial_index:
        self._trials.append(None)
    self._trials[trial_index] = trial

    if len(self._result_buffer) == len(self._gen_pos_to_trial_idx):
        self._flush_generation()

PyMOOOptimizerConfig

Bases: BaseModel

Configuration for PyMOO-based evolutionary optimizers.

Attributes:

Name	Type	Description
algorithm	Optional[PyMOOAlgorithm]	Optional evolutionary algorithm identifier. If omitted, AID2E infers "ga" for single-objective problems and "nsga2" for multi-objective problems.
pop_size	int	Population size (number of individuals per generation).
n_offsprings	Optional[int]	Number of offspring generated each generation. None defaults to pop_size.
crossover_prob	float	Simulated Binary Crossover (SBX) probability.
crossover_eta	float	SBX distribution index — larger values produce offspring closer to the parents.
mutation_eta	float	Polynomial mutation distribution index.
n_iterations	int	Number of generations to run when using this config in declarative/runtime-driven flows.
n_partitions	int	Reference-direction partitions for NSGA-III and MOEA/D. The total number of reference directions grows combinatorially with this value and n_objectives. Ignored for NSGA-II.
seed	Optional[int]	Random seed for reproducibility. None yields non-deterministic results.
verbose	bool	Whether PyMOO prints per-generation progress to stdout.

Examples:

>>> config = PyMOOOptimizerConfig(
...     pop_size=100,
...     seed=42,
... )
>>> config.algorithm is None
True
>>> config2 = PyMOOOptimizerConfig(algorithm="nsga3", n_partitions=12)

Notes

• ga is the recommended default for single-objective problems.
• NSGA-II is the recommended default for 2-objective problems.
• For 3+ objectives consider NSGA-III or MOEA/D — their reference direction structures are better suited to high-dimensional fronts.
• n_partitions has a strong effect on runtime for NSGA-III/MOEA/D; start with 12 for 2-3 objectives and reduce for 4+ objectives.

Source code in src/aid2e/optimizers/pymoo/config.py

class PyMOOOptimizerConfig(BaseModel):
    """Configuration for PyMOO-based evolutionary optimizers.

    Attributes:
        algorithm: Optional evolutionary algorithm identifier. If omitted,
            AID2E infers ``"ga"`` for single-objective problems and
            ``"nsga2"`` for multi-objective problems.
        pop_size: Population size (number of individuals per generation).
        n_offsprings: Number of offspring generated each generation.  ``None``
            defaults to ``pop_size``.
        crossover_prob: Simulated Binary Crossover (SBX) probability.
        crossover_eta: SBX distribution index — larger values produce offspring
            closer to the parents.
        mutation_eta: Polynomial mutation distribution index.
        n_iterations: Number of generations to run when using this config in
            declarative/runtime-driven flows.
        n_partitions: Reference-direction partitions for NSGA-III and MOEA/D.
            The total number of reference directions grows combinatorially with
            this value and ``n_objectives``.  Ignored for NSGA-II.
        seed: Random seed for reproducibility.  ``None`` yields non-deterministic
            results.
        verbose: Whether PyMOO prints per-generation progress to stdout.

    Examples:
        >>> config = PyMOOOptimizerConfig(
        ...     pop_size=100,
        ...     seed=42,
        ... )
        >>> config.algorithm is None
        True
        >>> config2 = PyMOOOptimizerConfig(algorithm="nsga3", n_partitions=12)

    Notes:
        - ``ga`` is the recommended default for single-objective problems.
        - NSGA-II is the recommended default for 2-objective problems.
        - For 3+ objectives consider NSGA-III or MOEA/D — their reference
          direction structures are better suited to high-dimensional fronts.
        - ``n_partitions`` has a strong effect on runtime for NSGA-III/MOEA/D;
          start with 12 for 2-3 objectives and reduce for 4+ objectives.
    """

    algorithm: Optional[PyMOOAlgorithm] = Field(
        default=None,
        description=(
            "Optional evolutionary algorithm. If omitted, AID2E infers 'ga' "
            "for single-objective problems and 'nsga2' for multi-objective problems."
        ),
    )
    pop_size: int = Field(
        default=100,
        ge=2,
        description="Population size — number of candidate solutions per generation.",
    )
    n_offsprings: Optional[int] = Field(
        default=None,
        ge=1,
        description=(
            "Number of offspring per generation. "
            "Defaults to pop_size when None."
        ),
    )
    crossover_prob: float = Field(
        default=0.9,
        ge=0.0,
        le=1.0,
        description="SBX crossover probability.",
    )
    crossover_eta: float = Field(
        default=15.0,
        gt=0.0,
        description="SBX crossover distribution index.",
    )
    mutation_eta: float = Field(
        default=20.0,
        gt=0.0,
        description="Polynomial mutation distribution index.",
    )
    n_iterations: int = Field(
        default=50,
        ge=1,
        description="Number of generations for runtime-driven optimization loops.",
    )
    n_partitions: int = Field(
        default=12,
        ge=1,
        description=(
            "Reference-direction partitions for NSGA-III and MOEA/D. "
            "Ignored for NSGA-II."
        ),
    )
    seed: Optional[int] = Field(
        default=None,
        description="Random seed. None means non-deterministic.",
    )
    verbose: bool = Field(
        default=False,
        description="Print per-generation statistics to stdout.",
    )

    def resolve_algorithm(self, n_objectives: int) -> PyMOOAlgorithm:
        """Resolve the algorithm for the given objective count.

        Args:
            n_objectives: Number of objectives in the optimization problem.

        Returns:
            Concrete PyMOO algorithm identifier.

        Raises:
            ValueError: If the configured explicit algorithm is incompatible
                with the objective count.
        """
        if n_objectives < 1:
            raise ValueError("n_objectives must be >= 1")

        if self.algorithm is None:
            return "ga" if n_objectives == 1 else "nsga2"

        if self.algorithm == "ga" and n_objectives != 1:
            raise ValueError(
                "PyMOO algorithm 'ga' only supports single-objective problems. "
                f"Received {n_objectives} objectives."
            )

        if self.algorithm in {"nsga2", "nsga3", "moead"} and n_objectives == 1:
            raise ValueError(
                f"PyMOO algorithm '{self.algorithm}' requires a multi-objective "
                "problem. Use 'ga' or omit 'algorithm' for single-objective optimization."
            )

        return self.algorithm

resolve_algorithm(n_objectives)

Resolve the algorithm for the given objective count.

Parameters:

Name	Type	Description	Default
n_objectives	int	Number of objectives in the optimization problem.	required

Returns:

Type	Description
PyMOOAlgorithm	Concrete PyMOO algorithm identifier.

Raises:

Type	Description
ValueError	If the configured explicit algorithm is incompatible with the objective count.

Source code in src/aid2e/optimizers/pymoo/config.py

def resolve_algorithm(self, n_objectives: int) -> PyMOOAlgorithm:
    """Resolve the algorithm for the given objective count.

    Args:
        n_objectives: Number of objectives in the optimization problem.

    Returns:
        Concrete PyMOO algorithm identifier.

    Raises:
        ValueError: If the configured explicit algorithm is incompatible
            with the objective count.
    """
    if n_objectives < 1:
        raise ValueError("n_objectives must be >= 1")

    if self.algorithm is None:
        return "ga" if n_objectives == 1 else "nsga2"

    if self.algorithm == "ga" and n_objectives != 1:
        raise ValueError(
            "PyMOO algorithm 'ga' only supports single-objective problems. "
            f"Received {n_objectives} objectives."
        )

    if self.algorithm in {"nsga2", "nsga3", "moead"} and n_objectives == 1:
        raise ValueError(
            f"PyMOO algorithm '{self.algorithm}' requires a multi-objective "
            "problem. Use 'ga' or omit 'algorithm' for single-objective optimization."
        )

    return self.algorithm

PyMOOProblem

Bases: Problem

PyMOO Problem that wraps an AID2E SearchSpace.

This problem is structural-only for ask/tell workflows. ``decode_x``
translates PyMOO float vectors into

human-readable AID2E parameter dicts — the same representation returned by suggest_candidates.

Parameters:

Name	Type	Description	Default
n_var	int	Number of continuous decision variables.	required
n_obj	int	Number of objectives.	required
xl	ndarray	Lower-bound array of shape (n_var,).	required
xu	ndarray	Upper-bound array of shape (n_var,).	required
param_items	List[Tuple[str, Any]]	Ordered list of (name, BaseParameter) pairs used for encoding/decoding.	required
objective_names	List[str]	Ordered objective names.	required

Source code in src/aid2e/optimizers/pymoo/optimizer.py

class PyMOOProblem(Problem):
    """PyMOO Problem that wraps an AID2E ``SearchSpace``.

        This problem is structural-only for ask/tell workflows. ``decode_x``
        translates PyMOO float vectors into
    human-readable AID2E parameter dicts — the same representation returned
    by ``suggest_candidates``.

    Args:
        n_var: Number of continuous decision variables.
        n_obj: Number of objectives.
        xl: Lower-bound array of shape ``(n_var,)``.
        xu: Upper-bound array of shape ``(n_var,)``.
        param_items: Ordered list of ``(name, BaseParameter)`` pairs used for
            encoding/decoding.
        objective_names: Ordered objective names.
    """

    def __init__(
        self,
        n_var: int,
        n_obj: int,
        xl: np.ndarray,
        xu: np.ndarray,
        param_items: List[Tuple[str, Any]],
        objective_names: List[str],
    ) -> None:
        """Initialise the PyMOO problem with search-space metadata."""
        super().__init__(n_var=n_var, n_obj=n_obj, xl=xl, xu=xu)
        self._param_items = param_items
        self._objective_names = objective_names

    def decode_x(self, x_row: np.ndarray) -> Dict[str, Any]:
        """Translate a PyMOO float vector into an AID2E parameter dictionary.

        Args:
            x_row: 1-D float array of length ``n_var``.

        Returns:
            Mapping of parameter names to decoded values (``float`` for
            ``RangeParameter``, choice value for ``ChoiceParameter``).
        """
        params: Dict[str, Any] = {}
        for i, (name, param) in enumerate(self._param_items):
            val = float(x_row[i])
            if isinstance(param, DesignRangeParameter):
                params[name] = val
            elif isinstance(param, DesignChoiceParameter):
                idx = int(round(val))
                idx = max(0, min(idx, len(param.choices) - 1))
                params[name] = param.choices[idx]
        return params

    def _evaluate(
        self, x: np.ndarray, out: dict, *args: Any, **kwargs: Any
    ) -> None:
        """Raise because AID2E always uses external evaluation.

        Args:
            x: Population matrix (unused).
            out: Output dictionary (unused).

        Raises:
            NotImplementedError: Always, because evaluation belongs to the
                workflow/scheduler layer in AID2E.
        """
        raise NotImplementedError(
            "PyMOOProblem is structural-only in ask/tell mode. "
            "Use PyMOOOptimizer.suggest_candidates() and "
            "PyMOOOptimizer.update_with_results() with external evaluation."
        )

__init__(n_var, n_obj, xl, xu, param_items, objective_names)

Initialise the PyMOO problem with search-space metadata.

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def __init__(
    self,
    n_var: int,
    n_obj: int,
    xl: np.ndarray,
    xu: np.ndarray,
    param_items: List[Tuple[str, Any]],
    objective_names: List[str],
) -> None:
    """Initialise the PyMOO problem with search-space metadata."""
    super().__init__(n_var=n_var, n_obj=n_obj, xl=xl, xu=xu)
    self._param_items = param_items
    self._objective_names = objective_names

decode_x(x_row)

Translate a PyMOO float vector into an AID2E parameter dictionary.

Parameters:

Name	Type	Description	Default
x_row	ndarray	1-D float array of length n_var.	required

Returns:

Type	Description
Dict[str, Any]	Mapping of parameter names to decoded values (float for
Dict[str, Any]	RangeParameter, choice value for ChoiceParameter).

Source code in src/aid2e/optimizers/pymoo/optimizer.py

def decode_x(self, x_row: np.ndarray) -> Dict[str, Any]:
    """Translate a PyMOO float vector into an AID2E parameter dictionary.

    Args:
        x_row: 1-D float array of length ``n_var``.

    Returns:
        Mapping of parameter names to decoded values (``float`` for
        ``RangeParameter``, choice value for ``ChoiceParameter``).
    """
    params: Dict[str, Any] = {}
    for i, (name, param) in enumerate(self._param_items):
        val = float(x_row[i])
        if isinstance(param, DesignRangeParameter):
            params[name] = val
        elif isinstance(param, DesignChoiceParameter):
            idx = int(round(val))
            idx = max(0, min(idx, len(param.choices) - 1))
            params[name] = param.choices[idx]
    return params

SearchSpace dataclass

Represent an optimization search space built from design parameters.

Attributes:

Name	Type	Description
parameters	Dict[str, BaseParameter]	Mapping of parameter names to typed design parameters.
constraints	List[ParameterConstraint]	Optional list of parameter constraints to enforce.
name	Optional[str]	Optional identifier for the search space.
source_config	Optional[DesignConfig]	Optional originating DesignConfig for traceability.

Examples:

>>> from aid2e.utilities.configurations.base_models import RangeParameter
>>> params = {
...     "x": RangeParameter(name="x", value=0.5, bounds=(0.0, 1.0)),
... }
>>> space = SearchSpace(parameters=params)
>>> space.validate({"x": 0.4})
(True, [])

Source code in src/aid2e/optimizers/base.py

@dataclass
class SearchSpace:
    """Represent an optimization search space built from design parameters.

    Attributes:
        parameters: Mapping of parameter names to typed design parameters.
        constraints: Optional list of parameter constraints to enforce.
        name: Optional identifier for the search space.
        source_config: Optional originating DesignConfig for traceability.

    Examples:
        >>> from aid2e.utilities.configurations.base_models import RangeParameter
        >>> params = {
        ...     "x": RangeParameter(name="x", value=0.5, bounds=(0.0, 1.0)),
        ... }
        >>> space = SearchSpace(parameters=params)
        >>> space.validate({"x": 0.4})
        (True, [])
    """

    parameters: Dict[str, BaseParameter]
    constraints: List[ParameterConstraint] = field(default_factory=list)
    name: Optional[str] = None
    source_config: Optional[DesignConfig] = None

    def __post_init__(self) -> None:
        """Normalize parameter and constraint inputs after initialization."""
        parsed_parameters: Dict[str, BaseParameter] = {}
        for param_name, param in self.parameters.items():
            if isinstance(param, BaseParameter):
                parsed = param
            elif isinstance(param, dict):
                param_data = dict(param)
                if "values" in param_data:
                    raise ValueError(
                        f"Parameter '{param_name}' uses retired key 'values'. "
                        "Use 'choices'."
                    )
                if "bounds" in param_data and "value" not in param_data:
                    raise ValueError(
                        f"Range parameter '{param_name}' must define an explicit "
                        "'value' alongside 'bounds'."
                    )
                if "choices" in param_data and "value" not in param_data:
                    raise ValueError(
                        f"Choice parameter '{param_name}' must define an explicit "
                        "'value' alongside 'choices'."
                    )
                parsed = parse_parameter(param_name, param_data)
            else:
                raise TypeError(
                    "Parameters must be BaseParameter instances or dictionaries"
                )

            if parsed.name != param_name:
                parsed = parsed.model_copy(update={"name": param_name})
            parsed_parameters[param_name] = parsed

        self.parameters = parsed_parameters

        constraints_input = self.constraints or []
        self.constraints = [
            c if isinstance(c, ParameterConstraint) else ParameterConstraint(**c)
            for c in constraints_input
        ]

    @classmethod
    def from_design_config(cls, design_config: DesignConfig) -> "SearchSpace":
        """Build a search space from a DesignConfig instance.

        Args:
            design_config: Fully validated design configuration containing
                parameters and optional parameter constraints.

        Returns:
            SearchSpace populated with flattened parameters and constraints.
        """

        return cls(
            parameters=design_config.get_flat_parameters(),
            constraints=design_config.parameter_constraints or [],
            name=getattr(design_config, "name", None),
            source_config=design_config,
        )

    def validate(self, param_values: Dict[str, Any]) -> Tuple[bool, List[str]]:
        """Check if parameter values satisfy all constraints (for non-Ax optimizers).

        This method is for optimizers that DON'T have native constraint support
        (e.g., random search, simple evolutionary algorithms). For optimizers
        with native constraint support (e.g., Ax), pass self.constraints directly
        to the optimizer instead of calling this method.

        Args:
            param_values: Mapping of qualified parameter names to concrete values.

        Returns:
            Tuple of ``(all_valid, failed_constraints)`` where ``all_valid`` is
            ``True`` when every constraint passes and ``failed_constraints``
            lists the names (or error strings) of failing constraints.

        Example:
            >>> # For optimizers WITHOUT constraint support
            >>> is_valid, failures = search_space.validate(candidate)
            >>> if not is_valid:
            ...     print(f"Constraint violations: {failures}")

            >>> # For Ax (HAS constraint support) - DON'T use this method
            >>> # Instead, pass search_space.constraints to Ax optimizer

        Notes:
            - Only use this for runtime checking with constraint-agnostic optimizers
            - Ax and similar optimizers handle constraints internally via self.constraints
            - Syntax validation already done at DesignConfig load time
        """

        if not self.constraints:
            return True, []

        failed: List[str] = []
        for constraint in self.constraints:
            try:
                if not constraint.evaluate(param_values):
                    failed.append(constraint.name)
            except Exception as exc:  # pragma: no cover - defensive logging
                failed.append(f"{constraint.name} (error: {exc})")

        return len(failed) == 0, failed

__post_init__()

Normalize parameter and constraint inputs after initialization.

Source code in src/aid2e/optimizers/base.py

def __post_init__(self) -> None:
    """Normalize parameter and constraint inputs after initialization."""
    parsed_parameters: Dict[str, BaseParameter] = {}
    for param_name, param in self.parameters.items():
        if isinstance(param, BaseParameter):
            parsed = param
        elif isinstance(param, dict):
            param_data = dict(param)
            if "values" in param_data:
                raise ValueError(
                    f"Parameter '{param_name}' uses retired key 'values'. "
                    "Use 'choices'."
                )
            if "bounds" in param_data and "value" not in param_data:
                raise ValueError(
                    f"Range parameter '{param_name}' must define an explicit "
                    "'value' alongside 'bounds'."
                )
            if "choices" in param_data and "value" not in param_data:
                raise ValueError(
                    f"Choice parameter '{param_name}' must define an explicit "
                    "'value' alongside 'choices'."
                )
            parsed = parse_parameter(param_name, param_data)
        else:
            raise TypeError(
                "Parameters must be BaseParameter instances or dictionaries"
            )

        if parsed.name != param_name:
            parsed = parsed.model_copy(update={"name": param_name})
        parsed_parameters[param_name] = parsed

    self.parameters = parsed_parameters

    constraints_input = self.constraints or []
    self.constraints = [
        c if isinstance(c, ParameterConstraint) else ParameterConstraint(**c)
        for c in constraints_input
    ]

from_design_config(design_config) classmethod

Build a search space from a DesignConfig instance.

Parameters:

Name	Type	Description	Default
design_config	DesignConfig	Fully validated design configuration containing parameters and optional parameter constraints.	required

Returns:

Type	Description
SearchSpace	SearchSpace populated with flattened parameters and constraints.

Source code in src/aid2e/optimizers/base.py

@classmethod
def from_design_config(cls, design_config: DesignConfig) -> "SearchSpace":
    """Build a search space from a DesignConfig instance.

    Args:
        design_config: Fully validated design configuration containing
            parameters and optional parameter constraints.

    Returns:
        SearchSpace populated with flattened parameters and constraints.
    """

    return cls(
        parameters=design_config.get_flat_parameters(),
        constraints=design_config.parameter_constraints or [],
        name=getattr(design_config, "name", None),
        source_config=design_config,
    )

validate(param_values)

Check if parameter values satisfy all constraints (for non-Ax optimizers).

This method is for optimizers that DON'T have native constraint support (e.g., random search, simple evolutionary algorithms). For optimizers with native constraint support (e.g., Ax), pass self.constraints directly to the optimizer instead of calling this method.

Parameters:

Name	Type	Description	Default
param_values	Dict[str, Any]	Mapping of qualified parameter names to concrete values.	required

Returns:

Type	Description
bool	Tuple of (all_valid, failed_constraints) where all_valid is
List[str]	True when every constraint passes and failed_constraints
Tuple[bool, List[str]]	lists the names (or error strings) of failing constraints.

Example

For optimizers WITHOUT constraint support

is_valid, failures = search_space.validate(candidate) if not is_valid: ... print(f"Constraint violations: {failures}")

For Ax (HAS constraint support) - DON'T use this method

Instead, pass search_space.constraints to Ax optimizer

Notes

• Only use this for runtime checking with constraint-agnostic optimizers
• Ax and similar optimizers handle constraints internally via self.constraints
• Syntax validation already done at DesignConfig load time

Source code in src/aid2e/optimizers/base.py

def validate(self, param_values: Dict[str, Any]) -> Tuple[bool, List[str]]:
    """Check if parameter values satisfy all constraints (for non-Ax optimizers).

    This method is for optimizers that DON'T have native constraint support
    (e.g., random search, simple evolutionary algorithms). For optimizers
    with native constraint support (e.g., Ax), pass self.constraints directly
    to the optimizer instead of calling this method.

    Args:
        param_values: Mapping of qualified parameter names to concrete values.

    Returns:
        Tuple of ``(all_valid, failed_constraints)`` where ``all_valid`` is
        ``True`` when every constraint passes and ``failed_constraints``
        lists the names (or error strings) of failing constraints.

    Example:
        >>> # For optimizers WITHOUT constraint support
        >>> is_valid, failures = search_space.validate(candidate)
        >>> if not is_valid:
        ...     print(f"Constraint violations: {failures}")

        >>> # For Ax (HAS constraint support) - DON'T use this method
        >>> # Instead, pass search_space.constraints to Ax optimizer

    Notes:
        - Only use this for runtime checking with constraint-agnostic optimizers
        - Ax and similar optimizers handle constraints internally via self.constraints
        - Syntax validation already done at DesignConfig load time
    """

    if not self.constraints:
        return True, []

    failed: List[str] = []
    for constraint in self.constraints:
        try:
            if not constraint.evaluate(param_values):
                failed.append(constraint.name)
        except Exception as exc:  # pragma: no cover - defensive logging
            failed.append(f"{constraint.name} (error: {exc})")

    return len(failed) == 0, failed

Trial dataclass

Capture the parameters and results of a single optimization trial.

Attributes:

Name	Type	Description
index	int	Unique trial identifier within the optimizer.
parameters	Dict[str, Any]	Parameter values evaluated during the trial.
metrics	Optional[Dict[str, float]]	Objective values produced by evaluation (if available).
metadata	Dict[str, Any]	Optional auxiliary metadata about the trial.
status	str	Lifecycle status such as pending, completed, or failed.

Examples:

>>> trial = Trial(
...     index=0,
...     parameters={"x": 0.5},
...     metrics={"loss": 0.1},
...     status="completed",
... )
>>> trial.metadata
{}

Source code in src/aid2e/optimizers/base.py

@dataclass
class Trial:
    """Capture the parameters and results of a single optimization trial.

    Attributes:
        index: Unique trial identifier within the optimizer.
        parameters: Parameter values evaluated during the trial.
        metrics: Objective values produced by evaluation (if available).
        metadata: Optional auxiliary metadata about the trial.
        status: Lifecycle status such as ``pending``, ``completed``, or ``failed``.

    Examples:
        >>> trial = Trial(
        ...     index=0,
        ...     parameters={"x": 0.5},
        ...     metrics={"loss": 0.1},
        ...     status="completed",
        ... )
        >>> trial.metadata
        {}
    """

    index: int
    parameters: Dict[str, Any]
    metrics: Optional[Dict[str, float]] = None
    metadata: Dict[str, Any] = None
    status: str = "pending"

    def __post_init__(self) -> None:
        """Normalize metadata and status values after initialization."""
        if self.metadata is None:
            self.metadata = {}
        normalized = str(self.status).strip().lower()
        self.status = normalized if normalized else TRIAL_STATUS_PENDING
        if self.status not in VALID_TRIAL_STATUSES:
            logger.warning(
                "Unknown trial status '%s'; keeping value as-is.",
                self.status,
            )

    def save_to_json(self, output_path: Union[str, Path]) -> Path:
        """Write the trial's design parameters to disk as pretty-printed JSON.

        This helper is intentionally design-point focused: it serializes only
        ``parameters`` so command-line evaluators can consume the resulting file
        as an input payload without needing optimizer metadata.

        Args:
            output_path: Target JSON path.

        Returns:
            Resolved path of the written file.
        """
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as handle:
            json.dump(dict(self.parameters or {}), handle, indent=2, sort_keys=True)
        return path

__post_init__()

Normalize metadata and status values after initialization.

Source code in src/aid2e/optimizers/base.py

def __post_init__(self) -> None:
    """Normalize metadata and status values after initialization."""
    if self.metadata is None:
        self.metadata = {}
    normalized = str(self.status).strip().lower()
    self.status = normalized if normalized else TRIAL_STATUS_PENDING
    if self.status not in VALID_TRIAL_STATUSES:
        logger.warning(
            "Unknown trial status '%s'; keeping value as-is.",
            self.status,
        )

save_to_json(output_path)

Write the trial's design parameters to disk as pretty-printed JSON.

This helper is intentionally design-point focused: it serializes only parameters so command-line evaluators can consume the resulting file as an input payload without needing optimizer metadata.

Parameters:

Name	Type	Description	Default
output_path	Union[str, Path]	Target JSON path.	required

Returns:

Type	Description
Path	Resolved path of the written file.

Source code in src/aid2e/optimizers/base.py

def save_to_json(self, output_path: Union[str, Path]) -> Path:
    """Write the trial's design parameters to disk as pretty-printed JSON.

    This helper is intentionally design-point focused: it serializes only
    ``parameters`` so command-line evaluators can consume the resulting file
    as an input payload without needing optimizer metadata.

    Args:
        output_path: Target JSON path.

    Returns:
        Resolved path of the written file.
    """
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(dict(self.parameters or {}), handle, indent=2, sort_keys=True)
    return path

compute_pareto_front(trials, objective_names, objective_directions=None)

Extract non-dominated (Pareto-optimal) trials from a collection of completed trials.

Objectives are compared using objective_directions. Minimization is the default; maximization objectives treat larger values as better. For single-objective problems the function returns the best trial under that objective's direction.

Parameters:

Name	Type	Description	Default
trials	List[Trial]	Iterable of Trial objects. Only trials whose status is "completed" and whose metrics dict is non-empty are considered.	required
objective_names	List[str]	Ordered list of objective metric keys that must be present in each trial's metrics dict.	required
objective_directions	Optional[Dict[str, Any]]	Optional mapping from objective name to "minimize" or "maximize". Missing objectives default to minimization.	None

Returns:

Type	Description
List[Trial]	List of non-dominated Trial objects ordered by their original
List[Trial]	position in trials. Returns an empty list when no completed
List[Trial]	trials are available.

Examples:

>>> from aid2e.optimizers.base import Trial, compute_pareto_front
>>> t1 = Trial(index=0, parameters={}, metrics={"f1": 1.0, "f2": 3.0}, status="completed")
>>> t2 = Trial(index=1, parameters={}, metrics={"f1": 2.0, "f2": 1.0}, status="completed")
>>> t3 = Trial(index=2, parameters={}, metrics={"f1": 0.5, "f2": 2.0}, status="completed")
>>> front = compute_pareto_front([t1, t2, t3], ["f1", "f2"])
>>> {t.index for t in front}
{0, 1, 2}

Notes

Uses NumPy for vectorised comparisons when available; falls back to a pure-Python O(n²) loop otherwise. For production workloads with thousands of trials, the NumPy path is strongly recommended.

Source code in src/aid2e/optimizers/base.py

def compute_pareto_front(
    trials: List["Trial"],
    objective_names: List[str],
    objective_directions: Optional[Dict[str, Any]] = None,
) -> List["Trial"]:
    """Extract non-dominated (Pareto-optimal) trials from a collection of completed trials.

    Objectives are compared using ``objective_directions``. Minimization is the
    default; maximization objectives treat larger values as better. For
    single-objective problems the function returns the best trial under that
    objective's direction.

    Args:
        trials: Iterable of Trial objects.  Only trials whose ``status`` is
            ``"completed"`` and whose ``metrics`` dict is non-empty are
            considered.
        objective_names: Ordered list of objective metric keys that must be
            present in each trial's ``metrics`` dict.
        objective_directions: Optional mapping from objective name to
            ``"minimize"`` or ``"maximize"``. Missing objectives default to
            minimization.

    Returns:
        List of non-dominated Trial objects ordered by their original
        position in ``trials``.  Returns an empty list when no completed
        trials are available.

    Examples:
        >>> from aid2e.optimizers.base import Trial, compute_pareto_front
        >>> t1 = Trial(index=0, parameters={}, metrics={"f1": 1.0, "f2": 3.0}, status="completed")
        >>> t2 = Trial(index=1, parameters={}, metrics={"f1": 2.0, "f2": 1.0}, status="completed")
        >>> t3 = Trial(index=2, parameters={}, metrics={"f1": 0.5, "f2": 2.0}, status="completed")
        >>> front = compute_pareto_front([t1, t2, t3], ["f1", "f2"])
        >>> {t.index for t in front}
        {0, 1, 2}

    Notes:
        Uses NumPy for vectorised comparisons when available; falls back to
        a pure-Python O(n²) loop otherwise.  For production workloads with
        thousands of trials, the NumPy path is strongly recommended.
    """
    completed = [
        t for t in trials
        if t is not None and t.status == "completed" and t.metrics
    ]
    if not completed:
        return []
    if len(completed) == 1:
        return completed

    directions = objective_directions or {}
    objective_signs = []
    for obj in objective_names:
        direction = getattr(directions.get(obj), "value", directions.get(obj, "minimize"))
        objective_signs.append(-1.0 if str(direction).lower() == "maximize" else 1.0)

    def score(trial: "Trial", obj: str, sign: float) -> float:
        value = trial.metrics.get(obj)
        return float("inf") if value is None else float(value) * sign

    try:
        import numpy as np

        n = len(completed)
        F = np.array(
            [
                [score(t, obj, sign) for obj, sign in zip(objective_names, objective_signs)]
                for t in completed
            ],
            dtype=float,
        )
        is_dominated = np.zeros(n, dtype=bool)
        for i in range(n):
            if is_dominated[i]:
                continue
            # Vectorised: does any other solution dominate i?
            # j dominates i iff F[j] <= F[i] (all) and F[j] < F[i] (any)
            dom_mask = np.all(F <= F[i], axis=1) & np.any(F < F[i], axis=1)
            dom_mask[i] = False  # Exclude self-comparison
            if np.any(dom_mask):
                is_dominated[i] = True

        return [t for t, dom in zip(completed, is_dominated) if not dom]

    except ImportError:  # pragma: no cover – NumPy should always be present
        logger.warning("NumPy not available; falling back to pure-Python Pareto computation.")
        n = len(completed)
        is_dominated = [False] * n
        for i in range(n):
            if is_dominated[i]:
                continue
            for j in range(n):
                if i == j or is_dominated[j]:
                    continue
                # Check whether j dominates i
                all_leq = all(
                    score(completed[j], obj, sign) <= score(completed[i], obj, sign)
                    for obj, sign in zip(objective_names, objective_signs)
                )
                any_lt = any(
                    score(completed[j], obj, sign) < score(completed[i], obj, sign)
                    for obj, sign in zip(objective_names, objective_signs)
                )
                if all_leq and any_lt:
                    is_dominated[i] = True
                    break

        return [t for t, dom in zip(completed, is_dominated) if not dom]