TomOpt: Differentiable Design for Muon Tomography

• TomOpt is a modular Python package for muon scattering tomography that integrates cosmic-muon transport, detector response, reconstruction, and task inference into a single differentiable pipeline.
• It employs end-to-end differentiable optimization, using gradient back-propagation to update detector design parameters against realistic task-aware loss functions under cost constraints.
• Empirical studies demonstrate that TomOpt can enhance detector performance and inference accuracy in applications such as industrial muography and cargo border control by fine-tuning geometry and budget allocation.

Searching arXiv for papers on TomOpt and related applications. TomOpt is a modular Python software package for optimizing the design of muon-scattering tomography detectors by treating the whole chain—from cosmic-muon transport, to detector response, to reconstruction and inference, to detector-parameter updates—as a differentiable computation graph. It was developed as a proof of principle that particle-physics instruments can be designed not just by hand-tuning geometric proxies such as track resolution or acceptance, but by directly optimizing for the actual end task the instrument must solve, under realistic constraints such as cost. In the original formulation, the application context is industrial muography, specifically estimating the fill level of a steel ladle in a metal refinery using naturally occurring cosmic-ray muons; a later study uses TomOpt as a machine-learning-driven, differentiable optimization engine for a cargo border control scanner (Strong et al., 2023, Zaher et al., 14 Jul 2025).

1. Conceptual scope and problem formulation

TomOpt targets muon scattering tomography (MST), in which cosmic-ray muons traverse dense or inaccessible objects and undergo multiple Coulomb scattering. The amount of angular deflection depends on the material properties, especially through the radiation length $X_0$ and therefore indirectly on atomic number and density. MST instruments place tracking detectors above and below a passive volume, reconstruct the incoming and outgoing muon trajectories, and infer internal structure from the observed scattering. TomOpt addresses the instrument-design problem for this setting: how detector geometry, panel spacing, active area, placement, and related specifications should be chosen so that the final tomography task—imaging, material discrimination, or a task-specific regression or classification objective—is solved as well as possible (Strong et al., 2023).

The central methodological claim is end-to-end differentiable and inference-aware optimization. Rather than optimizing against detached proxies, TomOpt builds differentiable models for muon generation and propagation, detector hit generation, trajectory fitting, scattering-feature extraction, inference of latent volume properties or task-level predictions, and loss evaluation. Because these stages form a differentiable computation graph, gradients of the final task loss can be back-propagated to detector-design variables. This makes the optimization “inference-aware”: detector updates occur in the presence of the actual reconstruction and task objective, rather than in a proxy-only loop (Strong et al., 2023).

A later application in cargo inspection preserves the same abstraction. There TomOpt is presented as one of two complementary optimization approaches, alongside a more detailed but slower GEANT4 study. Its role is again to formalize detector design as an optimization problem tied to downstream reconstruction fidelity, with detector geometry decisions connected directly to reconstruction quality and material-discrimination performance through a differentiable pipeline (Zaher et al., 14 Jul 2025).

2. Differentiable simulation pipeline

The muon-transport model begins from parameterized cosmic-muon flux models. TomOpt lets the user sample muons from either the Guan et al. or Shukla et al. surface-flux parameterizations. The passive volume is voxelised in 3D, and each voxel is assigned a material via its radiation length $X_0$. Muons are propagated in steps of length $dr$, and the default scattering model is the Particle Data Group small-angle Gaussian approximation: $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$ TomOpt samples the scattering angle in each transverse direction as

$\delta \theta = \theta_0 z_2,$

with $z_2 \sim \mathcal N(0,1)$, and the correlated lateral displacement as

$\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$

with $z_1 \sim \mathcal N(0,1)$. The same is done in two orthogonal directions in the muon frame, then rotated into the volume frame. This is an approximation chosen for speed and differentiability. The authors explicitly note that the PDG Gaussian model neglects hard-scattering tails and, when used in a stepwise propagation algorithm, underestimates scattering relative to Geant4-like transport (Strong et al., 2023).

