Notebook 03 — Neural Network Surrogate for Beam Physics

CAS 2026 · AI for Medical Accelerators

In this notebook you will:

1. Use Cheetah to generate a training dataset (random quadrupole settings → beam properties)
2. Train an MLP surrogate that predicts beam properties in < 1 ms (vs. 2 ms per Cheetah call)
3. Measure the speedup and evaluate accuracy (R², MAE)
4. Use the surrogate inside a Bayesian optimisation loop — replacing Cheetah entirely
5. Bonus: add uncertainty quantification with a simple ensemble

Key idea: Tracking codes can take minutes to hours per run. A surrogate trained on thousands of simulation runs replaces them in milliseconds — enabling real-time control and rapid optimisation.

import sys
if 'google.colab' in sys.modules:
    !pip install -q cheetah-accelerator scikit-optimize plotly

import plotly.io as pio; pio.renderers.default = "notebook_connected"
import time
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import cheetah
from skopt import gp_minimize
from skopt.space import Real
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from sklearn.metrics import r2_score

torch.manual_seed(42)
np.random.seed(42)

PROTON_REST_MASS_MEV = 938.272
TOTAL_ENERGY_EV = (150.0 + PROTON_REST_MASS_MEV) * 1e6

print('Imports OK')

Imports OK

1 · Rebuild the Cheetah lattice (same as notebook 02)

segment = cheetah.Segment(elements=[
    cheetah.Drift(length=torch.tensor(0.5)),
    cheetah.Quadrupole(length=torch.tensor(0.1), name="QF"),
    cheetah.Drift(length=torch.tensor(1.0)),
    cheetah.Quadrupole(length=torch.tensor(0.1), name="QD"),
    cheetah.Drift(length=torch.tensor(0.5)),
    cheetah.Screen(name="monitor"),
])

beam_in = cheetah.ParameterBeam.from_twiss(
    energy=torch.tensor(TOTAL_ENERGY_EV),
    beta_x=torch.tensor(1.0),
    beta_y=torch.tensor(1.0),
    emittance_x=torch.tensor(1e-6),
    emittance_y=torch.tensor(1e-6),
    species=cheetah.Species("proton"),
)

print('Lattice and beam ready.')

Lattice and beam ready.

2 · Generate training data with Cheetah

We sweep over random quadrupole strengths (k₁_QF, k₁_QD) and record the beam properties at the monitor. This is the slow step — each Cheetah call takes around 1 ms on a typical CPU. 5 000 training samples typically complete in a few seconds.

N_TRAIN = 5_000
N_TEST  = 1_000
K1_MIN, K1_MAX = -15.0, 15.0   # T/m² — matches notebook 02 parameter range

def run_cheetah(k1_qf: float, k1_qd: float) -> np.ndarray:
    """Single Cheetah evaluation → [σ_x, σ_y] in mm.
    Note: μ_x = μ_y = 0 for our symmetric on-axis beam, so we only predict beam sizes.
    """
    segment.QF.k1 = torch.tensor(k1_qf, dtype=torch.float32)
    segment.QD.k1 = torch.tensor(k1_qd, dtype=torch.float32)
    b = segment.track(beam_in)
    return np.array([
        b.sigma_x.item() * 1e3,   # mm
        b.sigma_y.item() * 1e3,
    ])

# Generate extra samples to account for divergent beam configurations (possible NaN)
N_RAW = int((N_TRAIN + N_TEST) * 1.2)
rng = np.random.default_rng(42)
k1_raw = rng.uniform(K1_MIN, K1_MAX, size=(N_RAW, 2)).astype(np.float32)

print(f'Generating {N_RAW} Cheetah evaluations...')
t0 = time.perf_counter()
y_raw = np.array([run_cheetah(*p) for p in k1_raw], dtype=np.float32)
t_gen = time.perf_counter() - t0
print(f'Done in {t_gen:.1f} s  ({t_gen/N_RAW*1000:.2f} ms/call)')

# Filter samples where beam diverged (Cheetah returns NaN for invalid configs)
valid = np.all(np.isfinite(y_raw), axis=1)
k1_samples_all = k1_raw[valid]
y_all = y_raw[valid]
print(f'Valid samples: {valid.sum()} / {N_RAW}  ({100*valid.mean():.1f}%)')
assert len(y_all) >= N_TRAIN + N_TEST, "Too few valid samples — increase N_RAW or reduce K1_MAX"

