Value proposition Detector simulations allow scientists to design detectors and perform physics analyses. The simulation toolkit that has been developed and performs particle physics simulations based on Monte Carlo (MC) methods is Geant4.
The detailed particle MC simulations are inherently slow. Simulations have a crucial role in HEP experiments, and at the same time are very resource-intensive from the computing perspective. Therefore, HEP community is highly motivated to explore fast alternatives, with deep learning based fast simulation being the most promising.
3DGAN consists of a fast alternative to MC, with remarkable results in terms of speed up. 3DGAN was the first effort where the detector output was generated employing three dimensional convolutions, an approach for retaining correlations in all three spatial dimensions .
Release Notes
The fast particle detector simulation with GAN thematic module consists of two inseparable components. These two components are the machine learning framework and the particle simulation framework. An implementation of the 3DGAN approach has been developed and a more detailed description follows. The code is available on GitHub and it has been tested and run on a single Linux node using GPU infrastructure.
3DGAN is being trained to produce images similar to the ones that are produced by Monte Carlo simulations. As the calorimeter detectors consist of layers of cells, those cells are modelled as monochromatic pixelated images with the cell energy depositions being the pixel intensities. 3DGAN consists of 2 networks, a generator and a discriminator, the two networks compete with each other trying to optimise a loss function until the convergence point, where the discriminator won’t be able to distinguish the images generated by the generator from the real images. Each network is being trained using 3-dimensional convolution layers to represent the 3 spatial dimensions of the calorimeter images.
The generator network implements stochasticity through a latent vector drawn from a Gaussian distribution. The generator input includes the primary particle’s initial energy and the angle at which it entered the detector, concatenated to the latent vector. The generator network then maps the input to a layer of linear neurons followed by 3D convolutional layers. The discriminator input is an image while the network has only 3D convolutional layers. Batch normalisation is performed between the layers and the LeakyRelu[1] activation function is used for the discriminator layers while the Relu13 activation function is used for the generator layers. The model’s loss function is the weighted sum of individual losses concerning the discriminator outputs and domain-related constraints, which are essential to achieve high-level agreement over the very large dynamic range of the image pixel intensity distribution in a HEP task. The training of this model was inspired by the concept of transfer learning. This means that the 3DGAN was trained first for images in a limited energy range and after the GAN converged, the same trained model was further trained with the data from the whole available energy range.
Currently, the 3DGAN training workflow consists of several other processes, the data pre-processing process, the model definition, and the training process. The validation and hyperparameter optimization processes are under research.
The dataset used for studying and developing the 3DGAN model (public dataset) consists of calorimeter 3D images/arrays of energy depositions with shape 51x51x25, which represent the particle showers. These images were created from simulations performed with Geant4 software. The output of the Geant4 simulation is ROOT[2] files, which need to be converted into a ML-friendly format HDF5 in order to train the model.
The preprocessing is responsible for preparing (cleaning, scaling, etc.) and converting into a suitable format (HDF5 format) the simulated data created by Geant4 (ROOT format). It also encodes the input information such as the calorimeter’s geometry identifier, the energy of the primary particle initiating the shower, the angle at which the particle enters the detector, and also its type and/or initial position. The pre-processed data are then passed to the GAN model (currently developed using Tensorflow v1 and v2[3] , as well as in PyTorch Lightning[4]) for training. The validation process will verify the performance through a set of physics-motivated steps, both at the single image quality level and at the sample level.
During pre-processing, simulation inputs are defined and encoded, i.e. the detector geometry, the energy and angle of the incoming particle. The performance of the model is evaluated during validation processes through the creation of histograms describing particle shower observables. Shower observables are among others, total energy distribution (sum of all cell energy deposits), cell energy deposits distribution, longitudinal profile which represents the energy deposited by a shower as a function of the depth of the calorimeter and lateral profile which represents the energy density distribution as a function of the radius of the calorimeter. Moreover, the physics-based validation process will include accuracy verification of those key distributions’ first moments and precise evaluation of the tails of distributions that usually require larger amounts of samples. The original data coming from Geant4 and the 3DGAN data distributions will be compared during this evaluation process.
Concerning the particle simulation framework, there have been testbeds developed at CERN that incorporate different ML models than the 3DGAN. Our current focus doesn’t include integrating the 3DGAN model in the simulation framework that uses the Geant4 environment. An example of the use of ML techniques for the fast detector simulation and how to incorporate inference libraries into Geant4 is the Par04 example developed by the Geant4 community and can be found on CERN Gitlab[5]. The ML model used in this example is a Variational Autoencoder (VAE), trained externally in Python on full Geant4 detector simulation response data.
The training of 3DGAN implemented in PyTorch Lightning and itwinai[6], was conducted on JSC computing resources, utilizing restricted energy range (100-200 GeV) data for 130 epochs and full energy range (2-500 GeV) data for 30 epochs per half of the dataset due to memory constraints. The model conditioning includes theta angles (60°-120°) and primary energy levels (2-500 GeV), ensuring uniform distribution across the dataset. The trained weights are available for inference. The dataset consists of electron showers in a 51×51×25 voxel grid.
The initial prototype was implemented in TensorFlow/Keras, and its parameterization was precisely replicated during the transition to PyTorch Lightning. While the PyTorch version does not perform as well as the Keras prototype, it still captures most underlying event generation patterns effectively.
Before inference, the training loss was closely monitored, with the best PyTorch weights selected based on validation loss. Both TensorFlow and PyTorch-trained weights can be used for inference, providing flexibility. The model’s performance was evaluated using physics analysis plots, including 2D projections, which compare GAN-generated events to Geant4 (G4) simulated data. The results, available for both frameworks, highlight the model’s capability to reproduce complex shower patterns.
The TensorFlow model exhibits a tighter, more concentrated core in the projections, closely matching the G4 simulation, particularly in the XZ and YZ planes. The energy deposition is more symmetric with smoother gradients around the core. The PyTorch model shows increased dispersion, with the core appearing broader and less uniform, especially in the YZ projection. This suggests the PyTorch model current state captures the overall shower pattern but struggles with fine-grained details. Though, both models capture the high-intensity core (red regions).
While the PyTorch model effectively replicates the general shower structure, the TensorFlow model demonstrates superior performance in terms of spatial accuracy, angular precision, and smoother energy deposition profiles. This aligns with the observation that the PyTorch version, despite capturing the main patterns, currently underperforms compared to the Keras prototype, especially in the angular reconstruction and shower shape fidelity. Further tuning of the PyTorch model — adjusting learning rates, loss functions, or regularization techniques — could help bridge this performance gap.
[1] https://en.wikipedia.org/wiki/Rectifier_(neural_networks)
[2] ROOT: https://root.cern/
[3] Tensorflow: https://www.tensorflow.org/
[4] PyTorch Lightning: https://lightning.ai/
[5] https://gitlab.cern.ch/geant4/geant4/-/tree/master/examples/extended/parameterisations/Par04
[6] https://github.com/interTwin-eu/itwinai
Future Plans
The developments of the thematic module are almost completed from the aspect of the 3DGAN component. Studies might be conducted on existing solutions for parallel training and hyperparameter optimization, but are not essential. The implementation of the CaloINN component is in progress. Simultaneously, the validation of the CaloINN algorithm on different detector simulations is also underway. In collaboration with the HEP community, different validation techniques are studied with the goal of identifying the technique most aligned with the needs of the fast detector simulation use case.
