Scattering-Based Machine Learning Algorithms for Momentum Estimation in Muon Tomography

Abstract

Muon tomography leverages the small, continuous flux of cosmic rays produced in the upper atmosphere to measure the density of unknown volumes. The multiple Coulomb scattering that muons undergo when passing through the material can either be leveraged or represent a measurement nuisance. In either case, the scattering dependence on muon momentum is a significant source of imprecision. This can be alleviated by including dedicated momentum measurement devices in the experiment, which have a potential cost and can interfere with measurement. An alternative consists of leveraging information on scattering withstood through a known medium. We present a comprehensive study of diverse machine-learning algorithms for this regression task, from classical feature engineering with a fully connected network to more advanced architectures such as recurrent and graph neural networks and transformers. Several real-life requirements are considered, such as the inclusion of hit reconstruction efficiency and resolution and the need for a momentum resolution prediction that can improve reconstruction methods.

1. Introduction

Muon tomography encompasses imaging methods based on the measurement of absorption or deflection of atmospheric muons. The two main electromagnetic phenomena muons are subjected to, absorption by energy loss and multiple Coulomb scattering through matter, can be separately exploited by the two types of muon imaging: muon transmission and muon scattering tomography. In the past decade, applications of muon tomography have expanded across a wide range of fields, including geosciences [1], nuclear safety [2], archaeology [3], civil engineering [4], and security applications [5], and the technology has been gradually moving from high-energy physics to concrete engineering applications. While significant advancements have been made in detector design and reconstruction algorithms, current muon detection systems still lack the ability to measure muon momentum and rely solely on track measurements for imaging tasks.

This paper is organized as follows: in Section 2, we introduce the origin and flux of muons, how their momentum impacts both scattering and transmission tomography, and how it can be measured through scattering in known material; in Section 3, we discuss multiple machine learning (ML) algorithms trained to exploit the detector hit information to regress the momentum and improve its reconstruction accuracy; finally, in Section 4, we study a few requirements needed to apply these models to the constraints of a real-life experiment.

2. Muon Scattering and Momentum

2.1. Muon Flux at Sea Level

Both types of muon scattering and transmission muon tomography take advantage of the atmospheric muon flux. Atmospheric muons are abundantly produced by the interaction of primary cosmic rays within the upper atmosphere and represent approximately 80% of all charged particles reaching sea level, with a flux through a horizontal surface of around one particle per cm² per minute. An example of the cosmic ray muon flux is given by the Gaisser parametrization [6]

$${\displaystyle \frac{dN}{dEd\Omega }}=0.14E^{-2.7}\left({\displaystyle \frac{1}{1+\frac{1.1Ecos\theta }{115 \mathrm{GeV}}}}-{\displaystyle \frac{0.054}{1+\frac{1.1Ecos\theta }{850 \mathrm{GeV}}}}\right){\mathrm{cm}}^2{\mathrm{s}}^{-1}{\mathrm{sr}}^{-1}{\mathrm{GeV}}^{-1}$$

(1)

At sea level, the muon energy spectrum through a horizontal surface peaks at 4 GeV, as shown in Figure 1 in a wide momentum interval from hundreds of MeV up to thousands of GeV The width of the momentum spectrum translates to different behavior in the interaction with matter. While high-momentum muons (100 GeV and above) are useless to muon scattering tomography as they traverse matter without undergoing measurable deflections, they are crucial to muon transmission tomography, as they can traverse large thicknesses of material and allow for probing kilometre-scale structures. The effect of the momentum on muon scattering and transmission tomography is further detailed in Section 2.2 and Section 2.3, respectively.

2.2. Muon Momentum in Muon Scattering Tomography

Like other charged particles, when going through matter, muons collide with electrons and nuclei, thereby losing energy and undergoing deflections. This process is called multiple Coulomb scattering and can be approximated by a Gaussian distribution. The differential distribution of scattered muons can therefore be written as

$${\displaystyle \frac{dN}{d\theta }}={\displaystyle \frac{\theta }{\sqrt{2\pi }{\theta }_0}}{e}^{-{\theta }^2/2{\theta }_0^2},$$

(2)

where the standard deviation ${\theta }_0$ is given by Lynch and Dahl [8]:

