clarm

---

Note: EIDR number O#4748 - Conditional latent autoregressive recurrent model (CLARM). This program is Open-Source under the BSD-3 License.

---

Conditional Latent Autoregressive Recurrent Model for spatiotemporal learning

CLARM is a unsupervised deep learning approach to learn spatiotemporal dynamics. CLARM consists of a CVAE and LSTM combined together in an autoregressive loop to learn spatial and temporal dynamics independently. Here, CVAE learns spatial correlations of the high dimensional objects in a low-dimensional latent space distribution and a RNN like Long Short-Term Memory (LSTM) learn the temporal dynamics within the latent space. Using CVAE part, new realistic samples can generated by conditionally sampling the latent space followed by decoding. Using LSTM part, the future states can be forecasted given few initial states as inputs. The proposed model addresses challenges faced by existing models in learning spatiotemporal dynamics, such as expensive computations in ConvLSTM and DCGAN, mixed spatial features and temporal dynamics in 3DCNN/4DCNN, computational complexity, lack of scalability, and robustness issues in GNNs, linear representation of data in the latent space through PCA-LSTM, and explainability concerns in the latent space with like AE-LSTM.

The CLARM is used to learn spatialtemporal evolution of 6D phase space of charged beams in particle accelerators. The model can generate phase space projections at various accelerator modules by sampling and decoding the latent space representation. The model also forecasts 6D phase space of charged particles in downstream modules from limited phase space information at upstream locations. This repository contains codes accompanying the paper. The dataset accompying the paper is available at Zenodo.

Latent space visualization:

Generative ability of CLARM:

Forecasting ability of CLARM:

For more information:

1. Arxiv Link of the paper: A conditional latent autoregressive recurrent model for generation and forecasting beam dynamics in particle accelerators (LA-UR-24-22420).
2. HPSim, an advanced, open-source tool developed at LANL, enables rapid, online simulations of multiple-particle beam dynamics. The basic source code of HPSim is available here (now inactive repository). The dataset is collected from the new version of HPSim, which is available with LANL (not opensourced).
3. The dataset accompying the paper is available at Zenodo (LA-UR-24-22495).

About the repository:

1. The python code is written using pytorch, cuda 11.8 and cudnn 8.x.
2. "clarm_main.py" is the main python file (run in spyder).
3. lib folder has .py files required to run "clarm_main.py"
4. "resnet5_classifer.py" is the resnet50 model used as a classifer to classify 6d phase space of charged particles beam into 48 classes (accelerating sections/modules). The lib folder contains "pretrainedresnet50.py", which is used in the "clarm_main.py" for calculating the fid, precision, and recall. The weights of trained resnet50 classifer is not uploaded to the insufficient space of the repository. Email at mrautela@lanl.gov.