GPU batched ensemble

Batches a (B simulations × N agents) ensemble into a single device pass — the natural fit for derivative-free calibration, where you evaluate thousands of small replicate runs.

This is not the single-run path. For one large simulation with the view API, use model.gpu().run() instead (GPU backend, Going faster (lanes)).

import ambr as am
import numpy as np

# Evaluate B parameter sets in one (B, N) pass (NumPy or CuPy).
B = 8
runner = am.GPUEnsembleRunner(am.BatchedWellMixedSIR())
traj = runner.run(
    n_agents=1_000,
    steps=20,
    params={
        "beta": np.linspace(0.1, 0.4, B),
        "gamma": np.full(B, 0.05),
        "i0_frac": np.full(B, 0.02),
    },
    seed=0,
)
# traj["I_frac"] has shape (B, steps)

GPU-batched ensemble runner for AMBER (CuPy, NumPy fallback).

Runs B independent simulations – different parameter sets and/or seeds – as a single (B, N) tensor batch. This is the natural GPU upgrade for calibration / sensitivity analysis / parameter sweeps (Experiment, SMACOptimizer, ParallelRunner), where you evaluate thousands of small replicate runs: batching collapses B*steps kernel launches into steps, saturates the device, and runs ~100-600x faster than looping (the speedup is largest when each individual run is small – exactly the calibration regime).

A batched model implements three methods over (B, N) arrays from the active array module (CuPy on GPU, else NumPy):

setup(B, N, params, rng) -> state (dict of (B, N) arrays) step(state, params, rng) -> state record(state) -> dict of per-run metrics, each shape (B,)

where params maps name -> (B, 1) array (one value per run, broadcast over the N agents). The runner stacks each recorded metric into a (B, steps) trajectory – so a calibration loss against a target curve is one more batched tensor op over all B candidates at once.

class ambr.gpu_ensemble.BatchedWellMixedSIR

Bases: object

Reference batched model: well-mixed (mean-field) SIR.

Per agent, per run: a susceptible becomes infected with probability beta * I_fraction (force of infection), an infected recovers with probability gamma. Per-run parameters beta, gamma, i0_frac are (B, 1) arrays. This is the canonical curve-fitting target in epidemiology and batches perfectly (O(B*N) per step).

INFECTED = 1

R = 2

S = 0

record(state)

setup(B, N, params, rng)

step(state, params, rng)

class ambr.gpu_ensemble.GPUEnsembleRunner(model)

Bases: object

Run a batched model’s B simulations together as a (B, N) batch.

run(n_agents: int, steps: int, params: Dict, seed: int = 0) → Dict

Run B simulations (B = len of each param array) for steps.

Returns a dict metric -> (B, steps) array of per-run trajectories.

ambr.gpu_ensemble.smac_batch_calibrate(model, param_bounds: Dict[str, Tuple[float, float]], loss_fn: Callable[[Dict], Any], n_agents: int, steps: int, *, rounds: int = 8, batch_size: int = 64, fixed_params: Dict[str, float] | None = None, seed: int = 0, quiet: bool = True) → Tuple[Dict[str, float] | None, list]

Batched Bayesian calibration: SMAC + the GPU ensemble.

Each round, SMAC’s ask/tell interface proposes batch_size distinct configurations; GPUEnsembleRunner evaluates all of them in one GPU pass as a (batch_size, N) tensor; the per-run losses are told back so SMAC’s surrogate improves for the next round. This combines SMAC’s sample-efficient search with GPU batch throughput – rounds * batch_size evaluations in only rounds sequential model-update steps and rounds GPU launches.

Parameters:

• model – a batched model (see module docstring).

• param_bounds – name -> (low, high) for the calibrated parameters.

• loss_fn – trajectories -> (B,) per-run losses (e.g. SSE vs a target).

• fixed_params – parameters held constant across all runs.

• rounds – search budget (total evals = rounds * batch_size).

• batch_size – search budget (total evals = rounds * batch_size).

Returns (best_config, history) where history is the per-round best loss.