$${\theta }_0={\displaystyle \frac{13.6\mathrm{MeV}}{p\beta c}}\sqrt{{\displaystyle \frac{x}{X_0}}}\left[1+0.038\ln \left({\displaystyle \frac{x}{X_0{\beta }^2}}\right)\right].$$

(3)

p and $\beta c$ are the muon’s momentum and speed, and $x/X_0$ is the thickness of the scattering medium in radiation length. Radiation length is defined as the mean length of the material at which the energy of an electron is reduced by a factor $1/e$, and it can be found in the PDG [9]. This equation only holds under the small-angle scatter assumption, as hard scatterings on nuclei can cause non-Gaussian tails.

The scattering angle is thus a function of both the material properties ($X_0$) and muon momentum (p). In the context of a muon scattering tomography experiment, individual deflections of high-momentum muons through dense material and low-momentum muons through light material are indistinguishable. Since image reconstruction and material inference algorithms used in muon scattering tomography utilize measured scattering angles as input, they suffer from a lack of knowledge on muon momentum.

Some of these algorithms, like Angle Statistic Reconstruction [10], Binned Clustered Algorithm (BCA) [11], or momentum-aware POCA [12,13,14], can utilize momentum knowledge if available. To illustrate the impact of including the true momentum—as used in the simulations—in the density reconstruction, an example of a scattering muon tomography setup simulation is shown in Figure 2, with the predictions of the BCA with the true momentum—as used in simulation—and without any momentum information. Their performances are quantified through the mean squared error $\mathrm{MSE}={\displaystyle \frac{1}{N}}{\sum }_{i=0}^N{(Y-\widehat{Y})}^2$ and mean absolute error $\mathrm{MAE}={\displaystyle \frac{1}{N}}{\sum }_{i=0}^N|Y-\widehat{Y}|$ of the normalized density prediction $\widehat{Y}$. The figure shows that including true momentum information during density inference not only reduces noise but also improves the quality of the inference. While the improvement depends on the algorithm, the integration of a momentum measurement system into any muon scattering tomography experiment is extremely relevant.

2.3. Muon Momentum in Muon Transmission Tomography

Muon transmission tomography produces radiographic images by measuring the absorption of the cosmic muon flux through an object. Several experiments in the geoscience sector intend to infer the in-depth kilometre-scale volcano structures like the Puy-de-Dome [16] in France, the Mount Vesuvius in Italy [17], or Soufrière in Guadeloupe [18]. Given the large thicknesses of rock that muons have to traverse before reaching the detection system, only high-momentum muons (roughly above 100 GeV) can traverse the object without decaying and thus constitute most of the signal. On the other hand, low-momentum muons (5 GeV and below), which represent around 90% of the total muon flux, can scatter and enter the detector acceptance without interacting with the region of interest, thus mimicking high-momentum muons and biasing the muon flux absorption measurement, as shown in Figure 3. Since the flux of high-zenith-angle muons (near-horizontally) decreases by 2 to 4 orders of magnitude after a few kilometres of rock, it becomes comparable with the flux of the non-penetrating muons and makes the rejection of such events determinant.

While some experiments use a lead block to shield their detectors from low-momentum muons [17,18], reject them from the goodness of the track fit [19], or rely on coarse time of flight measurements to reject background muon tracks [20], the inclusion of a muon momentum measurement module in the detection system could further improve the background rejection capabilities.

2.4. Muon Momentum Measurement

Many particle physics experiments rely on the measurement of charged particle momentum. While momentum inference from the particle trajectory in a strong magnetic field is used extensively in high-energy physics experiments at the LHC, it is not suitable for muon tomography. Generating a high magnetic field is extremely costly as it requires large and complex magnets powered by high-power sources, which also imposes strong logistical constraints. It also interferes with the image reconstruction algorithms, as it becomes harder to disentangle muon scattering from the deflections caused by the magnetic field.

Two main techniques have been investigated in the literature, some exploiting Cherenkov radiation of muons through dedicated detectors [21,22] and others exploiting multiple Coulomb scattering of muons through passive layers of high-density and high-Z materials (e.g., lead, iron) [19,23,24,25]. The work presented in this paper focuses on the latter technique.

