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