# Train / test split
X_train = torch.tensor(k1_samples_all[:N_TRAIN])
y_train = torch.tensor(y_all[:N_TRAIN])
X_test  = torch.tensor(k1_samples_all[N_TRAIN:N_TRAIN + N_TEST])
y_test  = torch.tensor(y_all[N_TRAIN:N_TRAIN + N_TEST])

print(f'X_train: {X_train.shape}  →  y_train: {y_train.shape}')

Generating 7200 Cheetah evaluations...

Done in 6.0 s  (0.83 ms/call)
Valid samples: 7200 / 7200  (100.0%)
X_train: torch.Size([5000, 2])  →  y_train: torch.Size([5000, 2])

Explore the training data

# Plot σ_x as a function of (k1_QF, k1_QD)
fig = go.Figure(go.Scatter(
    x=k1_samples_all[:500, 0], y=k1_samples_all[:500, 1],
    mode='markers',
    marker=dict(
        color=y_all[:500, 0],   # σ_x in mm
        colorscale='Viridis',
        showscale=True,
        colorbar=dict(title='σ_x (mm)'),
        size=5, opacity=0.8,
    ),
))
fig.update_layout(
    title='Beam size σ_x as a function of quadrupole strengths (500 samples)',
    xaxis_title='k₁_QF (T/m²)', yaxis_title='k₁_QD (T/m²)',
    paper_bgcolor='#0d1b2a', plot_bgcolor='#111f30',
    font=dict(color='#e2e8f0'),
)
fig.show()

3 · Train the MLP surrogate

A 3-layer network with Tanh activations (smoother than ReLU for physics regression). Input: (k₁_QF, k₁_QD) · Output: (σ_x, σ_y) in mm.

class BeamSurrogate(nn.Module):
    """MLP surrogate: (k1_QF, k1_QD) → (σ_x, σ_y) in mm."""

    def __init__(self, n_in: int = 2, n_out: int = 2, width: int = 128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_in, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, n_out),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


surrogate = BeamSurrogate()
n_params  = sum(p.numel() for p in surrogate.parameters())
print(f'Surrogate parameters: {n_params:,}')

Surrogate parameters: 33,666

# Normalise inputs and outputs for stable training
X_mean, X_std = X_train.mean(0), X_train.std(0)
y_mean, y_std = y_train.mean(0), y_train.std(0)

def normalise_X(x): return (x - X_mean) / X_std
def normalise_y(y): return (y - y_mean) / y_std
def denormalise_y(y): return y * y_std + y_mean

X_train_n = normalise_X(X_train)
y_train_n = normalise_y(y_train)
X_test_n  = normalise_X(X_test)

EPOCHS = 300
BATCH  = 256

loader    = DataLoader(TensorDataset(X_train_n, y_train_n), batch_size=BATCH, shuffle=True)
optimizer = torch.optim.Adam(surrogate.parameters(), lr=3e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)

train_losses, val_losses = [], []
for epoch in range(EPOCHS):
    surrogate.train()
    epoch_loss = 0.0
    for X_b, y_b in loader:
        pred = surrogate(X_b)
        loss = nn.functional.mse_loss(pred, y_b)
        optimizer.zero_grad(); loss.backward(); optimizer.step()
        epoch_loss += loss.item() * len(X_b)
    train_losses.append(epoch_loss / N_TRAIN)
    scheduler.step()

    surrogate.eval()
    with torch.no_grad():
        val_pred = surrogate(X_test_n)
        val_loss = nn.functional.mse_loss(val_pred, normalise_y(y_test)).item()
    val_losses.append(val_loss)

    if (epoch + 1) % 50 == 0:
        print(f'Epoch {epoch+1:3d}/{EPOCHS}  train={train_losses[-1]:.6f}  val={val_loss:.6f}')

print('Training complete.')

Epoch  50/300  train=0.000214  val=0.000193

Epoch 100/300  train=0.000044  val=0.000065

Epoch 150/300  train=0.000013  val=0.000008

Epoch 200/300  train=0.000003  val=0.000003

Epoch 250/300  train=0.000002  val=0.000002

Epoch 300/300  train=0.000001  val=0.000001
Training complete.

4 · Evaluate accuracy

surrogate.eval()
with torch.no_grad():
    y_pred = denormalise_y(surrogate(X_test_n)).numpy()
y_true = y_test.numpy()