Detector modeling is likewise differentiable. The basic unit is a detector panel with learnable position and span. Each panel has nominal spatial resolution $\sigma_{\max}$, nominal efficiency $\epsilon_{\max}$, and a cost per unit area. In the current implementation, five parameters per panel are learned: $X_0$0 position and $X_0$1 span. A hard-edged physical panel would create non-smooth acceptance changes and vanishing gradients for off-panel muons, so TomOpt uses a smooth surrogate detector model during optimization. In “optimisation mode,” panels record hits for all muons at that panel’s $X_0$2, but hit efficiency and hit resolution smoothly degrade with distance from the panel centre according to a pair of double-sigmoid functions in $X_0$3 and $X_0$4. The hit efficiency is modeled as

$X_0$5

This smooth extension is the key technical device enabling differentiability with respect to panel centre, area, and position (Strong et al., 2023).

Hits are passed to differentiable trajectory fitting. Incoming and outgoing tracks are fit analytically from the $X_0$6 hit positions, weighted by hit uncertainties and efficiencies. TomOpt then computes derived scattering variables such as incoming and outgoing angles and total angular deflection, and propagates uncertainties by automatic differentiation through the fitted track formulas. A first-stage scattering localization is performed using an extended Point-of-Closest-Approach (PoCA) algorithm, in which the scattering is approximated as if it happened at one point where the incoming and outgoing fitted trajectories come closest. PoCA coordinates and their uncertainties are differentiable functions of the measured hits and panel parameters (Strong et al., 2023).

In the cargo-scanning formulation, the same differentiable structure is emphasized. There the pipeline includes muon generation, propagation through detector planes and passive volume, a differentiable detector response model, track fitting, PoCA-based reconstruction, volume inference, loss evaluation, and backpropagation of gradients to detector design parameters. Detector hits are reconstructed by sampling $X_0$7 and $X_0$8 uncertainties from Gaussians centered on the true position with widths given by the local differentiable spatial resolution. The detector efficiency and position resolution are again modeled by a double-sigmoid surrogate function (Zaher et al., 14 Jul 2025).

3. Inference layer and task-aware loss design

TomOpt is modular at the volume-inference stage. The original paper presents a generic voxel-level radiation-length inference by inverting the scattering model. For muons whose PoCA lies in a voxel of height $X_0$9, the authors use

$dr$0

and then invert the PDG relation, ignoring the logarithmic correction term, to estimate $dr$1: $dr$2 which gives

$dr$3

In practice, the aggregation can be weighted by muon efficiency and PoCA-voxel probability. To account for PoCA uncertainty, the nominal PoCA is smeared with Gaussian PDFs in $dr$4, integrated over neighbouring voxels, and each muon contributes probabilistically to multiple voxels. This softens voxel assignment while remaining differentiable (Strong et al., 2023).

The authors stress that this PoCA-based $dr$5 inference is only a first-stage approximation and is biased because it assumes one effective scattering point. TomOpt can also support more advanced inference such as expectation-maximisation or deep-learning methods. This flexibility is part of the framework definition: the purpose is not tied to one reconstruction algorithm, but to the broader principle that whatever inference sits between measured hits and final task output should participate in the differentiable optimization (Strong et al., 2023).

The benchmark task in the original paper is fill-level estimation in a steel ladle, framed as a single-variable regression problem. For this task, the authors do not use full $dr$6 reconstruction. Instead they define a weighted PoCA average because the mean PoCA $dr$7-coordinate is monotonic with fill height. Each PoCA receives weight

$dr$8

where $dr$9 is the trajectory efficiency, $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$0 the PoCA $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$1-uncertainty, and $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$2 a transverse weighting factor suppressing PoCAs associated with ladle walls and emphasizing those near the central interior. The predicted vertical observable is

$\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$3

Because this is biased with respect to the true fill height, a linear de-biasing transform is applied,

$\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$4

where $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$5 and $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$6 are fitted from known examples and updated during optimisation (Strong et al., 2023).