Figure 4a gives an example of a momentum measurement module built with three dense material layers (e.g., lead) of thickness $\Delta s$ used to scatter the muons. The muon hits are measured by eight detector layers separated by a gap $\Delta g$. The simulated detector panels are non-physical infinitely thin detectors recording the exact position of the muon. As such, muon scattering only occurs when propagating through the lead panels. This is a reasonable approximation as detector panels are typically thinner and have a lower density—e.g., gas mixtures, scintillators—than the scattering material. The selection of tracks is limited to those leaving a hit in each of the eight detector panels. From those eight hits, four linear segments can be determined from the hits, and three scattering angles $\Delta {\theta }_i$ can be reconstructed. Contrary to the typical muon scattering tomography task, where one wants to infer the density from the measured scattering angle, momentum can be estimated by inverting Equation (3):

$$p={\displaystyle \frac{13.6\mathrm{MeV}}{{\theta }_0}}\sqrt{{\displaystyle \frac{x}{X_0}}}\left[1+0.038\ln \left({\displaystyle \frac{x}{X_0}}\right)\right],$$

(4)

where muons are assumed to be relativistic ($\beta \simeq 1$). In this context, the scattering material $X_0$ is pre-defined, and the distance traversed by the muon x is estimated from the track measurement. Measuring the scattering angle of the muons through N layers of material gives a rough estimate of the scattering angles’ standard deviation ${\theta }_0$:

$${\theta }_0={\displaystyle \frac{1}{\sqrt{2}}}{\theta }_{RMS}^{scatt}={\displaystyle \frac{1}{\sqrt{2}}}\sqrt{{\displaystyle \sum _{i=1}^N}{\left(d{\theta }_i^{scatt}\right)}^2}.$$

(5)

Together with Equation (4), the scattering angles’ standard deviation ${\theta }_0$ provides an estimation of the muon momentum, as illustrated in Figure 4b. The mismeasurement tail at low momentum originates from the fact that the muon energy loss through 4 cm of lead has the same order of magnitude as its momentum, for example, a 200 MeV muon loses about 55 MeV through 4 cm of lead. This could be circumvented by the use of Kalman filters, as demonstrated in Ref. [26].

While such an approach has satisfactory results, it remains suboptimal as only a fraction of the information accumulated by the detector is used. Only scattering angle information can be plugged in Equation (3), representing three measurable parameters, while the detectors measure the muon position at eight locations.

In Section 3, different machine learning algorithms are trained to improve the momentum prediction accuracy. The dataset used during the training follows the setup of Figure 4a. The CRY library [7] is used to generate the muon flux, and the muon propagation and scattering are handled by GEANT4 [15].

3. Machine Learning Methods

Although physically motivated, using the sole scattering angles in Equation (4) represents a limited use of the information available from the detector. At the lowest level, all the information that can be gathered is contained within the hits registered in the different active layers. The ambition of this section is to explore different ML algorithms and compare their performance.

Historically, the first ML algorithms used in particle physics were boosted decision trees (BDTs) and neural networks (NNs). These simple models could be easily trained with the limited training data of the time, but they required the physicist to carefully craft variables to optimize the discrimination power. Such variables were often physically motivated, for example, an invariant mass or specific angular distance. In a way, this is similar to the use of the scattering angles motivated by Equation (4). Then, riding on the deep learning revolution and helped by the ever-growing availability of high-fidelity simulations, physicists figured out that they could increase their discrimination performance by supplementing these manually crafted variables with a larger set of lower-level information and providing this set to a deep neural network (DNN). This was especially visible in the context of jet tagging in LHC experiments [27], a topic that was and still is a main driver of ML development. Including additional information from each particle identified within a jet provided significant performance improvement and opened the door to a wide range of measurements thought to be impractical until then.

Arguably, a second revolution occurred when physicists started leveraging the symmetries of the problem they were facing directly within the ML model architecture. One such model was based on convolutional neural networks (CNNs), which were initially developed in the context of image processing and leverage local information through a kernel of weights applied across the image. It represents a jet as an image, from the unfolded barrel detector, allowing us to exploit the geometrical information and the translational symmetry [28].