OUTPUT_NAMES = ['σ_x', 'σ_y']
print('Accuracy on held-out test set:')
print(f'{"Output":<6}  {"R²":>6}  {"MAE (mm)":>10}')
print('-' * 28)
for i, name in enumerate(OUTPUT_NAMES):
    r2  = r2_score(y_true[:, i], y_pred[:, i])
    mae = np.abs(y_true[:, i] - y_pred[:, i]).mean()
    print(f'{name:<6}  {r2:>6.4f}  {mae:>10.4f}')

Accuracy on held-out test set:
Output      R²    MAE (mm)
----------------------------
σ_x     1.0000      0.0011
σ_y     1.0000      0.0012

fig = make_subplots(rows=1, cols=2,
    subplot_titles=['σ_x: predicted vs. Cheetah', 'σ_y: predicted vs. Cheetah'])

for col, (i, name) in enumerate([(0, 'σ_x'), (1, 'σ_y')], 1):
    lo = min(y_true[:, i].min(), y_pred[:, i].min())
    hi = max(y_true[:, i].max(), y_pred[:, i].max())
    fig.add_trace(go.Scatter(
        x=y_true[:, i], y=y_pred[:, i], mode='markers',
        marker=dict(color='#22d3ee', size=3, opacity=0.5),
        name=name, showlegend=False,
    ), row=1, col=col)
    fig.add_trace(go.Scatter(
        x=[lo, hi], y=[lo, hi], mode='lines',
        line=dict(color='#f87171', dash='dash', width=1),
        name='perfect', showlegend=False,
    ), row=1, col=col)

fig.update_xaxes(title_text='Cheetah (mm)', showgrid=True, gridcolor='#1e3048')
fig.update_yaxes(title_text='Surrogate (mm)', showgrid=True, gridcolor='#1e3048')
fig.update_layout(
    height=420,
    paper_bgcolor='#0d1b2a', plot_bgcolor='#111f30',
    font=dict(color='#e2e8f0'),
)
fig.show()

5 · Speed comparison

N_BENCH = 1_000
bench_params = torch.FloatTensor(N_BENCH, 2).uniform_(K1_MIN, K1_MAX)

# Cheetah: serial calls (one at a time, as in a real control loop)
t0 = time.perf_counter()
for k1_qf, k1_qd in bench_params:
    segment.QF.k1 = k1_qf
    segment.QD.k1 = k1_qd
    _ = segment.track(beam_in)
t_cheetah = time.perf_counter() - t0

# Surrogate: batched inference
bench_n = normalise_X(bench_params)
t0 = time.perf_counter()
surrogate.eval()
with torch.no_grad():
    _ = denormalise_y(surrogate(bench_n))
t_surrogate = time.perf_counter() - t0

print(f'Cheetah   : {t_cheetah*1e3:7.1f} ms for {N_BENCH} evals  ({t_cheetah/N_BENCH*1e3:.3f} ms/call)')
print(f'Surrogate : {t_surrogate*1e3:7.1f} ms for {N_BENCH} evals  ({t_surrogate/N_BENCH*1e6:.1f} μs/call)')
print(f'Speedup   : {t_cheetah/t_surrogate:.0f}×')

# Bar chart
fig = go.Figure(go.Bar(
    x=['Cheetah (serial)', 'Surrogate (batched)'],
    y=[t_cheetah * 1e3, t_surrogate * 1e3],
    marker_color=['#f87171', '#22d3ee'],
    text=[f'{t_cheetah*1e3:.0f} ms', f'{t_surrogate*1e3:.1f} ms'],
    textposition='outside',
))
fig.update_layout(
    title=f'Wall time for {N_BENCH} evaluations',
    yaxis_title='Time (ms)',
    paper_bgcolor='#0d1b2a', plot_bgcolor='#111f30',
    font=dict(color='#e2e8f0'),
)
fig.show()

Cheetah   :   814.7 ms for 1000 evals  (0.815 ms/call)
Surrogate :     1.4 ms for 1000 evals  (1.4 μs/call)
Speedup   : 575×

6 · Surrogate-in-the-loop: BO with the surrogate instead of Cheetah

Now we replace Cheetah with our trained surrogate inside the Bayesian optimisation objective. The surrogate is already hundreds of times faster — enabling BO over thousands of evaluations in seconds.

Note on accuracy: the surrogate is trained on 5 000 uniform random samples. The tight-focus basin (σ < 2 mm) covers only ~0.3 % of the parameter space, so only ~16 training points land there. The BO may find a local minimum. In practice you would use active learning or targeted sampling near the optimum to overcome this — see the further-reading section.

