EIC code displayed by LXR master/​AID2E-framework/​docs/​optimizer-architecture.md	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

Warning, /AID2E-framework/docs/optimizer-architecture.md is written in an unsupported language. File is not indexed.

 # AID2E Optimizer Architecture
 
 ## Overview
 
 `aid2e.optimizers` exposes a shared optimizer interface and two concrete backends:
 Ax for Bayesian optimization and PyMOO for evolutionary optimization. The package
 also exports the shared search-space and trial data structures, plus the Pareto
 utility used by the base implementation.
 
 The current public export surface is:
 
 - `BaseOptimizer`
 - `SearchSpace`
 - `Trial`
 - `compute_pareto_front`
 - `AxOptimizer`
 - `AxOptimizerConfig`
 - `PyMOOProblem`
 - `PyMOOOptimizer`
 - `PyMOOOptimizerConfig`
 - `AID2EProblem`
 
 The design is intentionally split between:
 
 - shared bookkeeping and result utilities in `src/aid2e/optimizers/base.py`
 - Ax-specific configuration, symbol resolution, and candidate generation in
   `src/aid2e/optimizers/ax/`
 - PyMOO-specific ask/tell logic in `src/aid2e/optimizers/pymoo/`
 - canonical config and runtime builders in `src/aid2e/utilities/configurations/`
   and `src/aid2e/utilities/runtime_builders.py`
 
 ## Shared Abstractions
 
 ### `SearchSpace`
 
 `SearchSpace` is the optimizer-facing representation of the design domain. It
 stores:
 
 - `parameters`: a mapping from parameter name to `BaseParameter`
 - `constraints`: optional `ParameterConstraint` objects
 - `name`: optional search-space identifier
 - `source_config`: optional originating `DesignConfig`
 
 `SearchSpace.from_design_config()` flattens a validated `DesignConfig` into the
 optimizer form by using `DesignConfig.get_flat_parameters()` and forwarding the
 associated parameter constraints.
 
 The constructor accepts either `BaseParameter` objects or dictionaries that are
 parsed into parameters. The code rejects retired dictionary shapes, including the
 legacy `values` key, and requires a concrete `value` field alongside `bounds` or
 `choices`.
 
 `SearchSpace.validate()` is a backend-agnostic constraint check. It evaluates
 `ParameterConstraint` objects directly and is intended for backends that do not
 enforce constraints natively.
 
 ### `Trial`
 
 `Trial` is the in-memory record for one evaluated point. It stores:
 
 - `index`
 - `parameters`
 - `metrics`
 - `metadata`
 - `status`
 
 Statuses are normalized to lower case in `__post_init__`. The base code recognizes
 `pending`, `suggested`, `running`, `completed`, `failed`, `aborted`, and
 `cancelled`. Unknown values are preserved, but a warning is logged.
 
 ### `compute_pareto_front`
 
 `compute_pareto_front()` computes the non-dominated subset of completed trials
 using minimization semantics for every objective. It only considers trials whose
 status is `completed` and that have metrics. The function is used by the base
 implementation and is not Ax- or PyMOO-specific.
 
 ### `BaseOptimizer`
 
 `BaseOptimizer` owns the shared trial ledger and the default result utilities. It
 accepts either a `SearchSpace` or a `DesignConfig`, validates that the search
 space and objective list are non-empty, and stores the common `seed`.
 
 The backend-required contract is small. Only these methods remain abstract:
 
 - `suggest_candidates`
 - `update_with_results`
 - `serialize_state`
 - `load_state`
 
 Everything else is provided by the base class:
 
 - `get_trials()`
 - `set_trial_status()`
 - `get_optimization_results()`
 - `save_optimization_results()`
 - `seed_from_trials()`
 - `get_pareto_front()`
 - `get_best_trial()`
 
 `seed_from_trials()` is the backend-switch primitive. It appends external trials
 to the local history, assigns new sequential indices, and preserves the original
 source trial index in metadata under `source_index`. This updates history only;
 backends that need warm-start behavior must add their own override.
 
 `get_pareto_front()` delegates to `compute_pareto_front()`. `get_best_trial()`
 uses the first Pareto member for multi-objective cases and the lowest-valued
 trial for single-objective cases.
 
 `get_optimization_results()` and `save_optimization_results()` provide a stable
 JSON-friendly result export, including raw status and display status labels.
 
 ## Ax Backend
 
 ### Configuration
 
 `AxOptimizerConfig` is the backend config model for the Ax implementation. It is
 registered under the `ax` name in the optimizer config registry.
 
 Supported fields are:
 
 - `initialization_strategy`: `sobol`, `uniform`, or `center`
 - `generator`: currently validated to `BOTORCH_MODULAR`
 - `generator_kwargs`: runtime kwargs passed to the Ax generator spec
 - `generator_gen_kwargs`: generation-time kwargs passed through to Ax
 - `objective_thresholds`: optional multi-objective thresholds by metric name
 - `n_initial_samples`
 - `n_iterations`
 - `batch_size`
 - `seed`
 
 The config explicitly rejects retired fields such as `surrogate_model` and
 `acquisition_function`. The current surface expects the newer `generator` plus
 `generator_kwargs` split instead.
 
 `src/aid2e/optimizers/ax/_resolver.py` keeps the config YAML-friendly by resolving
 string names into the concrete Ax, BoTorch, and GPyTorch classes used by the
 modern Modular BoTorch generator. It resolves:
 
 - `acquisition_class`
 - `botorch_acqf_class`
 - `botorch_acqf_classes_with_options`
 - `surrogate_spec`
 
 The resolver currently knows about the supported model, transform, kernel,
 likelihood, and MLL classes listed in that module. When `surrogate_spec` is
 present, it also resolves nested `metric_to_model_configs` entries in addition
 to the top-level `model_configs`.
 
 ### Search-space and constraint mapping
 
 `AxOptimizer` converts `RangeParameter` values into Ax floating-point range
 parameters and `ChoiceParameter` values into Ax string choice parameters.
 
 Constraint handling is best effort. The backend attempts to translate
 `ParameterConstraint.rule` into an Ax linear inequality. Simple expressions are
 handled; unsupported expressions are logged and skipped rather than being forced
 into an incorrect shape.
 
 For multi-objective problems, the optimization config uses Ax objectives plus
 optional objective thresholds. The objective names are treated as minimization
 targets in the Ax model.
 
 ### Generation strategy
 
 The Ax backend uses the node-based Ax generation API. The constructor requires an
 Ax runtime that provides `CenterGenerationNode`, `GenerationNode`,
 `GeneratorSpec`, and `MinTrials`.
 
 The generation strategy is built as:
 
 - an initialization node using `sobol`, `uniform`, or `center`
 - a model-based node using `Generators.BOTORCH_MODULAR`
 
 `center` is handled as a center node followed by a Sobol initialization node when
 additional initialization samples are needed. `uniform` falls back to Sobol if
 the installed Ax runtime does not expose a uniform generator.
 
 The backend uses `resolve_generator_kwargs()` to turn YAML-friendly values into
 Ax runtime objects before passing them into the generator spec.
 
 ### Candidate lifecycle
 
 `suggest_candidates()` makes one Ax generation call for the requested batch, then
 normalizes the returned generator output into single-arm runs. Each generated arm
 creates a new Ax trial, is marked running in Ax, and is recorded in the base
 history as `suggested`.
 
 The method raises if Ax returns fewer generator runs than requested. The method
 also records strategy metadata such as the current Ax step index or node name
 when available.
 
 `update_with_results()` validates that every configured objective is present in
 the metrics payload, attaches the results to the Ax experiment as deterministic
 data, marks the Ax trial completed, and updates the shared `Trial` history.
 
 ### State model
 
 `serialize_state()` stores:
 
 - the search-space parameters and constraints
 - objective names
 - seed
 - backend config
 - recorded trials
 - the trial counter
 
 `load_state()` rebuilds the search space, Ax search space, optimization config,
 experiment, and generation strategy from the serialized payload, then replays
 completed trials into the experiment.
 
 The code does not serialize an Ax experiment object directly. The restoration path
 is therefore reconstruction-based rather than a byte-for-byte restore of Ax
 internals.
 
 ## PyMOO Backend
 
 ### Configuration
 
 `PyMOOOptimizerConfig` is registered under the `pymoo` name in the optimizer
 registry. It currently supports:
 
 - `algorithm`: optional explicit algorithm selection
 - `pop_size`
 - `n_offsprings`
 - `crossover_prob`
 - `crossover_eta`
 - `mutation_eta`
 - `n_iterations`
 - `n_partitions`
 - `seed`
 - `verbose`
 
 If `algorithm` is omitted, `resolve_algorithm()` selects `ga` for single-objective
 problems and `nsga2` otherwise. Explicit algorithms must match the objective
 count:
 
 - `ga` only for single-objective problems
 - `nsga2`, `nsga3`, and `moead` only for multi-objective problems
 
 ### `PyMOOProblem`
 
 `PyMOOProblem` is the public PyMOO `Problem` wrapper for the AID2E search space.
 It is structural only:
 
 - `decode_x()` converts a PyMOO float vector back into an AID2E parameter dict
 - `_evaluate()` always raises `NotImplementedError`
 
 The class is used for ask/tell workflows where evaluation happens outside the
 optimizer, not inside the PyMOO problem object.
 
 `AID2EProblem` remains available temporarily as a deprecated compatibility
 alias. Ax does not expose an equivalent public `Problem` wrapper.
 
 Choice parameters are encoded as continuous indices and rounded back to the
 nearest valid choice during decoding. Range parameters are passed through as
 floats. Any other parameter types are not supported by the current backend.
 
 ### Candidate lifecycle
 
 PyMOO follows an external-evaluation ask/tell flow:
 
 - `suggest_candidates()` calls `algorithm.ask()`
 - `update_with_results()` buffers one evaluation result
 - `_flush_generation()` calls `algorithm.tell()` once the full generation has
   reported back
 
 `suggest_candidates()` ignores `n_candidates` as a hard request. The actual batch
 size is determined by the PyMOO algorithm. When the hint does not match the
 generated batch, a debug message is logged.
 
 The backend keeps an in-flight generation buffer and refuses to start a new
 generation until the current one has been fully updated. Trials are inserted into
 the shared history with status `pending` when they are suggested and are updated
 to `completed` when results arrive.
 
 `seed_from_trials()` is overridden to guard against seeding while a generation is
 in flight. The method still only updates history; it does not warm-start the
 internal PyMOO population.
 
 ### Algorithm mapping
 
 The backend currently constructs one of four algorithms:
 
 - `ga`
 - `nsga2`
 - `nsga3`
 - `moead`
 
 `SBX`, polynomial mutation, and random sampling are used as the default operators.
 `nsga3` and `moead` build reference directions from `n_partitions` or `pop_size`
 as appropriate.
 
 ### Constraints and state
 
 Constraints are not forwarded into PyMOO at present. When the search space
 contains constraints, the optimizer logs a warning and continues, so constraint
 satisfaction is not guaranteed by this backend.
 
 `serialize_state()` stores the search space, config, objective names, trials,
 trial counter, generation count, resolved algorithm name, and, when possible, a
 base64-encoded pickle of the PyMOO algorithm state.
 
 `load_state()` reconstructs the search space and problem, then either restores
 the pickled algorithm state or rebuilds the algorithm from config and seed if the
 pickle is missing or cannot be loaded. Any in-flight generation state is cleared
 during load.
 
 ## Runtime Construction Path
 
 `OptimizerConfiguration` is the canonical top-level config object used by the
 framework. It stores the backend `name`, the optimizer `type`, and the raw
 `parameters` payload.
 
 The runtime builder resolves backend selection through `build_optimizer_from_config()`
 in `src/aid2e/utilities/runtime_builders.py`:
 
 - backend inference normalizes `name`, `type`, and the `algorithm` parameter and
   matches any of those tokens against the supported backend keywords
 - `ax`, `bo`, `mobo`, and `bayesian` map to the Ax backend
 - `pymoo`, `ga`, `nsga2`, `nsga3`, `moead`, and `evolutionary` map to the PyMOO backend
 
 For Ax, the builder instantiates `AxOptimizerConfig` from the raw parameter
 payload and passes `problem_cfg.design_config` into `AxOptimizer`. For PyMOO, it
 does the same with `PyMOOOptimizerConfig`. In both cases, the builder only
 constructs the optimizer object; batch sizing and iteration counts are consumed
 by the outer orchestration or runtime loop that drives repeated
 `suggest_candidates()` and `update_with_results()` calls.
 
 The optimizer constructors accept `DesignConfig` objects directly, so the runtime
 builder passes the design config rather than pre-building `SearchSpace`. The base
 class handles the conversion to `SearchSpace`.
 
 `optimization_registry.py` registers the backend config models lazily so the
 canonical optimizer config can validate backend-specific parameters without
 forcing every backend module to be imported up front.
 
 ## Code-Visible Limitations
 
 - The base class assumes minimization semantics throughout.
 - `compute_pareto_front()` only considers completed trials with metrics.
 - `AxOptimizerConfig.n_iterations` and `PyMOOOptimizerConfig.n_iterations` are
   configuration values used by outer execution flows; the optimizer classes do
   not enforce a stop condition internally, and the runtime builder does not
   consume them itself.
 - Ax constraint handling is best effort and only covers the constraint shapes
   that the parser can convert safely.
 - PyMOO does not currently enforce search-space constraints.
 - PyMOO only supports `RangeParameter` and `ChoiceParameter`.
 - The Ax backend currently requires the node-based Ax runtime; older Ax APIs are
   not used here.
 
 ## Extension Points
 
 To add a new backend, the current code path requires three pieces:
 
 - a new optimizer class implementing the four abstract methods on `BaseOptimizer`
 - a Pydantic config model registered through `optimization_registry.register()`
 - a runtime-builder branch that maps the canonical optimizer config to the new backend
 
 Shared history, Pareto computation, trial export, and backend-switch seeding can
 then be reused directly from `BaseOptimizer`.

This page was automatically generated by the 2.3.7 LXR engine. The LXR team