Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-11-09 09:29:52

0001 #!/usr/bin/env python
0002 """Example demonstrating BaseOptimizer and AxOptimizer integration with DTLZ2.
0003 
0004 This script shows how to:
0005 1. Load optimizer configuration from JSON file
0006 2. Define a search space using the new BaseOptimizer interface
0007 3. Create an AxOptimizer instance
0008 4. Run multi-objective optimization on DTLZ2 benchmark
0009 5. Serialize and deserialize optimizer state
0010 
0011 DTLZ2 is a multi-objective test problem with a known Pareto front.
0012 For 2 objectives and M variables, the optimal solutions lie on the unit sphere.
0013 
0014 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0015 Homepage: https://aid2e.github.io/aid2e-framework
0016 Repository: https://github.com/aid2e/AID2E-framework.git
0017 """
0018 
0019 import os
0020 import json
0021 import numpy as np
0022 from pathlib import Path
0023 from aid2e.optimizers import BaseOptimizer, SearchSpace, AxOptimizer, AxOptimizerConfig
0024 from aid2e.utilities.configurations import OptimizerConfiguration
0025 
0026 
0027 def dtlz2(x_dict, n_objectives=2):
0028     """DTLZ2 multi-objective test problem.
0029     
0030     Args:
0031         x_dict: Dictionary of parameters (x1, x2, ..., x10)
0032         n_objectives: Number of objectives (default: 2)
0033     
0034     Returns:
0035         Dictionary with objective values {f1, f2, ...}
0036     
0037     Notes:
0038         For 2 objectives, the Pareto front lies on the unit circle.
0039         Optimal solutions have g(x) = 0, where g is the sum of squared
0040         deviations from 0.5 for decision variables after the first M-1.
0041     """
0042     # Convert dict to array
0043     x = np.array([x_dict[f'x{i+1}'] for i in range(len(x_dict))])
0044     k = len(x) - n_objectives + 1
0045     
0046     # g function (auxiliary function)
0047     g = np.sum((x[n_objectives-1:] - 0.5) ** 2)
0048     
0049     # Compute objectives
0050     objectives = {}
0051     for i in range(n_objectives):
0052         f = 1.0 + g
0053         for j in range(n_objectives - i - 1):
0054             f *= np.cos(x[j] * np.pi / 2.0)
0055         if i > 0:
0056             f *= np.sin(x[n_objectives - i - 1] * np.pi / 2.0)
0057         objectives[f'f{i+1}'] = f
0058     
0059     return objectives
0060 
0061 
0062 def main():
0063     print("=" * 70)
0064     print("AID2E BaseOptimizer + AxOptimizer: DTLZ2 Example")
0065     print("=" * 70)
0066     
0067     # 1. Load configuration from JSON file
0068     print("\n1. Loading configuration from JSON file...")
0069     config_path = Path(__file__).with_name("dtlz2_ax_python_api.json")
0070     with open(config_path, 'r') as f:
0071         config_data = json.load(f)
0072     
0073     print(f"   Config file: {config_path.name}")
0074     print(f"   Problem: {config_data['name']}")
0075     
0076     # 2. Parse configuration using OptimizerConfiguration
0077     print("\n2. Parsing optimization configuration...")
0078     optimizer_payload = config_data.get("optimizer", config_data)
0079     opt_config = OptimizerConfiguration(**optimizer_payload)
0080     objective_names = config_data.get("objectives", [])
0081     n_iterations = config_data.get("n_iterations", 30)
0082     
0083     print(f"   Optimizer: {opt_config.name}")
0084     print(f"   Strategy: {opt_config.parameters.get('initialization_strategy', 'sobol')}")
0085     print(f"   Generator: {opt_config.parameters.get('generator', 'BOTORCH_MODULAR')}")
0086     
0087     # 3. Create AxOptimizerConfig from parsed parameters
0088     print("\n3. Creating AxOptimizerConfig...")
0089     optimizer_params = opt_config.parameters
0090     ax_config = AxOptimizerConfig(
0091         initialization_strategy=optimizer_params.get('initialization_strategy', 'sobol'),
0092         generator=optimizer_params.get('generator', 'BOTORCH_MODULAR'),
0093         generator_kwargs=optimizer_params.get('generator_kwargs', {}),
0094         generator_gen_kwargs=optimizer_params.get('generator_gen_kwargs', {}),
0095         objective_thresholds=optimizer_params.get('objective_thresholds'),
0096         n_initial_samples=optimizer_params.get('n_initial_samples', 10),
0097         n_iterations=n_iterations,
0098         batch_size=optimizer_params.get('batch_size', 3),
0099         seed=optimizer_params.get('seed', 42)
0100     )
0101     print(f"   Initial samples: {ax_config.n_initial_samples}")
0102     print(f"   Total iterations: {ax_config.n_iterations}")
0103     print(f"   Batch size: {ax_config.batch_size}")
0104     
0105     # 4. Define search space from configuration parameters
0106     print("\n4. Defining search space from configuration...")
0107     search_space_params = {}
0108     parameters = config_data.get('parameters', {})
0109     for param_name, param_config in parameters.items():
0110         search_space_params[param_name] = {
0111             "type": "range",
0112             "value": 0.5,
0113             "bounds": param_config["bounds"]
0114         }
0115     
0116     search_space = SearchSpace(parameters=search_space_params)
0117     print(f"   Parameters: {list(search_space.parameters.keys())}")
0118     print(f"   Total dimensions: {len(search_space.parameters)}")
0119     
0120     # 5. Create AxOptimizer instance
0121     print("\n5. Creating AxOptimizer...")
0122     optimizer = AxOptimizer(
0123         search_space=search_space,
0124         config=ax_config,
0125         objective_names=objective_names,
0126         seed=ax_config.seed
0127     )
0128     print(f"   Optimizer: {optimizer}")
0129     print(f"   Objectives: {objective_names}")
0130     print(f"   Inherits from BaseOptimizer: {isinstance(optimizer, BaseOptimizer)}")
0131     
0132     # 6. Run multi-objective optimization loop
0133     print("\n6. Running multi-objective optimization on DTLZ2...")
0134     n_iterations = 5
0135     
0136     for iteration in range(n_iterations):
0137         print(f"\n   Iteration {iteration + 1}/{n_iterations}:")
0138         
0139         # Suggest candidates
0140         candidates = optimizer.suggest_candidates(n_candidates=ax_config.batch_size)
0141         print(f"   - Suggested {len(candidates)} candidates")
0142         
0143         # Get current trial count before adding results
0144         trial_start_idx = len(optimizer.experiment.trials) - len(candidates)
0145         
0146         # Evaluate candidates using DTLZ2
0147         for idx, candidate in enumerate(candidates):
0148             # Evaluate DTLZ2
0149             objectives = dtlz2(candidate, n_objectives=len(objective_names))
0150             
0151             # Update optimizer with results
0152             trial_idx = trial_start_idx + idx
0153             optimizer.update_with_results(
0154                 trial_index=trial_idx,
0155                 parameters=candidate,
0156                 metrics=objectives
0157             )
0158             
0159             # Display results
0160             x_vals = [f"{candidate[f'x{i+1}']:.3f}" for i in range(3)]  # Show first 3
0161             obj_vals = [f"{objectives[obj]:.4f}" for obj in objective_names]
0162             print(f"     Trial {trial_idx}: x=[{', '.join(x_vals)}, ...] → {dict(zip(objective_names, obj_vals))}")
0163     
0164     # 7. Get Pareto front
0165     print("\n7. Retrieving Pareto front...")
0166     pareto_front = optimizer.get_pareto_front()
0167     print(f"   Pareto front size: {len(pareto_front)}")
0168     
0169     if pareto_front:
0170         print("\n   Pareto-optimal solutions:")
0171         for i, trial in enumerate(pareto_front[:5]):  # Show first 5
0172             obj_vals = [f"{trial.metrics[obj]:.4f}" for obj in objective_names]
0173             print(f"     Solution {i+1}: {dict(zip(objective_names, obj_vals))}")
0174         
0175         if len(pareto_front) > 5:
0176             print(f"     ... and {len(pareto_front) - 5} more solutions")
0177     
0178     # 8. Get best trial (representative from Pareto front)
0179     print("\n8. Best trial (from Pareto front):")
0180     best_trial = optimizer.get_best_trial()
0181     if best_trial:
0182         print(f"   Objectives: {best_trial.metrics}")
0183         print(f"   (For DTLZ2, optimal Pareto front is on unit sphere)")
0184     
0185     # 9. Get all trials
0186     print(f"\n9. Total trials evaluated: {len(optimizer.get_trials())}")
0187     
0188     # 10. Serialize and deserialize state
0189     print("\n10. Testing state serialization...")
0190     state = optimizer.serialize_state()
0191     print(f"    Serialized state has {len(state['trials'])} trials")
0192     
0193     # Create new optimizer and load state
0194     optimizer2 = AxOptimizer(
0195         search_space=search_space,
0196         config=ax_config,
0197         objective_names=objective_names,
0198         seed=ax_config.seed
0199     )
0200     optimizer2.load_state(state)
0201     print(f"    Loaded state: {len(optimizer2.get_trials())} trials restored")
0202     
0203     pareto_front2 = optimizer2.get_pareto_front()
0204     print(f"    Pareto front after reload: {len(pareto_front2)} solutions")
0205     
0206     print("\n" + "=" * 70)
0207     print("DTLZ2 multi-objective optimization completed successfully!")
0208     print("=" * 70)
0209 
0210 
0211 if __name__ == "__main__":
0212     main()