surrogate.eval()

def objective_surrogate(params: list[float]) -> float:
    """Use the surrogate instead of Cheetah."""
    x = torch.tensor(params, dtype=torch.float32).unsqueeze(0)
    with torch.no_grad():
        y_pred = denormalise_y(surrogate(normalise_X(x)))
    sigma_x, sigma_y = y_pred[0, 0].item(), y_pred[0, 1].item()
    return (sigma_x + sigma_y) * 1e-3   # back to metres for consistency

space = [Real(K1_MIN, K1_MAX, name="k1_QF"), Real(K1_MIN, K1_MAX, name="k1_QD")]

t0 = time.perf_counter()
bo_surrogate = gp_minimize(
    func=objective_surrogate,
    dimensions=space,
    n_calls=60, n_initial_points=12, noise=1e-10, random_state=42,
)
t_bo_surrogate = time.perf_counter() - t0

# Validate: run the BO result through Cheetah (ground truth)
sigma_ground_truth = sum(run_cheetah(*bo_surrogate.x))

print(f'BO (surrogate): best σ = {bo_surrogate.fun*1e3:.2f} mm  (wall time: {t_bo_surrogate:.2f} s)')
print(f'Ground truth  : σ = {sigma_ground_truth:.2f} mm  (Cheetah validation)')
print(f'  k1_QF = {bo_surrogate.x[0]:.2f},  k1_QD = {bo_surrogate.x[1]:.2f} T/m²')

BO (surrogate): best σ = 3.92 mm  (wall time: 15.77 s)
Ground truth  : σ = 3.92 mm  (Cheetah validation)
  k1_QF = 9.56,  k1_QD = -15.00 T/m²

Result interpretation: The surrogate BO found σ_x + σ_y ≈ 3.9 mm — Cheetah confirms the prediction is accurate at that point. Compare with direct Cheetah BO in notebook 02, which finds ≈ 2 mm — the 2× gap is because the narrow focusing basin is undersampled in the uniform training set. The per-evaluation surrogate speedup (several hundred ×, see the benchmark above) makes the 60 inner model calls essentially free; the surrogate-BO loop is then dominated by GP overhead, not beam tracking.

---

Bonus · Ensemble uncertainty quantification

A single surrogate gives predictions but no uncertainty estimate. An ensemble of 5 surrogates trained with different random seeds gives a cheap uncertainty estimate: high variance = the model is extrapolating, be cautious.

N_ENSEMBLE = 5
ensemble = []

for seed in range(N_ENSEMBLE):
    torch.manual_seed(seed)
    m = BeamSurrogate()
    opt = torch.optim.Adam(m.parameters(), lr=3e-3)
    sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=200)
    dl  = DataLoader(TensorDataset(X_train_n, y_train_n), batch_size=256, shuffle=True)
    m.train()
    for _ in range(200):
        for X_b, y_b in dl:
            loss = nn.functional.mse_loss(m(X_b), y_b)
            opt.zero_grad(); loss.backward(); opt.step()
        sch.step()
    m.eval()
    ensemble.append(m)
    print(f'  Model {seed+1}/{N_ENSEMBLE} trained')

# Predict with ensemble
with torch.no_grad():
    preds = torch.stack([denormalise_y(m(X_test_n)) for m in ensemble], dim=0)  # (5, N, 2)

mean_pred = preds.mean(0)
std_pred  = preds.std(0)

# Coverage: fraction of true values within ±2σ ensemble
coverage = ((y_test - mean_pred).abs() < 2 * std_pred).float().mean(0)
print('\nEnsemble coverage at ±2σ:')
for i, name in enumerate(OUTPUT_NAMES):
    print(f'  {name}: {coverage[i].item():.2%}')

  Model 1/5 trained

  Model 2/5 trained

  Model 3/5 trained

  Model 4/5 trained

  Model 5/5 trained

Ensemble coverage at ±2σ:
  σ_x: 67.00%
  σ_y: 71.10%

UQ map — visualising uncertainty across the parameter space

The ensemble std gives a continuous measure of confidence across the full (k₁_QF, k₁_QD) space — not just on the test set.

Key insight: the surrogate is most uncertain where:

• Training data was sparse (random uniform sampling may leave gaps)
• The beam is near a strongly focusing or defocusing configuration
• Multiple ensemble members disagree on the beam dynamics in that region

