trajcast.models-arxiv2025

TrajCast Models arXiv 2025

This repository comprises a collection of TrajCast models, a framework for forecasting molecular dynamics (MD) trajectories using autoregressive equivariant message-passing networks. Provided with a starting configuration comprising information about atom types, atomic positions, and velocities, TrajCast predicts displacements and new velocities for later state at time interval Δt. By rolling-out the predictions of TrajCast autoregressivley, a MD trajectory of the system of interest of arbitrary length can be generated. Naturally, using larger time intervals than classical MD simulations, TrajCast can generate long trajectories with fewer steps. We provide example of how this is and single step inference is done in this notebook.

Weight and Architecture

We provide weights for mimicking the temporal evolution of three different systems. All models comprise about 2.3M parameters, while being trained on a different forecasting time interval:

• paracetamol: Weights for gas-phase paracetamol using a time interval of 7 fs.
• quartz: Weights for crystalline quartz using a time interval of 30 fs.
• water: Weights for liquid bulk water using a time interval of 5 fs.

Note: We provide each model based on two different O(3) backends: e3nn and cuEquivariance. Choose the state dictionary and config.yaml dependent on whether you have CUDA and cuEquivariance installed. Please note that depending on the device used to initialize a model with the cuEquivariance backend, some parameter names may differ.

Below we provide an overview of our architecture. For more information we refer to our preprint and code.

<p align="center"> <img src="arch.svg"> </p>

(A) Autoregressive workflow: An atomistic system at time t0 is passed through an equivariant MPNN (grey box) to predict the new positions and velocities at time t1. Atomic attributes (positions, velocities, chemical elements) are encoded into initial features, which are refined over T message passing blocks. Estimates of the displacement and velocity vectors are generated based on the final features. These are then refined to ensure momentum conservation. The trajectory is built by rolling out predictions, where outputs from one step serve as inputs for the next. A thermostat ensures sampling from the canonical (NVT) ensemble at constant temperature T, with states following the Boltzmann distribution. (B) The embedding block encodes node and edge attributes and generates the initial features. (C) Messages are constructed by convolving latent features with filters derived from a learnable radial basis and the spherical harmonics expansion of edge vectors. (D) In the update block, messages from neighbors are pooled and combined via a tensor product with velocity vectors in a learnable radial and spherical harmonic basis. The result is passed through a non-linearity and added to the previous layer’s features, weighted by the node’s chemical element. (E) Conservation of total linear and angular momentum is enforced by adjusting the displacements and velocities.

Citation

If you decide to use this dataset, please consider citing our preprint

@misc{thiemann2025Force-Free,
      title={Force-Free Molecular Dynamics Through Autoregressive Equivariant Networks}, 
      author={Fabian L. Thiemann and Thiago Reschützegger and Massimiliano Esposito and Tseden Taddese and Juan D. Olarte-Plata and Fausto Martelli},
      year={2025},
      eprint={2503.23794},
      archivePrefix={arXiv},
      primaryClass={physics.comp-ph},
      url={https://arxiv.org/abs/2503.23794}, 
}

Contact

For more information or if you like to contribute, please reach out to [email protected]