In contrast, the use of recurrent neural networks (RNNs) [29,30], and by extension recursive neural networks [31], focused on the hierarchy relations between the jets’ constituents. While both approaches provided significant improvement, notably the ability to solely use low-level information and handle a variable number of inputs—a major limitation of DNNs—they both had limitations. For example, the binning into cells required it to produce an image coming with information loss in a CNN, and the ordering of the objects to be fed to an RNN can introduce hard-to-learn long-range dependencies, despite the introduction of long-short term memory (LSTM) and gated recurrent units (GRUs).

From these considerations started a wide search for models exploiting more useful symmetries and making use of as much information as physicists could consider. Many models were considered, but two of them are of particular note. The first is called ParticleNet [32] and expands the concept of convolution to a graph neural network (GNN). In the context of jet tagging, this is much more appropriate than a CNN because it can make better use of the sparsity of a jet representation. The other model is called ParT [33] and is built from the transformer architecture [34] that replaced RNNs thanks to their self-attention mechanism, better suited to handle any range of dependency. Notably, both models natively include a permutation symmetry, which is better suited to jets and could explain why they outperformed the previous ones.

This section will follow this historical evolution by starting with a DNN, augmenting the input variables, and then exploring the aforementioned more advanced architectures. The implementations of all models were performed using the standard Pytorch [35], PyG [36] for the GNN, and their training through lightning [37].

All model hyperparameters can be found in

Table A1. DNN hyperparameters.

Parameter	Value
Number of layers	6
Neurons per layer	64
Hidden activation	GELU
Output activation	None
Batchnorm	True

,

Table A2. RNN hyperparameters.

Parameter	Value
Embedding
Normalize inputs (batchnorm)	True
Neurons for each layer	32,64
Activation	ReLU
Recurrent cell
Cell type	GRU
Dimension	64
Layers	6
Bidirectional	False
Output (DNN)
Layers	3
Neurons per layer	64
Hidden activation	ReLU
Output activation	None

,

Table A3. Transformer hyperparameters.

Parameter	Value
Embedding/position encoding
Normalize inputs (batchnorm)	True
Neurons for each layer	32, 64
Activation	GELU
Encoder
Dimension	64
Encoder layers/class layers	6/2
Heads	8
Feed-forward dimension	128
Layernorm	True
Output (DNN)
Layers	3
Neurons per layer	64
Hidden activation	GELU
Output activation	None

and

Table A4. GNN hyperparameters.

Parameter	Value
Graph blocks
Block type	EdgeConv
Feature dimension	64
Number of blocks	6
DNN feedforward dimension	128
DNN batchnorm	True
Graphnorm	True
Activation	GELU
Node aggregation	Mean
Graph pooling
Aggregation	Mean
Jumping knowledge	Concatenation
Output (DNN)
Layers	3
Neurons per layer	64
Hidden activation	GELU
Output activation	None

. Additionally, all the models are trained with the RAdam optimizer [38], a batch size of 1024, mean squared error (MSE) loss function, and ReduceLROnPlateau learning rate scheduler—starting at ${10}^{-3}$ or ${10}^{-4}$ depending on the model. Around a million muons were generated in the setup of Figure 4a; 70% are used directly in the training, 10% are used for validation, and the remaining 20% to produce all the illustrations in the following sections, therefore providing unbiased performance estimation.

3.1. Deviations as Features

As a first attempt at leveraging more information than through Equation (4), a DNN will be trained on some features to try to regress the muon momentum. The momentum distribution follows the expected flux of Figure 1, convoluted with the detector acceptance. As this distribution spans many orders of magnitude, the actual target of the regression is ${log}_{10}P$, which will be discussed further in the following.

To have a direct comparison with Figure 4b, only the three scattering angles—as featured in Figure 4a—will be used as input. Similarly to the target, a ${log}_{10}$ operation is applied beforehand, and then they are preprocessed by removing the mean and scaling to unit variance. The result of the regression in Figure 5a shows that the DNN is able to exploit more information from these mere three inputs than Equation (4). There might, however, be more variables we can extract from the hits, for example, the deflection angles and the track residuals, as illustrated in Figure 6. When these features are added into the regression, while keeping the same model, ${log}_{10}$, and preprocessing, the gain is significant, as illustrated in Figure 5b.