This map can be used to decide where to collect additional Cheetah (or real-machine) data.

# 2-D uncertainty map: ensemble std across the full (k1_QF, k1_QD) parameter space
N_GRID = 50
k1_grid = np.linspace(K1_MIN, K1_MAX, N_GRID).astype(np.float32)
k1_grid_2d = np.array([[qf, qd] for qf in k1_grid for qd in k1_grid], dtype=np.float32)
grid_norm = normalise_X(torch.tensor(k1_grid_2d))

with torch.no_grad():
    preds_grid = torch.stack(
        [denormalise_y(m(grid_norm)) for m in ensemble], dim=0
    )  # (N_ENSEMBLE, N_GRID², 2)

std_grid = preds_grid.std(0).numpy().reshape(N_GRID, N_GRID, 2)  # (50, 50, 2)

fig = make_subplots(rows=1, cols=2,
    subplot_titles=['Ensemble std(σ_x)', 'Ensemble std(σ_y)'])

for col, (i, cbar_x) in enumerate([(0, 0.44), (1, 1.0)], 1):
    fig.add_trace(go.Heatmap(
        z=std_grid[:, :, i].T,
        x=k1_grid, y=k1_grid,
        colorscale='Viridis',
        colorbar=dict(title='Std (mm)', x=cbar_x, len=0.85),
    ), row=1, col=col)

fig.update_xaxes(title_text='k₁_QF (T/m²)')
fig.update_yaxes(title_text='k₁_QD (T/m²)')
fig.update_layout(
    height=420,
    title_text='Ensemble uncertainty map — regions of high variance = model is less confident',
    paper_bgcolor='#0d1b2a', plot_bgcolor='#111f30',
    font=dict(color='#e2e8f0'),
)
fig.show()

print('High-std regions = divergent ensemble predictions = treat surrogate output with caution there.')
print(f'In-domain mean ±2σ coverage (from previous cell): σ_x {coverage[0].item():.1%} / σ_y {coverage[1].item():.1%}')

High-std regions = divergent ensemble predictions = treat surrogate output with caution there.
In-domain mean ±2σ coverage (from previous cell): σ_x 67.0% / σ_y 71.1%

Monte-Carlo dropout — single-model uncertainty

The 5-model ensemble above is one way to estimate uncertainty. A cheaper single-model alternative is Monte-Carlo dropout (Gal & Ghahramani, 2016): leave dropout active at inference and average $T$ stochastic forward passes. This approximates Bayesian posterior sampling at $\sim1/M$ the training cost of an $M$-model ensemble.

# Monte-Carlo dropout: dropout left ON at inference approximates a Bayesian posterior
class MCDropoutSurrogate(nn.Module):
    def __init__(self, n_in=2, n_out=2, width=128, p=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_in, width), nn.Tanh(), nn.Dropout(p),
            nn.Linear(width, width), nn.Tanh(), nn.Dropout(p),
            nn.Linear(width, width), nn.Tanh(), nn.Dropout(p),
            nn.Linear(width, n_out),
        )
    def forward(self, x):
        return self.net(x)

torch.manual_seed(0)
mc_model = MCDropoutSurrogate(p=0.1)
opt = torch.optim.Adam(mc_model.parameters(), lr=3e-3)
sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=200)
dl  = DataLoader(TensorDataset(X_train_n, y_train_n), batch_size=256, shuffle=True)
mc_model.train()
for _ in range(200):
    for X_b, y_b in dl:
        loss = nn.functional.mse_loss(mc_model(X_b), y_b)
        opt.zero_grad(); loss.backward(); opt.step()
    sch.step()

# Predictive distribution: T stochastic passes with dropout ACTIVE
T = 50
mc_model.train()                      # keep dropout stochastic at inference
with torch.no_grad():
    mc_preds = torch.stack([denormalise_y(mc_model(X_test_n)) for _ in range(T)], dim=0)
mc_mean = mc_preds.mean(0)
mc_std  = mc_preds.std(0)

print(f'Monte-Carlo dropout (T={T} stochastic passes):')
print(f'{"Output":<6}  {"R2":>6}  {"MAE (mm)":>9}  {"mean s_MC (mm)":>15}')
print('-' * 42)
for i, name in enumerate(OUTPUT_NAMES):
    r2  = r2_score(y_test[:, i].numpy(), mc_mean[:, i].numpy())
    mae = (y_test[:, i] - mc_mean[:, i]).abs().mean().item()
    print(f'{name:<6}  {r2:>6.4f}  {mae:>9.4f}  {mc_std[:, i].mean().item():>15.4f}')

