| name | bayes-opt |
| description | Use when optimizing material compositions, polymer sequences, molecular structures, or any expensive-to-evaluate property where each evaluation requires simulation or experiment. Also use when the user mentions Bayesian optimization for materials, black-box optimization with categorical variables, or Pareto-front exploration for multi-objective material design. |
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
sys.path.insert(0, str(Path(__file__).parent / "scripts"))
from ax_optimizer import AxOptimizer
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]},
}
objective = {"name": "rg_error", "mode": "minimize"}
optimizer = AxOptimizer(param_space=param_space, objective=objective, max_trials=50)
result = optimizer.optimize(
evaluation_fn=lambda params: {"rg_error": run_simulation(params)},
)
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
{"type": "choice", "values": ["A", "B", "C"], "is_ordered": False}
{"type": "range", "bounds": [0.0, 1.0]}
{"type": "range", "bounds": [10, 100], "value_type": "int"}
{"type": "range", "bounds": [1e-10, 1e-5], "log_scale": True}
{"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)
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")
run_simulation("input.lmp")
rg = analyze_trajectory("dump.lammpstrj")
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)
Human-in-the-Loop
optimizer = AxOptimizer(
param_space=param_space, objective=objective,
human_in_the_loop=True,
auto_approve_first_n=5,
)
Generation Strategy
optimizer = AxOptimizer(
param_space=param_space, objective=objective,
generation_method="quality",
initialization_budget=10,
)
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