The application of a ${log}_{10}$ to both inputs and targets came from the empirical observation that spans several orders of magnitude and that the angles scale non-linearly from it. High-momentum muons suffer very few deflections, and it becomes hard for the DNN to distinguish them above 100 GeV. The ${log}_{10}$ was then chosen to disentangle them, effectively “zooming in” on that high-momentum region. This, however, should not come as a surprise since Equation (3) implies that ${log}_{10}\theta \propto -{log}_{10}P$, as shown in Figure 7.

The results of Figure 5 suggest that more information can be obtained by using raw information, namely, detector hits. However, the actual information about the muon momentum is contained in the scattering of the muon compared to a linear track, as implied in Equation (3). High-energy muons have hits that almost form a perfect line, while low-momentum muons typically exhibit a curve or helix around a linear track fit, as shown in Figure 8a in three dimensions, and their projections on both horizontal axes are shown in Figure 8b. This feature comes from the fact that muons are deflected within a cone whose zenith angle is determined by Equation (2) and with a uniform azimuthal angle, giving the impression of a rotation around the linear track axis.

Therefore, what actually matters are the distances from the hits to their orthogonal projections onto the linear track. Instead of considering only the absolute distance, however, to avoid any loss of information, the two-dimensional vector in the perpendicular plane of the track will be used and referred to as a deviation.

3.2. Deviations as a Sequence

In order to be fed to an RNN, an order needs to be provided for the deviations. Fortunately, a natural choice consists of following the trajectory of the muon, in this case, from top to bottom. To exploit the azimuthal symmetry of the scattering in Equation (2), the track direction is used as the basis z-axis, with the deviations expressed in polar coordinates (r,$\theta$) instead of Cartesian (x,y). This choice of cylindrical basis improved the regression performance when used with sequential algorithms. Additionally, to circumvent the logarithmic relationship with the momentum illustrated in Figure 6, the log-modulus function $f\left(x\right)=\mathrm{sign}\left(x\right){log}_{10}\left(\right|x|+1)$ is applied on the magnitude of the deviation r, while preserving the angle. This function is used instead of a pure logarithm to avoid changing the sign of the deviation when $\left|r\right|<1$, effectively “zooming” in on the high-energy muon while preserving the shape of the deviations, as illustrated in Figure A1. Finally, the tracks are rotated around their axes, so that the polar angle starts at $\theta =0$. After passing through the RNN layers, the embedding of the last object is fed to a DNN in a many-to-one scheme. The result of the regression when the cylindrical coordinates (r,$\theta$,z) are provided to the RNN is in Figure 9a.

The transformer on the other hand is naturally permutation-invariant, which needs to be broken to keep the hierarchy of the deviations along the muon track. In language processing, this is performed through position encoding. Several methods exist; initially, absolute encoding—using sine and cosine functions—is employed, before other methods such as relative encoding are used. In order to fully make use of the three-dimensional geometrical information, the positions of the hits projected on the linear track are fed to a DNN to provide a learnable position encoding. The deviations expressed as r and $\theta$ are used as features and encoded using a similar DNN embedding. Both positions and features are first passed through a batch normalization layer, and then their respective encodings are summed and passed through a transformer encoder. Similarly to CaiT vision transformer [39], a token is passed through several cross-attention layers together with the encoder output to summarize its information, and then it is passed to a DNN. The result of the regression is in Figure 9b.

Treating the deviations as a sequence could be exploited to identify upcoming tracks of low-momentum muons acting as backgrounds in the context of transmission tomography since they lose energy when they pass through dense material and their deviations grow larger toward their exit side. This is partially visible in Appendix A, with a visible upward trend in the deviations in r for low-momentum muons. This effect is, however, extremely limited for high-momentum muons, and it could potentially limit the usefulness of this feature. As this is more relevant as a background rejection proposal, it was deemed outside the scope of this study.

3.3. Deviations as a Geometry

