mpBAX - Multipoint Bayesian Algorithm Execution

A lightweight framework for model-based multi-objective optimization with built-in checkpointing and reproducibility.

Key Features:

• Config-First Design: Entire experiment specified in config (dict or YAML)
• Multi-Oracle Support: Optimize multiple independent objectives with different input dimensions
• Checkpointing: Automatic save/resume/rollback functionality
• Reproducible: Built-in seed management
• Extensible: Easy to add custom models, algorithms, and generators

Installation

Option 1: pip install (Recommended)

Install in development mode (editable):

cd /path/to/mpBAX
pip install -e .

Install normally:

pip install .

With optional PyTorch support (for DANetModel):

pip install -e ".[torch]"

With development tools (pytest):

pip install -e ".[dev]"

All optional dependencies:

pip install -e ".[all]"

Option 2: PYTHONPATH (No installation)

If you prefer not to install:

export PYTHONPATH=/path/to/mpBAX:$PYTHONPATH

Dependencies:

• Core: numpy>=1.20.0, pyyaml>=5.0
• Optional: torch>=1.13.0 (for DANetModel plugin)
• Dev: pytest>=7.0 (for running tests)

Quick Start

Create a single-file optimization script:

importnumpyasnpfrommpbax.core.engineimportEnginefrommpbax.core.modelimportDummyModelfrommpbax.core.algorithmimportGreedySampling# Define oracle functiondefquadratic(X: np.ndarray) ->np.ndarray:
"""X has shape (n, d), returns Y with shape (n, 1)"""returnnp.sum((X-0.5)**2, axis=1, keepdims=True)
# Configure optimizationconfig= {
'seed': 42,
'max_loops': 10,
'checkpoint': {'dir': 'checkpoints', 'freq': 1},
'oracles': [{
'name': 'quadratic',
'input_dim': 2,
'n_initial': 20,
'function': {'class': quadratic}, # Pass function directly'model': {'class': DummyModel} # Pass class directly
}],
'algorithm': {
'class': GreedySampling,
'params': {'input_dims': [2], 'n_propose': 10, 'n_candidates': 1000}
}
}
# Run optimizationengine=Engine(config)
engine.run()

That's it! mpBAX will:

1. Generate 20 initial samples
2. Evaluate them with quadratic function
3. Train a model on the data
4. Use the model to propose 10 new candidates
5. Repeat for 10 loops total
6. Save checkpoints every loop

For larger projects, use YAML configs and import paths instead of instances:

# config.yamloracles:
- function:
class: 'myproject.oracles.quadratic'# Import path string

See examples/ for more examples.

Flexible Config Patterns

mpBAX v2 introduces flexible parameter placement for cleaner configs:

Default Generator Shortcut:

# Use default uniform generator with custom number of samples'oracles': [{
'input_dim': 2,
'generate': {'params': {'n': 20}}, # No need to specify 'class'
...
}]

Parameter-Level Placement:

# Put input_dim where it's used (in model params)'oracles': [{
'generate': {'params': {'n': 20, 'd': 2}},
'model': {
'class': DummyModel,
'params': {'input_dim': 2} # Input dim specified here
}
}]

Custom Generator with All Params:

# Pass all params (n, d, custom) to generator'generate': {
'class': lhs_generator,
'params': {'n': 20, 'd': 4, 'criterion': 'maximin'}
}

Cleaner Top-Level Config:

# Use 'training' instead of 'model' to avoid confusionconfig= {
'training': {'mode': 'finetune'}, # Clear and unambiguous'oracles': [...]
}

All old config patterns remain supported for backward compatibility.

Examples

See examples/ directory:

• 01_basic_optimization.py - Simple single-objective example
• 02_multi_oracle.py - Multi-objective optimization
• 03_multi_output.py - Oracle with multiple outputs
• 04_checkpointing.py - Resume and rollback
• 05_custom_model.py - Implementing custom models
• 06_danet_model.py - Using DANetModel plugin

Key Concepts

Oracle Function Requirements

Oracle functions are the core of mpBAX - they represent your expensive simulation or objective function.

Shape Convention:

• Input X: Shape (n, d) where n = number of samples, d = input dimension
• Output Y: Shape (n, k) where k ≥ 1 (number of outputs per sample)

Example - Single Output:

defmy_oracle(X: np.ndarray) ->np.ndarray:
""" Args: X: Input with shape (n, d) Returns: Y: Output with shape (n, 1) """returnnp.sum(X**2, axis=1, keepdims=True) # Shape (n, 1)

Example - Multi-Output:

defmulti_output_oracle(X: np.ndarray) ->np.ndarray:
""" Args: X: Input with shape (n, d) Returns: Y: Output with shape (n, 3) # Three outputs per input """y1=np.sum(X**2, axis=1, keepdims=True)
y2=np.sum(X, axis=1, keepdims=True)
y3=np.max(X, axis=1, keepdims=True)
returnnp.hstack([y1, y2, y3]) # Shape (n, 3)

Important: Always use keepdims=True or reshape to ensure Y is 2D, never 1D.

Config-First Architecture

Everything is specified in a config dict/YAML:

• Oracle functions: Import paths (YAML) or direct instances (Python)
• Models: Per-oracle model configuration with hyperparameters
• Generators: Custom initial sample generators (optional)
• Algorithm: Proposal strategy with parameters

Multi-Oracle Optimization

Each oracle operates independently with its own:

• Input dimension (can differ between oracles)
• Number of initial samples
• Model and hyperparameters
• Data storage and checkpoints

Algorithms receive predictions from all models and propose candidates for each oracle.

Multi-Output vs Multi-Oracle

Multi-Output: Single oracle returns multiple values

deforacle(X): # Shape (n, d)returnY# Shape (n, 3) - three outputs per input

Multi-Oracle: Multiple independent oracles

config['oracles'] = [
{'name': 'obj1', 'input_dim': 2, ...},
{'name': 'obj2', 'input_dim': 3, ...}
]

Checkpointing System

Automatic checkpointing enables:

• Resume: Continue from latest checkpoint
• Rollback: Delete checkpoints after a specific loop
• Reproducibility: Complete experiment state saved

Directory structure:

checkpoints/
├── config.yaml
├── state.pkl
├── oracle_name_1/
│ ├── data_0.pkl, data_1.pkl, ...
│ ├── model_0_final.pkl, model_1_final.pkl, ...
│ └── model_0_best.pkl, ... (if checkpoint_mode='both')
└── oracle_name_2/
└── ...

Configuration Reference

Complete Example

Here's a complete config demonstrating all available options:

# Top-level settingsseed: 42max_loops: 10# Checkpointing configurationcheckpoint:
dir: 'checkpoints'freq: 1resume_from: null # or 'latest' or loop number# Global training settings (renamed from 'model' to avoid confusion)training:
mode: 'retrain'# 'retrain' or 'finetune'checkpoint_mode: 'final'# 'final', 'best', or 'both'# Oracle configurations (list - one per oracle)oracles:
- name: 'my_objective'input_dim: 4# Can also be in model.params or generate.params.dn_initial: 20# Can also be in generate.params.n# Oracle functionfunction:
class: 'myproject.oracles.my_oracle'# or direct instanceparams: {} # Empty dict calls factory with no args# Optional: Custom generator (null = uniform [0,1]^d)generate: null# Alternative patterns:# generate:# params: {n: 20} # Use default generator with custom n# generate:# class: 'myproject.generators.lhs_generator'# params: {n: 20, d: 4, criterion: 'maximin'} # All params here# Model for this oraclemodel:
class: 'DummyModel'# or 'DANetModel' or customparams: {} # Can include input_dim here instead of oracle level# Algorithm configurationalgorithm:
class: 'GreedySampling'params:
input_dims: [4]n_propose: 10n_candidates: 1000

Top-Level Config

Field	Type	Description
seed	int	Random seed for reproducibility
max_loops	int	Maximum optimization loops
checkpoint	dict	Checkpointing configuration
training	dict	Global training settings (formerly model)
oracles	list	Oracle configurations (see below)
algorithm	dict	Algorithm configuration

Note:training was renamed from model to avoid confusion with per-oracle model field. The old name is still supported with a deprecation warning.

Checkpoint Config

Field	Type	Default	Description
dir	str	required	Checkpoint directory path
freq	int	1	Save frequency (1=every loop)
resume_from	str/int/null	null	Resume point ('latest', loop number, or null)

Usage examples:

Resume from latest checkpoint:

config['checkpoint']['resume_from'] ='latest'

Resume from specific loop:

config['checkpoint']['resume_from'] =5

Rollback to previous checkpoint:

frommpbax.core.checkpointimportCheckpointManagermanager=CheckpointManager('checkpoints')
manager.delete_checkpoints_after(loop=3) # Delete loops 4+

Training Config (Global)

Field	Type	Default	Description
mode	str	'retrain'	Training mode: 'retrain' or 'finetune'
checkpoint_mode	str	'final'	Model checkpoint: 'final', 'best', or 'both'

Training Modes:

Retrain (default) - Creates fresh model instance each loop, trains on all accumulated data from scratch. Simple and suitable for most models.

config['training'] = {'mode': 'retrain'}

Finetune - Reuses model instance from previous loop, continues training without resetting weights. Preserves normalization parameters (e.g., X_mu, X_sigma from loop 0). Ideal for neural networks.

config['training'] = {'mode': 'finetune'}

Checkpoint Modes:

• 'final': Save model after training completes (used for resumption)
• 'best': Save best model if tracked via get_best_model_snapshot()
• 'both': Save both final and best models

Oracle Config

Field	Type	Description
name	str	Oracle identifier (used for checkpoint dirs)
input_dim	int (optional)	Input dimensionality (can be in model.params or generate.params.d)
n_initial	int (optional)	Number of initial samples (can be in generate.params.n)
function	dict	{'class': func_or_path, 'params': {}}
generate	dict/null	Custom generator or null for uniform [0,1]^d
model	dict	{'class': model_class_or_path, 'params': {}}

Flexible Parameter Placement:

The framework supports multiple ways to specify input_dim and n_initial:

# Pattern 1: Traditional (oracle level)oracles:
- input_dim: 4n_initial: 20# Pattern 2: input_dim in model.paramsoracles:
- model:
params: {input_dim: 4}# Pattern 3: n in generate.params (default generator shortcut)oracles:
- input_dim: 4generate:
params: {n: 20}# Pattern 4: All params in custom generatororacles:
- generate:
class: 'myproject.generators.custom'params: {n: 20, d: 4, scale: 0.5}model:
params: {input_dim: 4}

At least one location must specify each parameter. The framework validates that required params exist somewhere in the config.

Algorithm Config

Field	Type	Description
class	str/class	Algorithm class or import path
params	dict	Algorithm parameters

Built-in algorithms:

• RandomSampling: Random proposals
• GreedySampling: Greedy selection from candidates (requires n_candidates)

Advanced Features

Sample Weighting

Models can access loop tracking metadata to weight recent data. This is useful for adaptive optimization where recent samples are more informative about promising regions.

deftrain(self, X, Y, metadata=None):
ifmetadataand'loop_indices'inmetadata:
loop_indices=metadata['loop_indices']
max_loop=np.max(loop_indices)
# Weight recent data higher (e.g., 10x)weights=np.where(loop_indices==max_loop, 10.0, 1.0)
# Use weights in loss function

The DataHandler automatically tracks which loop each sample came from, making this information available to models during training.

See DANetModel for a complete implementation example using PyTorch with sample weighting.

Developer Guide

Project Structure

mpBAX/
├── mpbax/ # Core package
│ ├── core/ # Core components
│ │ ├── engine.py # Main orchestrator
│ │ ├── evaluator.py # Oracle wrapper
│ │ ├── data_handler.py # Data storage
│ │ ├── model.py # Model interface
│ │ ├── algorithm.py # Algorithm interface
│ │ └── checkpoint.py # Checkpointing
│ └── plugins/ # Extensions
│ └── models/ # Model plugins
│ └── da_net.py # Neural network model
├── examples/ # Example scripts
├── tests/ # Test suite
├── README.md
└── config.yaml # Config template

Design Principles

• Simplicity: Use the simplest and most robust approach
• Minimal Dependencies: numpy, pyyaml, pickle - that's it
• No Assumptions: Common API, no isinstance checks
• Shape Conventions: X is (n, d), Y is (n, k) where k ≥ 1

Extending mpBAX

Custom Model:

frommpbax.core.modelimportBaseModelclassMyModel(BaseModel):
def__init__(self, input_dim, my_param=1.0):
super().__init__(input_dim)
self.my_param=my_paramdeftrain(self, X, Y, metadata=None):
# Training logic hereself.is_trained=Truedefpredict(self, X):
# Prediction logic herereturnY_pred# Shape (n, k)

Use in config:

'model': {
'class': MyModel, # or 'myproject.models.MyModel''params': {'my_param': 2.0}
}

Custom Algorithm:

frommpbax.core.algorithmimportBaseAlgorithmclassMyAlgorithm(BaseAlgorithm):
def__init__(self, input_dims, n_propose, my_param=1.0):
super().__init__(input_dims, n_propose)
self.my_param=my_paramdefpropose(self, fn_pred_list):
# fn_pred_list: list of predict functions (one per oracle)# Return: list of X arrays (one per oracle)X_list= []
fori, fn_predinenumerate(fn_pred_list):
X_candidates=self.generate_candidates(self.input_dims[i])
# Selection logic hereX_proposed=self.select_best(X_candidates, fn_pred)
X_list.append(X_proposed)
returnX_list

Custom Generator:

Generators are functions that generate initial samples. They receive all params during the actual call (not during instantiation):

deflhs_generator(n, d, criterion='maximin'):
"""Latin Hypercube Sampling generator. Args: n: Number of samples to generate d: Input dimensionality criterion: LHS sampling criterion Returns: X with shape (n, d) """# Generate samples using LHSreturnX

Use in config:

'generate': {
'class': lhs_generator, # or import path'params': {
'n': 20, # Number of samples'd': 4, # Dimensionality'criterion': 'maximin'# Custom parameter
}
}

The generator receives all params specified in the config during the call. This allows flexible signatures beyond the standard (n, d) pattern.

Testing

Run tests:

# All tests
python -m pytest tests/
# Specific component
python tests/test_model.py
# With pytest
pytest tests/ -v

Test structure:

• test_data_handler.py - DataHandler tests
• test_evaluator.py - Evaluator tests
• test_model.py - Model interface tests
• test_algorithm.py - Algorithm tests
• test_checkpoint.py - Checkpointing tests
• test_engine.py - Integration tests

See tests/README.md for details.

Plugins

DANetModel - PyTorch neural network surrogate:

• Automatic input normalization (preserved across loops)
• Sample weighting for recent data
• Adaptive epochs (initial vs finetuning)
• Early stopping
• GPU support

See mpbax/plugins/models/README.md for full documentation.

Adding a Plugin:

1. Create file in mpbax/plugins/{category}/
2. Inherit from appropriate base class
3. Implement required methods
4. Document in plugin README