File indexing completed on 2026-08-12 08:24:57
0001 """Integration tests for BaseOptimizer and the strict Ax optimizer surface."""
0002
0003 import pytest
0004 from ax.adapter.registry import Generators
0005 from ax.generators.torch.botorch_modular.surrogate import SurrogateSpec
0006
0007 from aid2e.optimizers import (
0008 AxOptimizer,
0009 AxOptimizerConfig,
0010 BaseOptimizer,
0011 SearchSpace,
0012 Trial,
0013 )
0014 from aid2e.optimizers.ax import optimizer as ax_optimizer_module
0015
0016
0017 AX_NODE_RUNTIME_AVAILABLE = ax_optimizer_module.AX_NODE_STRATEGY_AVAILABLE
0018
0019
0020 def _search_space() -> SearchSpace:
0021 """Create a canonical search space used across optimizer tests."""
0022 return SearchSpace(
0023 parameters={
0024 "x": {"type": "range", "value": 0.5, "bounds": [0.0, 1.0]},
0025 "y": {"type": "range", "value": 0.0, "bounds": [-1.0, 1.0]},
0026 }
0027 )
0028
0029
0030 def test_base_optimizer_is_abstract():
0031 """Test that BaseOptimizer cannot be instantiated directly."""
0032 search_space = SearchSpace(
0033 parameters={"x": {"type": "range", "value": 0.5, "bounds": [0.0, 1.0]}}
0034 )
0035
0036 with pytest.raises(TypeError, match="Can't instantiate abstract class"):
0037 BaseOptimizer(search_space=search_space, n_objectives=1)
0038
0039
0040 def test_ax_optimizer_inherits_from_base():
0041 """Test that AxOptimizer inherits from BaseOptimizer."""
0042 assert issubclass(AxOptimizer, BaseOptimizer)
0043
0044
0045 def test_ax_optimizer_requires_node_generation_runtime():
0046 """Ax should fail fast when the installed runtime lacks node generation APIs."""
0047 config = AxOptimizerConfig(seed=42)
0048
0049 if AX_NODE_RUNTIME_AVAILABLE:
0050 optimizer = AxOptimizer(
0051 search_space=_search_space(),
0052 config=config,
0053 objective_names=["loss"],
0054 )
0055 assert optimizer.generation_strategy is not None
0056 return
0057
0058 with pytest.raises(RuntimeError, match="node-based generation API required"):
0059 AxOptimizer(
0060 search_space=_search_space(),
0061 config=config,
0062 objective_names=["loss"],
0063 )
0064
0065
0066 @pytest.mark.skipif(
0067 not AX_NODE_RUNTIME_AVAILABLE,
0068 reason="Installed Ax runtime lacks required node-based generation APIs.",
0069 )
0070 def test_ax_optimizer_with_config():
0071 """Test AxOptimizer integrates properly with the strict Ax config surface."""
0072 config = AxOptimizerConfig(
0073 initialization_strategy="sobol",
0074 generator="BOTORCH_MODULAR",
0075 n_initial_samples=10,
0076 n_iterations=50,
0077 batch_size=5,
0078 seed=42,
0079 )
0080
0081 optimizer = AxOptimizer(
0082 search_space=_search_space(),
0083 config=config,
0084 objective_names=["loss", "time"],
0085 )
0086
0087 assert optimizer.n_objectives == 2
0088 assert optimizer.seed == 42
0089 assert optimizer.config.initialization_strategy == "sobol"
0090 assert optimizer.config.generator == "BOTORCH_MODULAR"
0091 assert optimizer.objective_names == ["loss", "time"]
0092
0093
0094 @pytest.mark.skipif(
0095 not AX_NODE_RUNTIME_AVAILABLE,
0096 reason="Installed Ax runtime lacks required node-based generation APIs.",
0097 )
0098 def test_ax_optimizer_builds_default_sobol_to_mbm_strategy():
0099 """Test that the default strategy is Sobol followed by Modular BoTorch."""
0100 config = AxOptimizerConfig(
0101 initialization_strategy="sobol",
0102 n_initial_samples=6,
0103 seed=42,
0104 )
0105
0106 optimizer = AxOptimizer(
0107 search_space=_search_space(),
0108 config=config,
0109 objective_names=["f1", "f2"],
0110 )
0111
0112 names_and_nodes = optimizer.generation_strategy.nodes_by_name
0113 names = list(names_and_nodes.keys())
0114 nodes = list(names_and_nodes.values())
0115 assert names == ["Sobol", "ModularBoTorch"]
0116 sobol_spec = nodes[0].generator_specs[0]
0117 model_spec = nodes[1].generator_specs[0]
0118
0119 assert sobol_spec.generator_enum == Generators.SOBOL
0120 assert sobol_spec.generator_kwargs == {"seed": 42}
0121 assert nodes[0].transition_criteria[0].threshold == 6
0122 assert model_spec.generator_enum == Generators.BOTORCH_MODULAR
0123 assert model_spec.generator_kwargs == {}
0124
0125
0126 @pytest.mark.skipif(
0127 not AX_NODE_RUNTIME_AVAILABLE,
0128 reason="Installed Ax runtime lacks required node-based generation APIs.",
0129 )
0130 def test_ax_optimizer_resolves_symbolic_generator_kwargs():
0131 """Test that YAML-friendly generator kwargs are resolved at runtime."""
0132 config = AxOptimizerConfig(
0133 initialization_strategy="sobol",
0134 generator="BOTORCH_MODULAR",
0135 objective_thresholds={"f1": 1.0, "f2": 1.0},
0136 generator_kwargs={
0137 "botorch_acqf_class": "qLogNoisyExpectedHypervolumeImprovement",
0138 "surrogate_spec": {
0139 "model_configs": [{"botorch_model_class": "SingleTaskGP"}]
0140 },
0141 },
0142 generator_gen_kwargs={
0143 "model_gen_options": {
0144 "optimizer_kwargs": {"sequential": False, "num_restarts": 5}
0145 }
0146 },
0147 seed=42,
0148 )
0149
0150 optimizer = AxOptimizer(
0151 search_space=_search_space(),
0152 config=config,
0153 objective_names=["f1", "f2"],
0154 )
0155
0156 nodes = list(optimizer.generation_strategy.nodes_by_name.values())
0157 model_spec = nodes[1].generator_specs[0]
0158 assert model_spec.generator_kwargs["botorch_acqf_class"].__name__ == (
0159 "qLogNoisyExpectedHypervolumeImprovement"
0160 )
0161 assert isinstance(model_spec.generator_kwargs["surrogate_spec"], SurrogateSpec)
0162 assert optimizer.optimization_config.objective_thresholds[0].bound == 1.0
0163 assert model_spec.generator_gen_kwargs["model_gen_options"]["optimizer_kwargs"][
0164 "sequential"
0165 ] is False
0166
0167
0168 @pytest.mark.skipif(
0169 not AX_NODE_RUNTIME_AVAILABLE,
0170 reason="Installed Ax runtime lacks required node-based generation APIs.",
0171 )
0172 def test_ax_optimizer_supports_center_initialization():
0173 """Test center initialization inserts a center node before Sobol."""
0174 config = AxOptimizerConfig(
0175 initialization_strategy="center",
0176 n_initial_samples=4,
0177 seed=42,
0178 )
0179
0180 optimizer = AxOptimizer(
0181 search_space=_search_space(),
0182 config=config,
0183 objective_names=["loss"],
0184 )
0185
0186 nodes = list(optimizer.generation_strategy.nodes_by_name.values())
0187 assert nodes[0].__class__.__name__ == "CenterGenerationNode"
0188 assert nodes[1].name == "Sobol"
0189 assert nodes[1].transition_criteria[0].threshold == 3
0190
0191
0192 @pytest.mark.skipif(
0193 not AX_NODE_RUNTIME_AVAILABLE,
0194 reason="Installed Ax runtime lacks required node-based generation APIs.",
0195 )
0196 def test_ax_optimizer_supports_uniform_initialization():
0197 """Test uniform initialization uses Ax's UNIFORM generator when available."""
0198 config = AxOptimizerConfig(
0199 initialization_strategy="uniform",
0200 n_initial_samples=5,
0201 seed=42,
0202 )
0203
0204 optimizer = AxOptimizer(
0205 search_space=_search_space(),
0206 config=config,
0207 objective_names=["loss"],
0208 )
0209
0210 nodes = list(optimizer.generation_strategy.nodes_by_name.values())
0211 first_node_enum = nodes[0].generator_specs[0].generator_enum
0212 assert first_node_enum in {Generators.UNIFORM, Generators.SOBOL}
0213
0214
0215 @pytest.mark.skipif(
0216 not AX_NODE_RUNTIME_AVAILABLE,
0217 reason="Installed Ax runtime lacks required node-based generation APIs.",
0218 )
0219 def test_ax_optimizer_suggest_and_update():
0220 """Test basic suggest and update workflow."""
0221 config = AxOptimizerConfig(
0222 initialization_strategy="sobol",
0223 n_initial_samples=5,
0224 seed=42,
0225 )
0226
0227 optimizer = AxOptimizer(
0228 search_space=_search_space(),
0229 config=config,
0230 objective_names=["loss"],
0231 )
0232
0233 candidates = optimizer.suggest_candidates(n_candidates=3)
0234 assert len(candidates) == 3
0235 assert all("x" in c and "y" in c for c in candidates)
0236
0237 for i, candidate in enumerate(candidates):
0238 optimizer.update_with_results(
0239 trial_index=i,
0240 parameters=candidate,
0241 metrics={"loss": 0.5 * i},
0242 )
0243
0244 trials = optimizer.get_trials()
0245 assert len(trials) == 3
0246 assert all(t.status == "completed" for t in trials)
0247
0248 best = optimizer.get_best_trial()
0249 assert best is not None
0250 assert best.metrics["loss"] == 0.0
0251
0252
0253 @pytest.mark.skipif(
0254 not AX_NODE_RUNTIME_AVAILABLE,
0255 reason="Installed Ax runtime lacks required node-based generation APIs.",
0256 )
0257 def test_ax_optimizer_pareto_front_single_objective():
0258 """Test Pareto front for single objective returns best trial."""
0259 config = AxOptimizerConfig(seed=42)
0260 search_space = SearchSpace(
0261 parameters={"x": {"type": "range", "value": 0.5, "bounds": [0.0, 1.0]}}
0262 )
0263
0264 optimizer = AxOptimizer(
0265 search_space=search_space,
0266 config=config,
0267 objective_names=["loss"],
0268 )
0269
0270 candidates = optimizer.suggest_candidates(n_candidates=5)
0271 for i, candidate in enumerate(candidates):
0272 optimizer.update_with_results(
0273 trial_index=i,
0274 parameters=candidate,
0275 metrics={"loss": float(i)},
0276 )
0277
0278 pareto_front = optimizer.get_pareto_front()
0279 assert len(pareto_front) == 1
0280 assert pareto_front[0].metrics["loss"] == 0.0
0281
0282
0283 @pytest.mark.skipif(
0284 not AX_NODE_RUNTIME_AVAILABLE,
0285 reason="Installed Ax runtime lacks required node-based generation APIs.",
0286 )
0287 def test_ax_optimizer_serialize_deserialize():
0288 """Test state serialization and deserialization."""
0289 config = AxOptimizerConfig(
0290 seed=42,
0291 generator_kwargs={"botorch_acqf_class": "qLogNoisyExpectedImprovement"},
0292 )
0293 search_space = SearchSpace(
0294 parameters={"x": {"type": "range", "value": 0.5, "bounds": [0.0, 1.0]}}
0295 )
0296
0297 optimizer = AxOptimizer(
0298 search_space=search_space,
0299 config=config,
0300 objective_names=["loss"],
0301 )
0302
0303 candidates = optimizer.suggest_candidates(n_candidates=3)
0304 for i, candidate in enumerate(candidates):
0305 optimizer.update_with_results(
0306 trial_index=i,
0307 parameters=candidate,
0308 metrics={"loss": 0.1 * i},
0309 )
0310
0311 state = optimizer.serialize_state()
0312
0313 assert "search_space" in state
0314 assert "objective_names" in state
0315 assert "config" in state
0316 assert "trials" in state
0317 assert state["config"]["generator"] == "BOTORCH_MODULAR"
0318 assert state["config"]["generator_kwargs"]["botorch_acqf_class"] == (
0319 "qLogNoisyExpectedImprovement"
0320 )
0321
0322 optimizer2 = AxOptimizer(
0323 search_space=search_space,
0324 config=config,
0325 objective_names=["loss"],
0326 )
0327 optimizer2.load_state(state)
0328
0329 trials = optimizer2.get_trials()
0330 assert len(trials) == 3
0331 best = optimizer2.get_best_trial()
0332 assert best.metrics["loss"] == 0.0
0333
0334
0335 def test_search_space_and_trial_classes():
0336 """Test SearchSpace and Trial data classes."""
0337 search_space = SearchSpace(
0338 parameters={
0339 "x": {"type": "range", "value": 0.5, "bounds": [0.0, 1.0]},
0340 "y": {"type": "choice", "value": "a", "choices": ["a", "b", "c"]},
0341 }
0342 )
0343 assert len(search_space.parameters) == 2
0344 assert "x" in search_space.parameters
0345
0346 trial = Trial(
0347 index=0,
0348 parameters={"x": 0.5, "y": "a"},
0349 metrics={"loss": 0.1},
0350 status="completed",
0351 )
0352 assert trial.index == 0
0353 assert trial.status == "completed"
0354 assert trial.metadata == {}
0355
0356 trial2 = Trial(
0357 index=1,
0358 parameters={"x": 0.3},
0359 metadata={"note": "test"},
0360 )
0361 assert trial2.metadata["note"] == "test"
0362 assert trial2.metrics is None
0363
0364
0365 def test_search_space_rejects_legacy_parameter_shapes():
0366 """Legacy values/default coercions should no longer be accepted."""
0367 with pytest.raises(ValueError, match="retired key 'values'"):
0368 SearchSpace(
0369 parameters={"y": {"type": "choice", "value": "a", "values": ["a", "b"]}}
0370 )
0371
0372 with pytest.raises(ValueError, match="explicit 'value'"):
0373 SearchSpace(parameters={"x": {"type": "range", "bounds": [0.0, 1.0]}})