To use the geometric information in a CNN, the deviations of Figure A1 need to be binned into pixels. The main issue comes from the resolution of such pixels and the memory footprint. Figure 7 shows that no matter the basis used, the deviations span several orders of magnitude. Using a pixel size that would be fine enough to be able to measure high-momentum muon deviations—around 0.1 $\mathrm{mm}$ at 100 $\mathrm{GeV}$—while still having enough pixels to catch the range of low-momentum muons—a few $\mathrm{cms}$ at 1 $\mathrm{GeV}$—would mean thousands of pixels per dimension. This would not fit in memory to be used by a three-dimensional CNN. An attempt was made to use two-dimensional projections in the $x-y$ and $x-z$ planes, akin to multi-view CNNs from Ref. [40]. This essentially consists of turning the first two graphs of Appendix A into an image, with a single pixel per layer in the z direction to limit the image memory footprint. Still, to accommodate the different ranges, two CNNs per projection were needed with different resolutions, similar to the idea in Ref. [41]. While this setup did yield a decent regression, it was barely able to beat Equation (4). This was expected considering the information losses due to the projection and binning, as well as the sparsity of the image. However, it did prove that the geometric information could still be exploited.

This led to the use of a GNN instead, which, contrary to a CNN, is perfectly suited to sparse inputs. To build a graph from a track, the hits are rotated and then translated so that the track axis is vertical and centred on $(x=0,y=0)$. As also performed for the sequence processing, the log-modulus is applied to the distance between each hit and the track, while keeping the same angle to not modify the shape of the graph. The hits are then identified as the graph nodes, with the number of edges connecting each node with the closest neighbour treated as a hyperparameter. An example of graphs when each node is connected to the two closest nodes—unless the node is at the edges of the graph, in which case it only connects to its closest neighbour—is shown in Figure A2 with and without the log-modulus function.

Convolutions in a GNN can be reformulated in the message passing scheme, as implemented in PyG, where edge features of ${\mathbf{x}}_i^{\left(k\right)}$ at node i and layer k can be obtained from the previous layer $k-1$ as

$${\mathbf{x}}_i^{\left(k\right)}={\gamma }^{\left(k\right)}\left({\mathbf{x}}_i^{(k-1)},\underset{j\in \mathcal{N}\left(i\right)}{⨁}{\varphi }^{\left(k\right)}({\mathbf{x}}_i^{(k-1)},{\mathbf{x}}_j^{(k-1)},{\mathbf{e}}_{ij})\right),$$

(6)

where $\mathcal{N}\left(i\right)$ is the set of neighbours of node i, ${\mathbf{e}}_{ij}$ is the edge attribute between nodes i and j, ⨁ is a differentiable aggregation operation, and $\gamma$ and $\varphi$ are optional, complex, and differentiable functions, typically DNNs. In particular, for this application, the model used is edgeConv [42], specifically designed for point cloud convolutions, where Equation (6) is particularized to

$${\mathbf{x}}_i^{\left(k\right)}=\underset{j\in \mathcal{N}\left(i\right)}{⨁}{\varphi }^{\left(k\right)}({\mathbf{x}}_i^{(k-1)}\left|\right|{\mathbf{x}}_j^{(k-1)}-{\mathbf{x}}_i^{(k-1)}).$$

(7)

In this particular message passing mode, there is no edge attribute and the DNN $\varphi$ acts on the concatenated tensor of the node tensor and its difference with other connected nodes. For the graphs defined in Figure A2, x is just the position in the graph, and the mean pooling is used for ⨁.

Similarly to what was performed in Ref. [42], the output of each convolutional layer is concatenated to obtain a tensor with varying levels of abstraction, a method often referred to as jumping knowledge [43]. Residual connections are also included to help the training, as suggested in Ref. [44]. Before being fed to the DNN for the regression, the information from each node needs to go through a final differentiable aggregation operation. To that end, several options were tried, using single operators (mean, min, max, standard deviation) or a mix of them, concatenated. These operators are permutation-invariant, which, besides the convolution operation, implies that the deviations are treated as unordered, which explains the limited performance shown in Figure 9c. More complex and permutation-variant aggregation operations could potentially improve the regression [45,46].

3.4. Model Comparison

The real measure of the regression performance, besides the prediction bias, resides in the prediction resolution. This is illustrated as a function of the momentum in Figure 10 for the different models developed in this section. The sizes and inference times of the different models are in

Table 1. A comparison of the size of the different ML models in terms of the number of parameters and memory and inference time per muon (in a batch size of 1024 on a GPU).

Model	# Params	Size [MB]	Time [μs]
DNN	19 K	0.07	12
RNN	165 K	0.63	50
TNN	420 K	1.61	17
GNN	100 K	0.40	25

. The best model, the transformer (TNN), has a relative resolution below 10% and an inference time not far from the DNN thanks to GPU parallelization.

