utilities/workflows/toy_evaluator.py

0001 from __future__ import annotations
0002 from typing import Dict
0003 import numpy as np
0004 from aid2e.optimizers.base import SearchSpace
0005 from aid2e.optimizers.ax import AxOptimizer, AxOptimizerConfig
0006
0007
0008 def eval_epic_b0(design: Dict[str, float]) -> Dict[str, float]:
0009     # placeholder for B0 resolution evaluation, to be replaced by real computations
0010     z1 = float(design["b0_tracker.layer1_z_cm"])
0011     z2 = float(design["b0_tracker.layer2_z_cm"])
0012     z3 = float(design["b0_tracker.layer3_z_cm"])
0013     z4 = float(design["b0_tracker.layer4_z_cm"])
0014     d12, d23, d34 = (z2-z1), (z3-z2), (z4-z3)
0015     lever = (z4 - z1)
0016     nonu = float(np.std([d12, d23, d34]))
0017     toy_res = float(lever - 10.0 * nonu)
0018     return {"b0_resolution": toy_res}
0019
0020 def run_epic_b0_toy_optimization(config, verbosity=0):
0021     # runner for toy optimization of ePIC B0 tracker design,
0022     # using Ax optimizer with a simple proxy objective function (to be replaced by real computations later)
0023     dc = config.problem.design_config
0024     params = {}
0025     for name in dc.get_parameter_names():
0026         bounds = dc.get_parameter_bounds(name)
0027         if bounds is not None:
0028             params[name] = {"type": "range", "bounds": [float(bounds[0]), float(bounds[1])]}
0029
0030     search_space = SearchSpace(parameters=params)
0031
0032     objective_names = [o.name for o in config.problem.objectives]
0033     optimizer_params = dict(config.optimizer.parameters or {})
0034
0035     ax_cfg = AxOptimizerConfig(
0036         initialization_strategy=optimizer_params.get("initialization_strategy", "sobol"),
0037         generator=optimizer_params.get("generator", "BOTORCH_MODULAR"),
0038         generator_kwargs=optimizer_params.get("generator_kwargs", {}),
0039         generator_gen_kwargs=optimizer_params.get("generator_gen_kwargs", {}),
0040         objective_thresholds=optimizer_params.get("objective_thresholds"),
0041         n_initial_samples=optimizer_params.get("n_initial_samples", 10),
0042         n_iterations=optimizer_params.get("n_iterations", 100),
0043         batch_size=optimizer_params.get("batch_size", 1),
0044         seed=optimizer_params.get("seed", 42),
0045     )
0046
0047     optimizer = AxOptimizer(
0048         search_space=search_space,
0049         config=ax_cfg,
0050         objective_names=objective_names,
0051         seed=42,
0052     )
0053
0054     trial = 0
0055     n_total = ax_cfg.n_initial_samples + ax_cfg.n_iterations * ax_cfg.batch_size
0056
0057     while trial < n_total:
0058         candidates = optimizer.suggest_candidates(n_candidates=ax_cfg.batch_size)
0059
0060         for design_point in candidates:
0061
0062             failures = []
0063             for c in (dc.parameter_constraints or []):
0064                 expr = str(c.rule)
0065                 for k in sorted(design_point.keys(), key=len, reverse=True):
0066                     expr = expr.replace(k, f"design_point[{k!r}]")
0067
0068                 if not bool(eval(expr, {"__builtins__": {}}, {"design_point": design_point})):
0069                     failures.append(c.name)
0070
0071             ok = len(failures) == 0
0072
0073             if not ok:
0074                 metrics = {"b0_resolution": -1e9}
0075             else:
0076                 metrics = eval_epic_b0(design_point)
0077
0078             optimizer.update_with_results(trial, design_point, metrics)
0079
0080             trial += 1
0081             if trial >= n_total:
0082                 break
0083
0084     best = optimizer.get_best_trial()
0085
0086     print("Best trial:")
0087     print(best.parameters)
0088     print(best.metrics)