File indexing completed on 2026-08-12 08:24:55
0001 """Minimal Ax optimizer-only DTLZ2 example."""
0002
0003 from __future__ import annotations
0004
0005 import math
0006 import sys
0007 from pathlib import Path
0008 from typing import Any, Dict, List
0009
0010 REPO_ROOT = Path(__file__).resolve().parents[2]
0011 if str(REPO_ROOT) not in sys.path:
0012 sys.path.insert(0, str(REPO_ROOT))
0013
0014 from aid2e.utilities import build_optimizer_from_config
0015 from aid2e.utilities.configurations import load_config
0016
0017
0018 DEFAULT_CONFIG_PATH = REPO_ROOT / "examples" / "optimizers" / "dtlz2_ax_optimizer_only.yml"
0019
0020
0021 def ordered_dtlz_vector(parameters: Dict[str, Any]) -> List[float]:
0022 """Return ordered DTLZ decision variables from a flat parameter dict."""
0023 indexed: List[tuple[int, float]] = []
0024 for key, value in parameters.items():
0025 short_key = key.split(".")[-1].split("__")[-1]
0026 if short_key.startswith("x") and short_key[1:].isdigit():
0027 indexed.append((int(short_key[1:]), float(value)))
0028 indexed.sort(key=lambda item: item[0])
0029 return [value for _, value in indexed]
0030
0031
0032 def dtlz2_objectives(x: List[float]) -> Dict[str, float]:
0033 """Compute the 2-objective DTLZ2 function."""
0034 g = sum((value - 0.5) ** 2 for value in x[1:])
0035 factor = 1.0 + g
0036 f1 = factor * math.cos(x[0] * math.pi / 2.0)
0037 f2 = factor * math.sin(x[0] * math.pi / 2.0)
0038 return {"f1": float(f1), "f2": float(f2)}
0039
0040
0041 def main(argv: List[str]) -> int:
0042 config_file = Path(argv[1]).resolve() if len(argv) >= 2 else DEFAULT_CONFIG_PATH
0043 config = load_config(str(config_file))
0044 optimizer_config = config.optimizer.parse_algorithm_params()
0045 if optimizer_config is None:
0046 raise RuntimeError("No registered optimizer config model found for Ax example.")
0047
0048 optimizer = build_optimizer_from_config(config.problem, config.optimizer)
0049
0050 print("Backend: ax")
0051 print(f"Config: {config_file}")
0052 print(f"Objectives: {[objective.name for objective in config.problem.objectives]}")
0053 print(f"Design variables: {len(config.problem.design_config.get_flat_parameters())}")
0054 print(f"Generator: {optimizer_config.generator}")
0055 print(f"Initial samples: {optimizer_config.n_initial_samples}")
0056 print(f"Optimization iterations: {optimizer_config.n_iterations}")
0057 print(f"Batch size: {optimizer_config.batch_size}")
0058 print(
0059 f"\n{'Trial':<6} {'Phase':<10} {'x1':<10} {'x2':<10} {'x3':<10} "
0060 f"{'f1':<12} {'f2':<12}"
0061 )
0062 print("-" * 80)
0063
0064 remaining_init = optimizer_config.n_initial_samples
0065 init_batch = 0
0066 while remaining_init > 0:
0067 init_batch += 1
0068 current_batch = min(optimizer_config.batch_size, remaining_init)
0069 candidates = optimizer.suggest_candidates(n_candidates=current_batch)
0070 start_index = len(optimizer.get_trials()) - len(candidates)
0071 for offset, parameters in enumerate(candidates):
0072 trial_index = start_index + offset
0073 metrics = dtlz2_objectives(ordered_dtlz_vector(parameters))
0074 optimizer.update_with_results(
0075 trial_index=trial_index,
0076 parameters=parameters,
0077 metrics=metrics,
0078 )
0079 vector = ordered_dtlz_vector(parameters)
0080 print(
0081 f"{trial_index + 1:<6} {'init':<10} "
0082 f"{vector[0]:<10.4f} {vector[1]:<10.4f} {vector[2]:<10.4f} "
0083 f"{metrics['f1']:<12.6f} {metrics['f2']:<12.6f}"
0084 )
0085 remaining_init -= current_batch
0086
0087 for iteration in range(optimizer_config.n_iterations):
0088 candidates = optimizer.suggest_candidates(
0089 n_candidates=optimizer_config.batch_size
0090 )
0091 start_index = len(optimizer.get_trials()) - len(candidates)
0092 for offset, parameters in enumerate(candidates):
0093 trial_index = start_index + offset
0094 metrics = dtlz2_objectives(ordered_dtlz_vector(parameters))
0095 optimizer.update_with_results(
0096 trial_index=trial_index,
0097 parameters=parameters,
0098 metrics=metrics,
0099 )
0100 vector = ordered_dtlz_vector(parameters)
0101 print(
0102 f"{trial_index + 1:<6} {f'iter-{iteration + 1}':<10} "
0103 f"{vector[0]:<10.4f} {vector[1]:<10.4f} {vector[2]:<10.4f} "
0104 f"{metrics['f1']:<12.6f} {metrics['f2']:<12.6f}"
0105 )
0106
0107 results = optimizer.get_optimization_results()
0108 pareto_front = optimizer.get_pareto_front()
0109
0110 print("\nSummary")
0111 print(f"Trials recorded: {results['n_trials']}")
0112 print(f"Pareto points: {len(pareto_front)}")
0113 return 0
0114
0115
0116 if __name__ == "__main__":
0117 raise SystemExit(main(sys.argv))