4. Real-Life Experiment Considerations

Besides good performance, there are a few other considerations for an ML model to be truly useful in a real-life environment. Section 3 only focused on an ideal detector, assuming all hits were detected and their position measured with infinite precision. In addition, the predicted momentum does not come with any prediction resolution, which would improve muon tomography inference. This section will address how these experimental factors can be measured or improved by re-training the models of Section 3 with the same simulated data, but this time incorporating a more realistic detector response.

4.1. Hit Detection Efficiency

The fact that some hits might not be detected needs to be accounted for in the momentum inference. This can be circumvented by applying a veto on any track that does not leave a hit in every detector layer. However, this would increase the necessary acquisition time to keep the same number of tracks. For example, with eight detector layers and a 95% efficiency—a value in the ballpark of most muon tomography detectors—34% of the tracks would have to be discarded. Additionally, relaxing the trigger selection to events having fewer hits than the total number of panels could increase the overall acceptance, as muon tracks missing the last few panels could now be recorded.

To study the effect of missing hits, several datasets are derived from the one used to train the models in Section 3. For each track, one or several hits are excluded from the linear fit and flagged as missing for the different momentum inference methods. With eight detector planes, it will be assumed that a veto is placed on tracks with less than six hits; therefore, datasets with two, one, and no missing hits will be included in the training. Their performance is in Figure 11.

Any hit missed on either side of a scattering material block will make it impossible to measure the scattering angle, but Equations (4) and (5) can still be used with fewer values, albeit with limited precision. However, in some unfortunate cases, no scattering angles can be measured and no momentum can be inferred, which is another limitation of that method.

The DNN model is more resilient to missed hits because more deflection angles and residuals can be measured. The main drawback of DNNs is that the number of input features is fixed. When a hit is missed, one alternative is to use a zero value for the features associated with the missed hit; this is called zero-padding. An increase in prediction resolution, as observed in Figure 11, is expected from the loss of information, but the mixed training setting seems to be able to mitigate that.

The transformer, as well as the GNN, can handle variable-length inputs and does not necessitate zero-padding. There is, however, a larger degradation in the resolution of the transformer compared to the DNN, which indicates that the model struggles to handle the change in the deviations’ representation. Comparing the case with no missed hit with Figure 10 additionally shows that the mixed training causes confusion. This seems to be linked to the cylindrical representation—and especially the polar angle—on Figure A1 that presents large variations when the track parameters obtained from a fit with fewer hits are modified. The GNN, on the other hand, even if less performant overall, suffers less from that effect as it uses a Cartesian representation (Appendix A). The DNN is also less affected because it relies more on the deviation angles.

4.2. Hit Spatial Resolution

To simulate the hit resolution, a Gaussian smearing is applied to the hit position in both the x and y directions. The width is chosen as 1 $\mathrm{mm}$. This is a typical, albeit slightly pessimistic, value for resistive plate and drift wire chamber detectors. The features and deviations are extracted from these smeared detector hits, and the different models are then trained on this modified dataset. The regression results are shown in Figure 12. The loss of performance matches the expectation from Figure 7c, as even for a muon with infinite momentum, the smeared hits will emulate deviations of a muon with a momentum of a couple dozen GeVs. While ML models perform better than Figure 12b because they leverage more features, they hit a plateau relatively early in the training, and it is unlikely that the information loss from the training can be recovered. In the transformer case in Figure 12c, if the training is performed for longer, the overall performance decreases, but it is able to predict higher-momentum muons.

4.3. Momentum Resolution

As shown in Figure 2, the density estimation performance can already gain from the momentum knowledge. Supplementing the prediction of our models with a measure of uncertainty could further increase that accuracy, especially in methods using statistical inference.

One option to include the resolution prediction is through the loss function. Instead of producing a single value, the ML model can include predictions for the quantiles of the output distribution. For example, one can add two quantile losses $\rho$ to the MSE loss, to regress the 25% and 75% quantiles as

$$\begin{array}{c}\hfill \mathrm{Loss}(y,\widehat{y})=\mathrm{MSE}(y,\widehat{y})+{\rho }_{25\%}(y,\widehat{y_{25\%}})+{\rho }_{75\%}(y,\widehat{y_{75\%}})\end{array}$$