mc_cov = ((y_test - mc_mean).abs() < 2 * mc_std).float().mean(0)
print('\nMC-dropout coverage at +/-2 sigma:')
for i, name in enumerate(OUTPUT_NAMES):
    print(f'  {name}: {mc_cov[i].item():.2%}')

print('\nSingle-model UQ at ~1/N_ENSEMBLE the cost of the deep ensemble; both flag\n'
      'the high-uncertainty corners of the (k1_QF, k1_QD) domain.')

Monte-Carlo dropout (T=50 stochastic passes):
Output      R2   MAE (mm)   mean s_MC (mm)
------------------------------------------
σ_x     0.9993     0.0221           0.1284
σ_y     0.9993     0.0220           0.1265

MC-dropout coverage at +/-2 sigma:
  σ_x: 99.80%
  σ_y: 99.60%

Single-model UQ at ~1/N_ENSEMBLE the cost of the deep ensemble; both flag
the high-uncertainty corners of the (k1_QF, k1_QD) domain.

Explainability — feature importance via gradient attribution

Feature importance answers: "which inputs does this surrogate depend on most?"

We compute the mean absolute gradient |∂σ_x/∂k₁| and |∂σ_y/∂k₁| across 200 test samples.

This is the physics sanity check for the surrogate. In a FODO doublet every quadrupole acts in both planes (focusing in one, defocusing in the other), so σ_x and σ_y should each depend on both k₁_QF and k₁_QD — a roughly uniform sensitivity matrix, not a diagonal one. A near-zero block would instead suggest the model missed the dual-plane optics or latched onto a spurious correlation.

# Feature importance: mean |∂σ_x/∂k1| and |∂σ_y/∂k1| across 200 test samples
surrogate.eval()
importance = np.zeros((2, 2))   # (n_outputs, n_inputs)

for x_row in X_test_n[:200]:
    x = x_row.unsqueeze(0).clone().requires_grad_(True)
    y_pred_local = surrogate(x)   # (1, 2)
    for out_i in range(2):
        y_pred_local[0, out_i].backward(retain_graph=(out_i == 0))
        importance[out_i] += x.grad.abs().squeeze(0).detach().numpy()
        if x.grad is not None:
            x.grad.zero_()

importance /= 200.0

fig = go.Figure(go.Heatmap(
    z=importance,
    x=['k₁_QF', 'k₁_QD'],
    y=['σ_x', 'σ_y'],
    colorscale='Blues',
    text=[[f'{v:.4f}' for v in row] for row in importance],
    texttemplate='%{text}',
    showscale=True,
    colorbar=dict(title='Mean |∂output/∂input|'),
))
fig.update_layout(
    title='Feature importance — gradient attribution (mean |∂output/∂input|)',
    paper_bgcolor='#0d1b2a', plot_bgcolor='#111f30',
    font=dict(color='#e2e8f0'),
    height=300,
)
fig.show()

print('Physics check:')
print('  Each quadrupole acts in BOTH planes (focus in one, defocus in the other),')
print('  so σ_x and σ_y should both depend on both k₁_QF and k₁_QD — a near-uniform matrix.')
print('  A near-zero block would signal missed optics or a spurious correlation.')

Physics check:
  Each quadrupole acts in BOTH planes (focus in one, defocus in the other),
  so σ_x and σ_y should both depend on both k₁_QF and k₁_QD — a near-uniform matrix.
  A near-zero block would signal missed optics or a spurious correlation.

---

Summary

	Cheetah	Surrogate	Speedup
Per-call time	~0.8 ms	~0.001 ms	~600×
Differentiable	✓	✓	—
Works offline	✓	✓ (after training)	—
Works on real machine	✗	✓ (if trained on real data)	—
Uncertainty	✗	✓ (ensemble)	—

The surrogate trained on Cheetah is useful for prototyping. In production:

• Train on real machine data (or on more expensive codes like TraceWin / Geant4)
• Update the surrogate online as the machine drifts
• Use the ensemble std to flag when you're extrapolating outside the training domain

Further reading:

• Kaiser et al., PRAB 27, 054601 (2024) — Cheetah
• Edelen et al., NeurIPS ML4PS (2020) — surrogates in accelerator control
• AccML living review: aghribi.github.io/acc-ml-living-review