IonSim

Differentiable JAX digital twin of a trapped ion, with Bayesian optimal experiment design for closed-loop calibration. The simulator models exact time evolution of a two-level atom coupled to quantized motion via laser pulses; the calibration stack wraps it in a particle filter and an experiment optimizer that designs each round's pulse sequence to maximally inform the posterior over the ion's physical parameters.

Before each measurement round, the optimizer uses gradient descent over a differentiable pulse-sequence parameterization to find the frequency, duration, and power profile that will tell it the most about the ion's parameters, given what it already knows. A particle filter with MCMC rejuvenation tracks the posterior, and a closed-loop duration schedule ramps sequence length as uncertainty shrinks so the design stays alias-free. On simulated ions drawn uniformly from the prior:

Accuracy target	Adaptive (EIG)	Textbook	Speedup
10 % of prior	500 shots	500 shots
5 %	700	1,500	2.1x
2 %	1,300	2,800	2.2x
1 %	2,100	4,600	2.2x

The gap widens at tighter targets because the adaptive strategy frontloads information: it learns 2.8 bits in its first 10 rounds versus 1.7 for the textbook protocol (66% more), then tapers once the posterior is already tight. Random designs never reach 1%.

What the optimizer actually designs

A textbook calibration cycles through three fixed phases: spectroscopy (sweep frequency to find the carrier), Rabi flopping (drive on-carrier to measure coupling), sideband flopping (drive the red sideband for eta and n_bar). Each phase runs one single-pulse measurement that targets one or two parameters and ignores the rest. The adaptive optimizer instead designs a multi-pulse sequence each round, choosing every frequency, duration, and phase to maximize expected information gain (EIG) over all five parameters at once. Here is what each strategy proposes at three stages of a representative calibration run, scored by EIG at the same posterior state:

Round 4, spectroscopy phase (24% error remaining)

	Pulse sequence	EIG
Textbook	2.4 µs @ -0.9 trap (single pulse, sweep near carrier)	0.58
Adaptive	0.2 µs @ -0.8 trap, 0.4 µs idle, 0.3 µs @ -1.0 trap, 0.5 µs @ +0.6 trap	1.30 (2.2x)

Round 21, Rabi phase (3.9% error remaining)

	Pulse sequence	EIG
Textbook	0.8 µs on carrier (single Rabi flop)	0.08
Adaptive	3.5 µs @ -0.5 trap, 2.1 µs idle, 2.5 µs @ +0.1 trap, 2.8 µs @ -0.4 trap	0.32 (3.9x)

Round 41, sideband phase (1.7% error remaining)

	Pulse sequence	EIG
Textbook	11.2 µs @ -1.0 trap (red sideband flop)	0.09
Adaptive	6.6 µs @ +1.6 trap, 6.8 µs idle, 6.8 µs @ +0.1 trap, 9.0 µs @ -0.1 trap	0.78 (9.0x)

Frequencies shown as detuning from the carrier in units of the trap frequency (so +1.0 = blue sideband, -1.0 = red sideband). The advantage grows as uncertainty shrinks: once the posterior is tight, single-purpose probes waste most of their budget re-measuring parameters that are already known, while the optimizer finds multi-frequency sequences that extract information about everything simultaneously.

Key Features

• Accurate Physics: Exact Hamiltonian evolution without the Lamb-Dicke approximation, cross-checked against QuTiP.
• Differentiable: Full JAX + Equinox integration. Gradients flow through the simulator into the design optimizer (L-BFGS / Adam over EIG or Fisher objectives) and into the MALA resampler that rejuvenates the particle filter.
• Bayesian Optimal Design: Pluggable design objectives (expected information gain, Fisher A/D/E/smooth-E criteria), duration-budget schedules, and experiment-selection strategies, compared head-to-head in the experiments/ harness.

Installation

Clone the repository and install with uv:

git clone <repository>
cd IonSim
uv sync

For optional GPU acceleration:

uv sync --extra nvidia   # CUDA
uv sync --extra amd      # ROCm

See the JAX installation page for supported accelerators by platform.

Quick Start

Simulating an Ion

import jax.numpy as jnp
from ionsim.core.controls import PulseSequence
from ionsim.digital_twin.physics import StaticTrappedIon, get_ground_state

ion = StaticTrappedIon(
    coupling_strength=0.2,       # Rabi frequency (MHz)
    eta=0.25,                    # Lamb-Dicke parameter
    trap_freq=2.,                # Trap frequency (MHz)
    atomic_transition_freq=1207.4958430,  # MHz
    n_cutoff=20,                 # Hilbert space truncation
)

rho0 = get_ground_state(ion.n_array, n_bar=0.5)

pulse = PulseSequence(
    durations=[0.1, 0.2, 0.1],           # µs
    frequencies=[1207., 0., 1207.],       # MHz
    powers=[1., 0., 1.],
    phases=[jnp.pi, 0., 0.],
)