The task loss is deliberately task-aware. Rather than optimize mean squared error directly, the original paper uses a “spread over range” objective: $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$7 where $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$8 is the number of fill levels, $\theta_0 = \frac{0.0136}{p}\sqrt{n_{X_0}\left(1 + 0.038\ln n_{X_0}\right), \qquad n_{X_0} = \frac{d}{X_0}.$9 is the standard deviation of predictions at fixed true fill height, and $\delta \theta = \theta_0 z_2,$0 is the standard deviation of the mean prediction across fill heights. The numerator rewards small within-class spread; the denominator rewards large separation between classes. This encodes a measurement criterion in which repeated scans of the same fill level should cluster tightly and different fill levels should remain separable (Strong et al., 2023).

The cargo-inspection study uses two different task-driven objective functions within TomOpt. The first is an $\delta \theta = \theta_0 z_2,$1-inference-based objective using the mean squared error between inferred and true voxel radiation lengths, evaluated on normalized inferred and ground-truth $\delta \theta = \theta_0 z_2,$2 values because the estimator is biased, especially for low-$\delta \theta = \theta_0 z_2,$3 material like air. The second is a BCA-based objective, where the loss is the MSE between inferred BCA voxel scores and the logarithmic ground-truth $\delta \theta = \theta_0 z_2,$4, both normalized. In that work, the optimization target is explicitly reconstruction fidelity of the inferred 3D cargo image rather than a purely geometric criterion (Zaher et al., 14 Jul 2025).

4. Optimization variables, constraints, and training mechanics

TomOpt’s optimization formulation is task-aware and constraint-aware. There is no universal loss: classification may use cross-entropy, regression may use MSE or another statistic, and the total objective includes detector constraints such as budget. The original paper describes two budget treatments. In budget-penalisation mode, a differentiable cost term is added that remains small below a target budget and rises rapidly above it. In fixed-budget mode, the total cost is held exactly constant, and the optimiser instead learns a fractional budget allocation per panel. Panel areas are rescaled so that each panel’s area-cost matches its learned share of the total cost. This removes the need to tune a performance-versus-cost coefficient, because the loss becomes purely the task loss over a budget-constrained parameterisation (Strong et al., 2023).

The optimised variables can therefore include panel $\delta \theta = \theta_0 z_2,$5 placement, $\delta \theta = \theta_0 z_2,$6 span, and budget assignment. In the present implementation the design space is limited to horizontal panels with continuous resolution and efficiency models. The number of panels per layer is fixed because integer-valued architectural choices are not naturally differentiable (Strong et al., 2023).

Parameter updates use standard gradient-descent optimisers from PyTorch, such as SGD or Adam. The paper describes the optimization as the usual backpropagation-based minimisation of detector parameters $\delta \theta = \theta_0 z_2,$7,

$\delta \theta = \theta_0 z_2,$8

where $\delta \theta = \theta_0 z_2,$9 denotes the learnable detector parameters and $z_2 \sim \mathcal N(0,1)$0 a learning rate or adaptive equivalent. The package also includes practical stabilization machinery: one-cycle learning-rate scheduling, warm-up estimation of physically meaningful per-parameter learning rates, NaN-gradient suppression, panel-centering constraints, and a schedule that gradually sharpens the sigmoid panel model so that optimisation starts smooth and ends closer to physical detector behaviour (Strong et al., 2023).

The cargo-border-control study retains gradient descent as the core update mechanism,

$z_2 \sim \mathcal N(0,1)$1

but adds a Bayesian Optimization extension because some reconstruction-driven objectives are too noisy for plain gradient descent to behave reliably. The BO framework is implemented using Ax, which internally uses BoTorch. A Gaussian Process is used as the surrogate model and the acquisition function is Expected Improvement. In that study, BO replaces only the final optimizer; the simulation, reconstruction, and objective evaluation remain those of TomOpt (Zaher et al., 14 Jul 2025).

The reported search strategy is hybrid. The first 12 trials are generated using a Sobol quasi-random sequence to cover the search space uniformly and bootstrap the GP, after which Ax switches to a GP+EI strategy. Elsewhere in the plots the first 24 trials are referred to as Sobol-sampled; the paper does not reconcile this discrepancy, so it should be regarded as an inconsistency in reporting. BO is run for 100 evaluation trials, with the search space consisting of detector $z_2 \sim \mathcal N(0,1)$2 positions (Zaher et al., 14 Jul 2025).

5. Benchmark studies and empirical behavior

The original TomOpt benchmark is fill-level estimation in a steel ladle at a metal refinery. The passive volume is a rectangular ladle containing solid steel walls, liquid steel, a slag layer, and air. The detector consists of four square panels above and four below the ladle. The benchmark uses eight fill heights, 10 cm voxelisation, and 1000 muons per evaluation. The initial detector is intentionally poor: panels are offset in $z_2 \sim \mathcal N(0,1)$3-$z_2 \sim \mathcal N(0,1)$4 away from the ladle centre and packed too tightly in $z_2 \sim \mathcal N(0,1)$5, which degrades useful acceptance and PoCA resolution (Strong et al., 2023).

Empirically, TomOpt moves the panels toward the ladle centre in $z_2 \sim \mathcal N(0,1)$6, increases their separation in $z_2 \sim \mathcal N(0,1)$7, and adjusts budget allocation across panels. Stage one optimisation improves the mean MSE after de-biasing from $z_2 \sim \mathcal N(0,1)$8 to $z_2 \sim \mathcal N(0,1)$9, and the average absolute error from 16 cm to 2.8 cm. A second refinement stage, with panel $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$0 positions fixed at the intuitive central location and only $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$1 positions plus budget assignment trainable, further improves performance to $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$2 mean MSE and 2.5 cm average absolute error. The final TomOpt-designed detector slightly but consistently outperforms two sensible human-designed baselines: one maximizing long-baseline tracking precision and one using more even panel spacing to balance precision and angular acceptance (Strong et al., 2023).

These results are presented as an illustration of competing design tradeoffs. Wider $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$3-separation improves track and PoCA resolution, but too much separation reduces acceptance for high-angle muons. Larger panel area improves useful flux, but panels extending beyond the relevant target footprint admit more muons influenced by walls and background structure. A fixed total budget means area added to one panel must be taken from another. TomOpt balances these effects because all are visible to the final loss through the differentiable chain (Strong et al., 2023).

The cargo-border-control study uses TomOpt on a scanner with a pair of hodoscopes above and below a passive volume, with each hodoscope measuring $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$4. Initially they are offset by $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$5 in both $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$6 and $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$7, so the passive volume is only partially covered. The passive volume is $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$8, representing a lorry-like cargo region with an iron base and random $\delta x = \theta_0\left(z_1\frac{d}{\sqrt{12} + z_2\frac{d}{2}\right),$9 blocks of iron, beryllium, and air, plus a $z_1 \sim \mathcal N(0,1)$0 uranium block inserted with probability 0.5 (Zaher et al., 14 Jul 2025).

With the BCA objective, the best observed BO objective decreases by 13.3% over 100 trials. The final optimized configuration combines more complete upper coverage with lower overlap. With the $z_1 \sim \mathcal N(0,1)$1-MSE objective, the best observed objective decreases by 22.8% over 100 trials, but the authors interpret the optimization as less reliable because the objective is much noisier. Evaluated over 100 simulated passive volumes—50 without uranium and 50 with an additional randomly placed uranium block—the optimized configuration yields only a slight improvement in MSE for the $z_1 \sim \mathcal N(0,1)$2 method, with SSIM remaining similar across initial, optimized, and baseline configurations. For the BCA method, the optimized design gives consistent improvements in both MSE and SSIM and reaches performance comparable to the baseline detector. A plausible implication is that TomOpt’s effectiveness is strongly contingent on the stability and informativeness of the reconstruction objective supplied to the optimization loop (Zaher et al., 14 Jul 2025).

6. Software architecture, novelty claim, and limitations

TomOpt is Python-based, highly modular, and implemented on top of PyTorch for automatic differentiation. The package organizes the workflow into interchangeable modules for muon generation, passive-volume modelling, detector panels and hit recording, tracking and PoCA, inference, loss definition, and optimisation callbacks. The callback system supports muon resampling, de-biasing correction, population-level loss computation, NaN protection, physically calibrated learning-rate setup, panel-centering regularisation, metric logging, checkpointing, one-cycle learning-rate scheduling, and sigmoid-smoothness scheduling (Strong et al., 2023).

The novelty claim is specific: TomOpt is presented as the first demonstration of end-to-end-differentiable and inference-aware optimization of particle-physics instruments. The contrast is with a traditional workflow in which one proposes a geometry, simulates it, reconstructs events, evaluates a few proxy metrics, and manually iterates or uses black-box search. In that arrangement, reconstruction and task inference are often frozen, simplified, or absent from the loop, which can misalign the optimization with the actual scientific objective. TomOpt differs by embedding inference in the loop and propagating gradients from the end metric back to physical design parameters (Strong et al., 2023).

The package has clear limitations. The transport model is approximate: pure Gaussian scattering, no hard tails, and no muon energy loss or decay, which overestimates detected flux for low-energy muons. The step size $z_1 \sim \mathcal N(0,1)$3 affects scattering fidelity, and the authors recommend choosing it around 1% of the passive-volume size. Detector modelling is simplified: only horizontal panels are supported, with continuous resolution and uniform granularity; realistic signal formation and channelisation are not modeled. Current cost modeling is basic, and the presented volume inference assumes perfect knowledge of muon momentum, which is unrealistic in most MST systems. Computationally, although the steel-ladle benchmark runs on a CPU laptop in tens of minutes, larger detectors and richer inference models will be more demanding (Strong et al., 2023).

The cargo study adds further practical limitations. The strongest one is the noise and bias in the reconstruction objectives, especially the $z_1 \sim \mathcal N(0,1)$4-based one, whose inferred radiation lengths are systematically biased low, particularly for air. Another limitation is the then-current hodoscope implementation: plates are treated as unique plates, so placing two plates at the same $z_1 \sim \mathcal N(0,1)$5 causes problems for track fitting, forcing the authors to stagger hodoscopes in $z_1 \sim \mathcal N(0,1)$6 to avoid overlap. The authors also note that the BO-optimized layouts did not significantly surpass a human-designed baseline, indicating that current inference methods may not be sensitive enough to subtle parameter changes (Zaher et al., 14 Jul 2025).

7. Applications and projected development

TomOpt’s demonstrated applications are industrial muography and cargo border control. In the steel-ladle study, the downstream task is fill-level regression. In the SilentBorder-related cargo study, the downstream task is 3D cargo-image reconstruction and material discrimination, evaluated through MSE and SSIM under $z_1 \sim \mathcal N(0,1)$7- and BCA-based inference objectives. In both cases, detector placement and coverage emerge as major design levers, but the preferred geometry depends on the task definition and the reconstruction model (Strong et al., 2023, Zaher et al., 14 Jul 2025).

The original paper identifies several future directions: open-source release; stronger and more realistic inference, including deep neural networks; richer detector geometry such as side detectors and rotated planes; better transport modeling, potentially via differentiable parametric surrogates of Geant; better technology and cost models; and a broader suite of benchmark tasks. It also notes that, beyond differentiable optimization, TomOpt can serve as a fast simulator for developing advanced non-differentiable inference algorithms for muon tomography (Strong et al., 2023).

The cargo study reinforces these directions. TomOpt is positioned as the optimization-centric counterpart to a higher-fidelity GEANT4 framework: TomOpt is faster and differentiable, enabling optimization workflows that would be impractical with raw GEANT4, while GEANT4 remains the tool for detailed validation and broader detector-physics studies. The authors propose future replacement or augmentation of current reconstruction losses with deep-learning-based volume inference, either as post-processing on reconstruction outputs or as fully end-to-end learned inference directly from minimally processed scattering information (Zaher et al., 14 Jul 2025).

Taken together, these works define TomOpt as a differentiable design-optimization framework for muon-scattering tomography detectors whose central contribution is not a particular detector geometry or a particular reconstruction algorithm, but a methodological shift: detector design is optimized against the actual downstream tomography task, with reconstruction and inference retained inside the optimization loop under explicit geometry and budget constraints (Strong et al., 2023).