A comprehensive framework toward the seamless integration of muon reconstruction algorithms with machine learning

Muon-scattering tomography (MST) utilizes naturally occurring cosmic-ray muons to reveal the three-dimensional composition of concealed volumes, such as cargo containers in the maritime domain, reducing the need for artificial radiation sources. The reconstruction methods of current state-of-the-art systems rely on geometry-based approaches, such as the Point of Closest Approach (PoCA) algorithm, whose strong heuristics blur fine structures and introduce high-frequency noise. Statistical Expectation–Maximization (EM) reconstruction methods can recover these lost details but are traditionally ruled out for real-time application given their high computational and numerical demands. We introduce a comprehensive framework for MST reconstruction in PyTorch, including traditional and fast, but inaccurate geometry-based methods, as well as a highly optimized EM solver within a single, end-to-end differentiable pipeline. Using parallelism and graphics processing unit (GPU) acceleration, our framework overcomes the aforementioned computational obstacles. As a benchmark, the EM solver is tested on several MST scenarios generated with Geant4. Image quality metrics shows its superiority over traditional reconstruction algorithms, while retaining a per-iteration latency of 0.8s at a 1cm voxel resolution on standard GPUs.