Supplemental Materials: The memory scaling of reverse-mode differentiation in particle accelerator simulations with space charge

This dataset includes the necessary material to produce the results in the paper "The memory scaling of reverse-mode differentiation in particle accelerator simulations with space charge". This includes several data collection scripts, the data collected using those scripts, and the code used to analize and produce the figures from that data. The required benchmarking code for Cheetah is included in this archive. This code is based on Cheetah v0.7.4, and includes the changes required for both the memory allocation/trace info and the memory use per section of Space Charge. The official Github repository for Cheetah can be found here: https://github.com/desy-ml/cheetah. All the simulations were run on Perlmutter at NERSC (LBNL) using the NVIDIA A100 (80GB) GPU partition. A GPU is required for correctly benchmarking memory use in Cheetah. Archive structure:- `README.md`: Overall description- `data_analysis.py`: The analysis script, responsible for cleaning the raw data in `data/memory_data.csv` and generating the regression fits in `data/analyzed_data.csv`, as well as plotting the fits of each segment and saving those figures into `outputs/`.- `data_collection.py`: The data collection script, responsible of running space charge kick simulations in a variety of configurations and saving the memory usage info in `data/memory_data.csv`- `generate_memory_plots.py`: The memory plot script, responsible for running the accelerator simulation with space charge in differentiable and non-differentiable mode, and saving the side-by-side comparison of the memory usage graphs in `outputs/`- `regression.py`: A utility file that contains the code for the weighted least-squares regression used to calculate the coefficients of the memory use.- `simulation.py`: A utility file that contains the code for running both full accelerator simulations and single space charge kicks and extracting the memory use data.- `analyzed_data.csv`: A csv containing the coefficients of the memory use of each section of space charge kick and their standard deviations.- `memory_data.csv`: A csv containing the data of the memory use of each section of space charge kick in a variety of configurations.- `cheetah_benchmark_B1.ipynb`: The code for the benchmark from Appendix B1.- `cheetah_benchmark_B2.ipynb`: The code for the benchmark from Appendix B2.- `cheetah.tar`: A modified version of Cheetah based on v0.7.4, which includes the changes required for both the memory allocation/trace info and the memory use per section of Space Charge.