Differentiable Surrogate for Detector Simulation and Design with Diffusion Models

This repository contains the code for the paper:

"Differentiable Surrogate for Detector Simulation and Design with Diffusion Models"
by Xuan Tung Nguyen et al. Machine Learning: Science and Technology, 7(2), 025061 (2026) https://doi.org/10.1088/2632-2153/ae5c56

The project provides a conditional denoising-diffusion probabilistic model (DDPM) to simulate electromagnetic calorimeter showers. The model can generate high-fidelity, differentiable shower distributions conditioned on detector geometry, material, and incoming particle energy.

Table of Contents

• Installation
• Data
• Training
• Evaluation
• Citation
• Acknowledgments

---

Installation

git clone https://github.com/X-T-Nguyen/Diffusion-Surrogate-Detector-Design.git
cd Diffusion-Surrogate-Detector-Design
conda create -n diff-surrogate python=3.11
conda activate diff-surrogate
pip install -r requirements.txt

Data

The dataset used in this work is publicly available on Zenodo: https://doi.org/10.5281/zenodo.17105137

Training

Pre-training:

python MainCondition.py

Post-training:

python fine_tune.py

Evaluation:

The evaluation scripts are provided to assess model performance and generate key analysis outputs. These include:

Visual comparison between generated and ground-truth showers

python shower_plot.py

Computation and plotting of longitudinal and transverse energy profiles

python edep_plot.py

Evaluation of physical fidelity metrics

python metric_plot.py

Gradient-based analysis comparing the foundation model and the post-trained model

python grad_plot.py

Citation

If you use this code in your research, please cite:

@article{nguyen2026diffsurrogate,
  title   = {Differentiable Surrogate for Detector Simulation and Design with Diffusion Models},
  author  = {Nguyen, Xuan Tung and Chen, Long and Dorigo, Tommaso and Gauger, Nicolas R. and Vischia, Pietro and Nardi, Federico and Awais, Muhammad and Hanif, Hamza and Abbas, Shahzaib and Kapoor, Rukshak},
  journal = {Machine Learning: Science and Technology},
  volume  = {7},
  number  = {2},
  pages   = {025061},
  year    = {2026},
  doi     = {10.1088/2632-2153/ae5c56},
  publisher = {IOP Publishing}
}

Acknowledgments

This work was carried out within the MODE Collaboration, and we thank its members for valuable discussions.

We acknowledge funding and computing support from the German National High Performance Computing (NHR) association (Center NHR South-West), the Alliance for High Performance Computing in Rhineland-Palatinate (AHRP) via the Elwetritsch cluster at RPTU Kaiserslautern-Landau, and the Artemisa computing infrastructure funded by the European Union ERDF and the Comunitat Valenciana.

Pietro Vischia was supported by the Ramón y Cajal programme (Project No. RYC2021-033305-I) funded by MCIN/AEI/10.13039/501100011033 and by the European Union NextGenerationEU/PRTR.

We also acknowledge technical support from the Instituto de Física Corpuscular (IFIC, CSIC–UV).