EIC code displayed by LXR master/​AID2E-framework/​examples/​dtlz2_ax_optimizer_showcase.py	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2026-08-12 08:24:55

 """DTLZ2 Optimization with Ax Bayesian Optimizer.
 
 This example demonstrates the DAG Executor with Ax optimizer running the DTLZ2 problem
 in two different workflow configurations:
 
 Case 1: Single Branch - Both objectives computed in one Python function
 Case 2: Separate Branches - Each objective computed in a separate branch
 
 Configuration:
 - 10 Sobol initialization points
 - 10 Bayesian optimization iterations
 - Batch size of 3 (3 parallel evaluations per iteration)
 - Total evaluations: 10 + (10 * 3) = 40 points
 
 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
 """
 
 import numpy as np
 import json
 from pathlib import Path
 from typing import Dict, Any, List
 
 from aid2e.utilities.workflows import (
     DAGExecutor,
     WorkflowDefinition,
     BranchDefinition,
     StageDefinition,
     JobDefinition,
     JobContext,
 )
 from aid2e.utilities.configurations.objectives import (
     ObjectiveDefinition,
     ObjectiveDirection,
 )
 from aid2e.optimizers.base import SearchSpace
 from aid2e.optimizers.ax import AxOptimizer, AxOptimizerConfig
 
 
 # =============================================================================
 # DTLZ2 Problem Implementation (same as before)
 # =============================================================================
 
 def dtlz2_both_objectives(x: List[float]) -> Dict[str, float]:
     """Compute both DTLZ2 objectives in one function."""
     x = np.array(x)
     g = np.sum((x[1:] - 0.5) ** 2)
     f1 = (1 + g) * np.cos(x[0] * np.pi / 2) * np.cos(x[1] * np.pi / 2)
     f2 = (1 + g) * np.cos(x[0] * np.pi / 2) * np.sin(x[1] * np.pi / 2)
     return {"f1": float(f1), "f2": float(f2)}
 
 
 def dtlz2_f1_only(x: List[float]) -> float:
     """Compute only f1 objective of DTLZ2."""
     x = np.array(x)
     g = np.sum((x[1:] - 0.5) ** 2)
     f1 = (1 + g) * np.cos(x[0] * np.pi / 2) * np.cos(x[1] * np.pi / 2)
     return float(f1)
 
 
 def dtlz2_f2_only(x: List[float]) -> float:
     """Compute only f2 objective of DTLZ2."""
     x = np.array(x)
     g = np.sum((x[1:] - 0.5) ** 2)
     f2 = (1 + g) * np.cos(x[0] * np.pi / 2) * np.sin(x[1] * np.pi / 2)
     return float(f2)
 
 
 # =============================================================================
 # Python Evaluator Wrappers
 # =============================================================================
 
 def evaluate_both_objectives_wrapper(context: JobContext) -> Dict[str, float]:
     """Wrapper to evaluate both objectives from JobContext."""
     design_point = context.design_point
     x = [design_point['x1'], design_point['x2'], design_point['x3']]
     objectives = dtlz2_both_objectives(x)
     context.add_log(f"Design point: {x}")
     context.add_log(f"Objectives: {objectives}")
     context.xcom_push("objectives", objectives)
     return objectives
 
 
 def evaluate_f1_wrapper(context: JobContext) -> float:
     """Wrapper to evaluate f1 from JobContext."""
     design_point = context.design_point
     x = [design_point['x1'], design_point['x2'], design_point['x3']]
     f1 = dtlz2_f1_only(x)
     context.add_log(f"Design point: {x}")
     context.add_log(f"f1 = {f1}")
     context.xcom_push("f1", f1)
     return f1
 
 
 def evaluate_f2_wrapper(context: JobContext) -> float:
     """Wrapper to evaluate f2 from JobContext."""
     design_point = context.design_point
     x = [design_point['x1'], design_point['x2'], design_point['x3']]
     f2 = dtlz2_f2_only(x)
     context.add_log(f"Design point: {x}")
     context.add_log(f"f2 = {f2}")
     context.xcom_push("f2", f2)
     return f2
 
 
 # =============================================================================
 # Workflow Definitions
 # =============================================================================
 
 def create_single_branch_workflow() -> WorkflowDefinition:
     """Create workflow with single branch computing both objectives."""
     compute_job = JobDefinition(
         name="compute_objectives",
         command="python",
         payload={
             "evaluator_type": "python",
             "python_callable": evaluate_both_objectives_wrapper,
             "op_args": (),
             "op_kwargs": {},
         },
     )
     
     eval_stage = StageDefinition(name="evaluate", jobs=[compute_job])
     main_branch = BranchDefinition(name="main", stages=[eval_stage])
     
     workflow = WorkflowDefinition(
         name="dtlz2_ax_single_branch",
         description="DTLZ2 with Ax optimizer - single branch",
         branches=[main_branch],
         objectives=[
             ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE),
             ObjectiveDefinition(name="f2", direction=ObjectiveDirection.MINIMIZE),
         ],
     )
     
     return workflow
 
 
 def create_separate_branches_workflow() -> WorkflowDefinition:
     """Create workflow with separate branches for each objective."""
     # Branch 1: f1
     f1_job = JobDefinition(
         name="compute_f1",
         command="python",
         payload={
             "evaluator_type": "python",
             "python_callable": evaluate_f1_wrapper,
             "op_args": (),
             "op_kwargs": {},
         },
     )
     f1_stage = StageDefinition(name="evaluate_f1", jobs=[f1_job])
     f1_branch = BranchDefinition(name="f1_branch", stages=[f1_stage])
     
     # Branch 2: f2
     f2_job = JobDefinition(
         name="compute_f2",
         command="python",
         payload={
             "evaluator_type": "python",
             "python_callable": evaluate_f2_wrapper,
             "op_args": (),
             "op_kwargs": {},
         },
     )
     f2_stage = StageDefinition(name="evaluate_f2", jobs=[f2_job])
     f2_branch = BranchDefinition(name="f2_branch", stages=[f2_stage])
     
     workflow = WorkflowDefinition(
         name="dtlz2_ax_separate_branches",
         description="DTLZ2 with Ax optimizer - separate branches",
         branches=[f1_branch, f2_branch],
         objectives=[
             ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE),
             ObjectiveDefinition(name="f2", direction=ObjectiveDirection.MINIMIZE),
         ],
     )
     
     return workflow
 
 
 # =============================================================================
 # Optimization Functions
 # =============================================================================
 
 def run_case_1_single_branch():
     """Run Case 1: Single branch with Ax optimizer."""
     print("\n" + "="*80)
     print("CASE 1: Single Branch with Ax Bayesian Optimizer")
     print("="*80)
     
     # Create workflow
     workflow = create_single_branch_workflow()
     print(f"\n✓ Workflow: {workflow.name}")
     print(f"  Description: {workflow.description}")
     print(f"  Branches: {len(workflow.branches)}")
     print(f"  Objectives: {[obj.name for obj in workflow.objectives]}")
     
     # Create executor
     executor = DAGExecutor(
         workflow=workflow,
         base_output_dir="/tmp/dtlz2_ax_optimization/case1",
         log_level="WARNING",
     )
     
     # Create search space
     search_space = SearchSpace(
         parameters={
             "x1": {"type": "range", "bounds": [0.0, 1.0]},
             "x2": {"type": "range", "bounds": [0.0, 1.0]},
             "x3": {"type": "range", "bounds": [0.0, 1.0]},
         }
     )
     
     # Create Ax optimizer configuration
     ax_config = AxOptimizerConfig(
         initialization_strategy="sobol",
         n_initial_samples=10,
         batch_size=3,
         generator="BOTORCH_MODULAR",
         seed=42,
     )
     
     # Create Ax optimizer
     optimizer = AxOptimizer(
         search_space=search_space,
         config=ax_config,
         objective_names=["f1", "f2"],
         seed=42,
     )
     
     print(f"\n✓ Optimizer: Ax Bayesian Optimizer")
     print(f"  Initialization: {ax_config.initialization_strategy} ({ax_config.n_initial_samples} points)")
     print(f"  Generator: {ax_config.generator}")
     print(f"  Batch Size: {ax_config.batch_size}")
     print(f"  Total Iterations: 10 Bayesian iterations")
     
     # Optimization loop
     print(f"\n{'Iter':<6} {'Batch':<6} {'x1':<10} {'x2':<10} {'x3':<10} {'f1':<12} {'f2':<12} {'Phase':<15}")
     print("-" * 95)
     
     trial_index = 0
     
     # Phase 1: Sobol initialization (10 points in batches of 3)
     n_sobol_batches = int(np.ceil(ax_config.n_initial_samples / ax_config.batch_size))
     for batch in range(n_sobol_batches):
         # Get batch size (might be less than 3 for last batch)
         batch_size = min(ax_config.batch_size, ax_config.n_initial_samples - batch * ax_config.batch_size)
         
         # Suggest candidates
         candidates = optimizer.suggest_candidates(n_candidates=batch_size)
         
         # Evaluate each candidate
         for i, design_point in enumerate(candidates):
             objectives = executor.execute(design_point)
             optimizer.update_with_results(trial_index, design_point, objectives)
             
             print(f"{trial_index+1:<6} {batch+1:<6} {design_point['x1']:<10.4f} {design_point['x2']:<10.4f} "
                   f"{design_point['x3']:<10.4f} {objectives.get('f1', 0):<12.6f} "
                   f"{objectives.get('f2', 0):<12.6f} {'Sobol Init':<15}")
             
             trial_index += 1
     
     # Phase 2: Bayesian optimization (10 iterations x 3 points = 30 points)
     n_bayesian_iterations = 10
     for iteration in range(n_bayesian_iterations):
         # Suggest batch of candidates
         candidates = optimizer.suggest_candidates(n_candidates=ax_config.batch_size)
         
         # Evaluate each candidate
         for i, design_point in enumerate(candidates):
             objectives = executor.execute(design_point)
             optimizer.update_with_results(trial_index, design_point, objectives)
             
             print(f"{trial_index+1:<6} {iteration+1:<6} {design_point['x1']:<10.4f} {design_point['x2']:<10.4f} "
                   f"{design_point['x3']:<10.4f} {objectives.get('f1', 0):<12.6f} "
                   f"{objectives.get('f2', 0):<12.6f} {'Bayesian':<15}")
             
             trial_index += 1
     
     # Get Pareto front
     pareto_front = optimizer.get_pareto_front()
     
     print(f"\n✓ Optimization Complete!")
     print(f"  Total evaluations: {trial_index}")
     print(f"  Sobol initialization: {ax_config.n_initial_samples}")
     print(f"  Bayesian iterations: {n_bayesian_iterations}")
     print(f"  Pareto front points: {len(pareto_front)}")
     
     print(f"\nPareto Front (non-dominated points):")
     print(f"{'Trial':<8} {'x1':<10} {'x2':<10} {'x3':<10} {'f1':<12} {'f2':<12}")
     print("-" * 72)
     for trial in pareto_front[:10]:  # Show top 10
         dp = trial.parameters
         obj = trial.metrics if trial.metrics else {}
         print(f"{trial.index:<8} {dp.get('x1', 0):<10.4f} {dp.get('x2', 0):<10.4f} {dp.get('x3', 0):<10.4f} "
               f"{obj.get('f1', 0):<12.6f} {obj.get('f2', 0):<12.6f}")
     
     return optimizer, executor
 
 
 def run_case_2_separate_branches():
     """Run Case 2: Separate branches with Ax optimizer."""
     print("\n" + "="*80)
     print("CASE 2: Separate Branches with Ax Bayesian Optimizer")
     print("="*80)
     
     # Create workflow
     workflow = create_separate_branches_workflow()
     print(f"\n✓ Workflow: {workflow.name}")
     print(f"  Description: {workflow.description}")
     print(f"  Branches: {len(workflow.branches)} ({[b.name for b in workflow.branches]})")
     print(f"  Objectives: {[obj.name for obj in workflow.objectives]}")
     
     # Create executor
     executor = DAGExecutor(
         workflow=workflow,
         base_output_dir="/tmp/dtlz2_ax_optimization/case2",
         log_level="WARNING",
     )
     
     # Create search space
     search_space = SearchSpace(
         parameters={
             "x1": {"type": "range", "bounds": [0.0, 1.0]},
             "x2": {"type": "range", "bounds": [0.0, 1.0]},
             "x3": {"type": "range", "bounds": [0.0, 1.0]},
         }
     )
     
     # Create Ax optimizer configuration (same as Case 1)
     ax_config = AxOptimizerConfig(
         initialization_strategy="sobol",
         n_initial_samples=10,
         batch_size=3,
         generator="BOTORCH_MODULAR",
         seed=42,  # Same seed for reproducibility
     )
     
     # Create Ax optimizer
     optimizer = AxOptimizer(
         search_space=search_space,
         config=ax_config,
         objective_names=["f1", "f2"],
         seed=42,
     )
     
     print(f"\n✓ Optimizer: Ax Bayesian Optimizer")
     print(f"  Initialization: {ax_config.initialization_strategy} ({ax_config.n_initial_samples} points)")
     print(f"  Generator: {ax_config.generator}")
     print(f"  Batch Size: {ax_config.batch_size}")
     print(f"  Total Iterations: 10 Bayesian iterations")
     
     # Optimization loop
     print(f"\n{'Iter':<6} {'Batch':<6} {'x1':<10} {'x2':<10} {'x3':<10} {'f1':<12} {'f2':<12} {'Phase':<15}")
     print("-" * 95)
     
     trial_index = 0
     
     # Phase 1: Sobol initialization
     n_sobol_batches = int(np.ceil(ax_config.n_initial_samples / ax_config.batch_size))
     for batch in range(n_sobol_batches):
         batch_size = min(ax_config.batch_size, ax_config.n_initial_samples - batch * ax_config.batch_size)
         candidates = optimizer.suggest_candidates(n_candidates=batch_size)
         
         for i, design_point in enumerate(candidates):
             objectives = executor.execute(design_point)
             optimizer.update_with_results(trial_index, design_point, objectives)
             
             print(f"{trial_index+1:<6} {batch+1:<6} {design_point['x1']:<10.4f} {design_point['x2']:<10.4f} "
                   f"{design_point['x3']:<10.4f} {objectives.get('f1', 0):<12.6f} "
                   f"{objectives.get('f2', 0):<12.6f} {'Sobol Init':<15}")
             
             trial_index += 1
     
     # Phase 2: Bayesian optimization
     n_bayesian_iterations = 10
     for iteration in range(n_bayesian_iterations):
         candidates = optimizer.suggest_candidates(n_candidates=ax_config.batch_size)
         
         for i, design_point in enumerate(candidates):
             objectives = executor.execute(design_point)
             optimizer.update_with_results(trial_index, design_point, objectives)
             
             print(f"{trial_index+1:<6} {iteration+1:<6} {design_point['x1']:<10.4f} {design_point['x2']:<10.4f} "
                   f"{design_point['x3']:<10.4f} {objectives.get('f1', 0):<12.6f} "
                   f"{objectives.get('f2', 0):<12.6f} {'Bayesian':<15}")
             
             trial_index += 1
     
     # Get Pareto front
     pareto_front = optimizer.get_pareto_front()
     
     print(f"\n✓ Optimization Complete!")
     print(f"  Total evaluations: {trial_index}")
     print(f"  Sobol initialization: {ax_config.n_initial_samples}")
     print(f"  Bayesian iterations: {n_bayesian_iterations}")
     print(f"  Pareto front points: {len(pareto_front)}")
     
     print(f"\nPareto Front (non-dominated points):")
     print(f"{'Trial':<8} {'x1':<10} {'x2':<10} {'x3':<10} {'f1':<12} {'f2':<12}")
     print("-" * 72)
     for trial in pareto_front[:10]:  # Show top 10
         dp = trial.parameters
         obj = trial.metrics if trial.metrics else {}
         print(f"{trial.index:<8} {dp.get('x1', 0):<10.4f} {dp.get('x2', 0):<10.4f} {dp.get('x3', 0):<10.4f} "
               f"{obj.get('f1', 0):<12.6f} {obj.get('f2', 0):<12.6f}")
     
     return optimizer, executor
 
 
 # =============================================================================
 # Main Entry Point
 # =============================================================================
 
 if __name__ == "__main__":
     print("\n" + "="*80)
     print("DTLZ2 Multi-Objective Bayesian Optimization with Ax")
     print("="*80)
     print("\nConfiguration:")
     print("  • 10 Sobol initialization points")
     print("  • 10 Bayesian optimization iterations")
     print("  • Batch size of 3 (parallel evaluations)")
     print("  • Total evaluations: 10 + (10 × 3) = 40 points")
     print("\nWorkflow Configurations:")
     print("  1. Single branch - both objectives in one Python function")
     print("  2. Separate branches - each objective in different branch")
     print("\nDTLZ2 Problem:")
     print("  Variables: x1, x2, x3 in [0, 1]")
     print("  Objectives: f1, f2 (minimize both)")
     print("  Optimal Pareto front: x1 in [0, 1], x2 = x3 = 0.5")
     
     try:
         # Run both cases
         optimizer1, executor1 = run_case_1_single_branch()
         optimizer2, executor2 = run_case_2_separate_branches()
         
         # Summary
         print("\n" + "="*80)
         print("COMPARISON SUMMARY")
         print("="*80)
         
         print("\nCase 1 (Single Branch):")
         print(f"  Workflow: {executor1.workflow.name}")
         print(f"  Branches: {len(executor1.workflow.branches)}")
         print(f"  Total evaluations: {len(optimizer1.get_trials())}")
         print(f"  Pareto front size: {len(optimizer1.get_pareto_front())}")
         print(f"  Output directory: {executor1.output_dir}")
         
         print("\nCase 2 (Separate Branches):")
         print(f"  Workflow: {executor2.workflow.name}")
         print(f"  Branches: {len(executor2.workflow.branches)}")
         print(f"  Total evaluations: {len(optimizer2.get_trials())}")
         print(f"  Pareto front size: {len(optimizer2.get_pareto_front())}")
         print(f"  Output directory: {executor2.output_dir}")
         
         print("\n" + "="*80)
         print("✅ Ax Optimization Showcase Complete!")
         print("="*80)
         print("\nKey Results:")
         print("  • Both workflow configurations use identical Ax optimizer")
         print("  • Bayesian optimization with SAASBO surrogate model")
         print("  • qNEHVI acquisition for multi-objective optimization")
         print("  • Batch optimization with 3 parallel evaluations per iteration")
         print("  • Sobol initialization ensures good space exploration")
         
     except ImportError as e:
         print("\n" + "="*80)
         print("⚠️  ERROR: Ax Platform Not Installed")
         print("="*80)
         print(f"\n{e}")
         print("\nTo run this showcase, install Ax:")
         print("  pip install ax-platform")
         print("\nAlternatively, use the simple random optimizer showcase:")
         print("  python examples/dtlz2_optimizer_showcase.py")

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