Bayes Opt

Material Property Optimizer

Bayesian optimization for material and molecular properties using Ax (v1.2.4) and BoTorch.

When to Use

• Optimizing polymer sequences for target properties (Rg, end-to-end distance)
• Material composition optimization (alloys, mixtures)
• Molecular structure optimization
• Any expensive black-box function where evaluations require simulations/experiments
• Problems with mixed discrete/continuous parameter spaces
• Multi-objective optimization with Pareto front exploration

When NOT to use:

• Cheap-to-evaluate functions (use scipy.optimize or grid search)
• Problems with known analytical gradients (use gradient-based optimizers)
• Pure hyperparameter tuning for ML models (use Optuna or Ray Tune)
• Problems with >50 dimensions (GP surrogate scales poorly)

Quick Start

import sys
from pathlib import Path

# If running inside the skill directory:
sys.path.insert(0, str(Path(__file__).parent / "scripts"))
# If running from a project directory:
# sys.path.insert(0, str(Path.home() / ".claude/skills/bayes-opt/scripts"))

from ax_optimizer import AxOptimizer

# 1. Define parameter space
param_space = {
    "bead_0": {"type": "choice", "values": ["A", "B", "C"], "is_ordered": False},
    "bead_1": {"type": "choice", "values": ["A", "B", "C"], "is_ordered": False},
    "temperature": {"type": "range", "bounds": [300.0, 500.0]},
}

# 2. Define objective
objective = {"name": "rg_error", "mode": "minimize"}

# 3. Create optimizer and run
optimizer = AxOptimizer(param_space=param_space, objective=objective, max_trials=50)

result = optimizer.optimize(
    evaluation_fn=lambda params: {"rg_error": run_simulation(params)},
)

# 4. Get best result
best_params, best_value = optimizer.get_best_parameters()

Quick Reference

| Feature | API | |---------|-----| | Minimize | {"name": "energy", "mode": "minimize"} | | Maximize | {"name": "conductivity", "mode": "maximize"} | | Hit target | {"name": "rg", "mode": "target", "target_value": 5.0} | | Multi-objective | [{"name": "strength", "mode": "maximize"}, {"name": "cost", "mode": "minimize"}] | | Pareto front | optimizer.get_pareto_frontier() | | Checkpoint | optimizer.save_checkpoint("state.json") | | Resume | AxOptimizer.load_checkpoint("state.json") | | Generation strategy | generation_method="quality" or "fast" or "random_search" |

Parameter Types

# Categorical (polymer beads, crystal structures)
{"type": "choice", "values": ["A", "B", "C"], "is_ordered": False}

# Continuous (temperature, composition)
{"type": "range", "bounds": [0.0, 1.0]}

# Integer
{"type": "range", "bounds": [10, 100], "value_type": "int"}

# Log-scale (spans orders of magnitude)
{"type": "range", "bounds": [1e-10, 1e-5], "log_scale": True}

# Fixed (not optimized, injected into every trial)
{"type": "fixed", "value": 1.0}

See references/parameter-config.md for complete reference.

Multi-Objective Optimization

Pass a list of objectives to get Pareto-optimal solutions:

objectives = [
    {"name": "strength", "mode": "maximize"},
    {"name": "cost", "mode": "minimize"},
]

optimizer = AxOptimizer(param_space=param_space, objective=objectives, max_trials=50)
result = optimizer.optimize(evaluation_fn=evaluate)

# Get Pareto frontier
for params, values, trial_idx, arm_name in optimizer.get_pareto_frontier():
    print(f"params={params}, values={values}")

LAMMPS Integration

from ax_optimizer import AxOptimizer

optimizer = AxOptimizer(
    param_space={f"bead_{i}": {"type": "choice", "values": ["A", "B"], "is_ordered": False}
                 for i in range(20)},
    objective={"name": "rg", "mode": "target", "target_value": 15.0},
)

def evaluate_lammps(params):
    sequence = [params[f"bead_{i}"] for i in range(20)]
    write_lammps_input(sequence, "input.lmp")  # user-defined
    run_simulation("input.lmp")                 # user-defined
    rg = analyze_trajectory("dump.lammpstrj")   # user-defined
    return {"rg": rg}

optimizer.optimize(evaluate_lammps, max_trials=50)

See scripts/lammps_interface.py for a helper class that generates bead-spring polymer LAMMPS inputs.

Advanced Features

Parallel Evaluation

trials = optimizer.get_next_trials(batch_size=4)
for trial_idx, params in trials:
    submit_job(trial_idx, params)  # submit all 4 simultaneously

Human-in-the-Loop

optimizer = AxOptimizer(
    param_space=param_space, objective=objective,
    human_in_the_loop=True,   # confirm before each trial
    auto_approve_first_n=5,   # auto-run first 5 (exploration)
)

Generation Strategy

optimizer = AxOptimizer(
    param_space=param_space, objective=objective,
    generation_method="quality",    # 'quality', 'fast', 'random_search'
    initialization_budget=10,       # number of initial Sobol trials
)

Dependencies

pip install ax-platform==1.2.4 botorch gpytorch

Common Mistakes

| Mistake | Fix | |---------|-----| | from material_property_optimizer import ... | No pip package exists. Use sys.path + from ax_optimizer import AxOptimizer | | Multi-objective as {"objectives": [...]} | Pass a list directly: objective=[{...}, {...}] | | generation_strategy=custom_gs | Use generation_method="quality" instead. Custom GenerationStep not supported. | | Checkpoint with .pkl extension | Checkpoints are JSON-based. Use .json extension. | | Composition fractions unconstrained | Ax has no built-in simplex constraint. Optimize N-1 fractions, derive the last, return penalty for invalid. | | log_scale on RangeParameterConfig | User config uses log_scale: True; the optimizer converts to scaling="log" internally. | | FixedParameterConfig | Does not exist in Ax 1.2.4. Use {"type": "fixed", "value": ...} in param_space. |

Troubleshooting

| Issue | Solution | |-------|----------| | GP fails to fit | Increase initialization_budget (more random trials) | | Optimization stuck | Check parameter bounds, widen if too tight | | is_ordered warning for choice params | Set "is_ordered": False explicitly for categorical variables | | Memory issues | Reduce batch size, save checkpoints periodically | | Multi-objective no Pareto front | Need enough trials (50+) for meaningful Pareto exploration |

Configuration Files

• references/parameter-config.md - All parameter types and options
• references/examples.md - Complete examples (polymer, alloy, multi-objective)
• references/api-reference.md - Full API documentation
• scripts/ax_optimizer.py - Core optimizer implementation
• scripts/lammps_interface.py - LAMMPS input generation helper

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: WuGroup-XJTLU
• Source: WuGroup-XJTLU/cc-skills-ZhenghaoWu-Group
• License: Apache-2.0

Install and usage instructions live in the source repository linked above.