result = ion.execute_sequence(rho0, pulse)

Running a Calibration Loop

import jax
import jax.numpy as jnp
from ionsim.optimal_calibration.hardware import SimulatedHardware
from ionsim.optimal_calibration.particle_filter import ParticleFilter
from ionsim.optimal_calibration.strategies import ManualStrategy
from ionsim.optimal_calibration.calibration_loop import run_calibration
from ionsim.optimal_calibration.parameters import IonParameter, array_by_parameter

true_theta = array_by_parameter({
    IonParameter.RABI_FREQ: 0.2, IonParameter.ETA: 0.1,
    IonParameter.TRAP_FREQ: 1.5, IonParameter.CARRIER_FREQ: 1207.2,
    IonParameter.N_BAR: 0.1,
})

bounds = array_by_parameter({
    IonParameter.RABI_FREQ: (0.1, 0.4), IonParameter.ETA: (0.0, 0.35),
    IonParameter.TRAP_FREQ: (1.3, 3.0), IonParameter.CARRIER_FREQ: (1204.2, 1210.2),
    IonParameter.N_BAR: (0.0, 0.3),
})

key = jax.random.key(0)
hardware = SimulatedHardware(true_theta=true_theta, n_cutoff=8)
pf = ParticleFilter(particle_count=2000, parameter_bounds=bounds, random_key=key)
strategy = ManualStrategy("textbook", ("spectroscopy", "rabi", "sideband"), total_duration_budget=2.0)

result = run_calibration(
    hardware=hardware, strategy=strategy, particle_filter=pf,
    n_cutoff=8, key=key, max_rounds=20, shots_per_round=1000,
)

print("Final estimate:", result.thetas[-1])
print("Final std:     ", result.stds[-1])

Experiments

The experiments/ harness compares calibration policies head-to-head on the simulated ion. Each experiment is a config file declaring a base configuration and a list of arms; the runner checkpoints per seed and writes a raw .npz plus a diagnostic dashboard.

uv run python -m experiments run eig_gmm --quick     # runs very small version to make sure things compile/run
uv run python -m experiments run eig_gmm             # normal run
uv run python -m experiments --help                  # all commands and flags

The output lands in results/<name>/: manifest.json, arrays.npz, dashboard.png.

To add an experiment, copy the closest configs/*.py and edit the arm list. The filename must match name.

EXPERIMENT = Experiment(
    name="my_experiment",
    question="Does X beat Y?",
    base=dict(n_iterations=30, resampler="seeded-mala"),
    quick=QUICK | dict(n_reps=300),
    arms=(
        ArmSpec("baseline", "textbook"),
        ArmSpec("mine", "adaptive", dict(objective="eig", n_theta_samples=50)),
    ),
    baseline="baseline",
)

An arm is a label, a strategy (adaptive, textbook, rabi, spectroscopy, random, two-stage), and the config overrides that define it. Everything else (objective, criterion, method, ...) is a config field, not a strategy.

Experiment	Question
adaptive_vs_manual	Does adaptive design beat naive fixed protocols, shot for shot?
eig_vs_fisher	Does the design criterion matter, holding the estimator fixed?
eig_gmm	Does a GMM resample surrogate fix carrier fringe-lock?
theta_samples_sweep	How many posterior draws does the adaptive design step need?
optimality_criteria	Which Fisher scalarization (A/D/E/smooth-E) designs better pulses?
budget_schedule	Does ramping the duration budget close the gap to textbook?
ramped_vs_textbook	Does a posterior-ramped duration budget beat the textbook recipe?
pruned_uncapped_ramp	Does the uncapped 2 µs posterior ramp beat textbook with a pruned optimizer?
bang_bang	Does releasing power after a bang-bang search improve calibration?

Notebooks

• calibration_w_particle_filter.ipynb seeing whether the particle filter plus pulse-sequence designer converges on the trap parameters.
• fisher_info.ipynb comparing the classical information gain from measurement sequences against the theoretical quantum limit. Do we lose anything by restricting to the Z-basis (bright-dark)?
• qutip_benchmark.ipynb compare the QuTiP implementation with the JAX version.

Package Structure

ionsim/
├── core/                   # Pulse-sequence representation shared across everything
├── digital_twin/           # JAX physics simulator (+ QuTiP reference for cross-checks)
└── optimal_calibration/    # Particle filter, experiment optimizer, and calibration loop

Running Tests

uv run pytest                            # full suite
uv run pytest tests/test_physics.py -v   # just the physics checks

Contributing

1. Branch from master: git checkout -b feature/my-feature
2. Add tests in tests/ for new functionality
3. Run uv run pytest to ensure all tests pass
4. Use ruff for linting; type hints throughout