(8)

$$\begin{array}{c}\hfill {\rho }_{\tau }(y,\widehat{y})=\left\{\begin{array}{c}\tau (y-\widehat{y}),\widehat{y}\le y,\hfill \\ (1-\tau )(\widehat{y}-y),\widehat{y}>y,\hfill \end{array}\right.\end{array}$$

(9)

where y is the prediction target, $\widehat{y}$ is the ML model prediction, $\tau$ is the quantile value, and $\widehat{y_{25\%}}$ and $\widehat{y_{75\%}}$ are the 25% and 75% quantile predictions, respectively. From $y={log}_{10}p$, the corresponding quantiles for the momentum can be obtained, and the resolution is defined as

$$\widehat{\Delta P}={\displaystyle \frac{\widehat{P_{75\%}}-\widehat{P_{25\%}}}{\widehat{P}}}.$$

(10)

The DNN model of Section 3.1 is modified with three output nodes—for $\widehat{y}$, $\widehat{y_{25\%}}$, and $\widehat{y_{75\%}}$—and the extended loss function of Equation (9). Any of the models discussed in Sections Section 3.2 and Section 3.3 could be modified easily, but the larger resolution of the DNN prediction allows a clearer illustration of the method, as shown in Figure 13 for both perfect and imperfect detectors. The resolutions can then be extracted using Equation (10), as shown in Figure 14.

These results show that with a very small modification of the training procedure, one can obtain a resolution as an ansatz for a prediction uncertainty. This resolution, however, has to be understood as a distribution property, not as an individual uncertainty for each muon. However, contrary to what was performed in Section 3.4, it can be obtained on a per-muon basis, not measured on a certain interval of predicted momentum. One possible extension of the method is to get rid of the MSE component and use multiple quantile predictions to obtain some information about the $f\left(p\right|\widehat{p})$ posterior distribution, although generative methods for probability density estimation might be more appropriate.

5. Outlook

Muon tomography, in both the transmission and scattering modes, is becoming a valuable tool for both scientific and civil purposes. As a non-invasive, large-scale, and passive detection method, it occupies a unique place in the wide range of imaging techniques. One factor limiting muon tomography performance is the dependence on the unknown momentum of the muon.

There are many ways to obtain that information, and in this paper we focused on inverting the scattering problem, using materials with known parameters to infer the momentum from the scatterings. We have compared the classical way of using physics’ first principles to several machine learning algorithms. We have shown that they not only improve the performance but are also able to deal more efficiently with experimental constraints of data taking, together with a proposal to augment the reconstruction algorithms with a prediction uncertainty.

To go beyond the stage of the proof of concept presented in this paper, there are several avenues we would like to address. First, the hyperparameters of the advanced algorithms have been manually selected and could be further optimized by a wider automatic search. In particular, the graph neural network presented here is only one model out of a wide range of models proposed in the literature. It might also be worth exploring whether both the sequence and geometry aspects of the deviations could be tackled by a combined model. Second, the use of the deviations from a linear track may have been successful, but it still represents a processing of the raw hit position data that the ML algorithms might not need. There might be a way to relax that step to improve both performance and inference time. Third, the use of the quantile regression can be used to obtain an estimate of the prediction resolution, and by extension a point-wise description of the posterior distribution. This would probably be better tackled by a generative model more suited to probability density estimation. Finally, the work presented in this paper is still far from practical application. Secondary charged particles such as electrons can be produced when muons pass through dense material; the additional hits they would generate in the detector have not been taken into account in our simulations. This effect, together with additional scattering, should also be reflected in more realistic detection planes. However, all these considerations—in simulation, tracking, and rejection algorithms—would need to be tailored for a specific detector technology. Additionally, the detector setup used throughout this paper is not particularly compact nor cheap in terms of material, which might be a limiting factor for its practical use. In order to accommodate the logistical constraints of an in situ deployment, the detector setup has to move toward a simpler configuration, with fewer detector panels and thinner absorbers. The methods developed in this paper should be more thoroughly tested on many different setups with possibly fewer detection layers and a smaller detector volume. And while active experiments can perform their own training on their specific geometry, it would be worth exploring a general-purpose algorithm that could be conditioned on a geometry that might not have been seen during training.