Reactive molecular dynamics framework for topology-changing simulations using JAX-MD.
The package supports:
- Classical OPLS-AA force fields
- Stochastic topology-changing reactions
- Metropolis-based reaction acceptance
- Rate-based reaction kinetics
- Activation-energy-driven reaction kinetics (rs@MD / Heuer-style)
- Optional MACE-JAX / MLIP energy evaluation for the Metropolis reaction acceptance step
The current implementation includes a LiPF6 decomposition example:
LiPF6 -> LiF + PF5
The LiPF6 reaction is controlled by a single reaction coordinate
sigma = d(P-F) - d(Li-F)
where the selected F atom is the leaving fluorine. Negative sigma values are reactant-like, while positive sigma values are product-like. The reaction code no longer uses independent hard Li-F or P-F distance gates to decide reactivity; the individual distances are used only to compute sigma and for diagnostic logging.
Clone the repository:
git clone https://github.com/BAMeScience/reactive_md_project.git
cd reactive_md_projectInstall in editable mode:
pip install -e .or with uv:
uv pip install -e .Install the package together with the test/development tools:
uv pip install -e .
uv pip install pytest ruffIf you prefer optional dependency groups, move the dev dependency list into [project.optional-dependencies] in pyproject.toml, then use:
uv pip install -e ".[dev]"For MACE-assisted reaction acceptance:
uv pip install -e ".[mace]"The currently tested dependency stack is:
jax 0.10.1
jaxlib 0.10.1
flax 0.12.2
haiku 0.0.16
For GPU support, install the appropriate CUDA-enabled JAX version, for example:
uv pip install -U "jax[cuda13]"The package provides the command-line executable:
reactive-mdDisplay all available options:
reactive-md --helpThe code requires:
- A LAMMPS data file
- A LAMMPS settings file
Example:
reactive-md \
--data system.data \
--settings system.in.settingsThe LiPF6 decomposition reaction uses the Fattebert-inspired coordinate:
sigma = d(P-F) - d(Li-F)
Interpretation:
sigma < 0: the leaving F is closer to P than to Li, so the geometry is PF6-like.sigma = 0: the leaving F is equally distant from P and Li.sigma > 0: the leaving F is closer to Li than to P, so the geometry is product-like.
Candidate reactions are ranked by sigma. The reaction probability is derived from sigma using a smooth sigmoid-like factor controlled by:
--sigma-mid
--sigma-width
--sigma-mid is the sigma value where the geometric reaction factor is 0.5. --sigma-width controls how sharply the probability changes around sigma_mid.
Example:
reactive-md \
--data system.data \
--settings system.in.settings \
--reaction-mode rate \
--reaction-rate-ps 0.01 \
--sigma-mid 0.0 \
--sigma-width 0.2r_pf_probe is not a reaction gate. It is only used after a reaction has been accepted to move the leaving F away from P before product-side FIRE relaxation.
Topology changes are proposed using the reaction coordinate σ and accepted according to a hybrid kinetic/thermodynamic Metropolis scheme.
The reaction coordinate σ is first evaluated to determine whether a reaction trial should be attempted. Only if this kinetic condition is satisfied is the expensive trial topology generated, FIRE-relaxed, and evaluated with the selected energy model (classical force field or MACE). The final thermodynamic acceptance is then determined using the Metropolis criterion based on the energy difference ΔE. Trial reactions yielding non-finite energy differences are automatically rejected.
reactive-md \
--data system.data \
--settings system.in.settings \
--reaction-mode metropolis \
--sigma-mid 0.0 \
--sigma-width 0.2In classical mode, the Metropolis energy difference is computed with the classical force field. If MACE support is enabled, the Metropolis energy difference is evaluated with MACE-JAX instead.
Topology changes are accepted according to a prescribed base reaction rate modulated by the sigma-dependent geometric factor:
reactive-md \
--data system.data \
--settings system.in.settings \
--reaction-mode rate \
--reaction-rate-ps 0.01 \
--sigma-mid 0.0 \
--sigma-width 0.2The effective rate is
k_eff = k_base * f(sigma)
and the reaction probability during one reactive interval is
p = 1 - exp(-k_eff * dt)
where:
k_baseis the prescribed or activation-energy-derived base rate in ps^-1.f(sigma)is the sigma-dependent geometric factor.dtis the reactive check interval.
Instead of specifying a reaction rate directly, provide an activation barrier:
reactive-md \
--data system.data \
--settings system.in.settings \
--reaction-mode rate \
--activation-energy-eV 0.20 \
--sigma-mid 0.0 \
--sigma-width 0.2The base reaction rate is computed using transition-state theory:
k(T) = (kB * T / h) * exp(-Ea / (kB * T))
and then modulated by the sigma-dependent geometric factor.
To evaluate reaction energies with MACE-JAX:
reactive-md \
--data system.data \
--settings system.in.settings \
--reaction-mode metropolis \
--use-mace-mcThe molecular dynamics trajectory still uses the classical force field. MACE is used only for the Metropolis reaction acceptance energy evaluation. The FIRE relaxation after the topology proposal currently uses the product-side classical force field.
Write a LAMMPS-style trajectory:
reactive-md \
--data system.data \
--settings system.in.settings \
--dump-file traj.dumpCompressed output:
reactive-md \
--data system.data \
--settings system.in.settings \
--dump-file traj.dump.gzRecord accepted reactions:
reactive-md \
--data system.data \
--settings system.in.settings \
--event-log-file reactions.csvRecord sigma-ranked reaction candidates:
reactive-md \
--data system.data \
--settings system.in.settings \
--candidate-log-file candidates.csv \
--candidate-log-top-n 20The candidate log contains sigma-ranked candidates. It no longer contains passes_lif, passes_pf, or passes_all, because those belonged to the old hard-distance-gated model.
Useful diagnostic columns include:
step, mode, rank, pf6_index, li_idx, leave_F, d_lif, d_pf, sigma, accepted
Run tests:
pytestRun only the reaction-coordinate tests:
pytest tests/test_reaction_sigma.pyRun with verbose output:
pytest -vreactive_md/
├── config.py
├── forcefield.py
├── reaction.py
├── md_driver.py
├── main.py
├── lammps_io.py
├── extract_params_oplsaa.py
├── reactions/
│ ├── lipf6.py
│ └── templates_pf5.py
tests/
├── test_candidate_selection.py
├── test_config.py
├── test_forcefield.py
├── test_lipf6.py
├── test_main.py
├── test_main_cli.py
├── test_md_driver.py
├── test_metropolis.py
├── test_reaction_coordinate.py
├── test_template_pf5.py
└── test_topology_change.py
| Test | Purpose |
| ----------------------------- | ----------------------------------------------- |
| `test_candidate_selection.py` | Reaction candidate identification and ranking. |
| `test_config.py` | Configuration defaults and validation. |
| `test_forcefield.py` | Force-field construction and evaluation. |
| `test_lipf6.py` | LiPF₆ topology and chemistry utilities. |
| `test_main.py` | Core simulation workflow. |
| `test_main_cli.py` | Command-line interface. |
| `test_md_driver.py` | MD driver and reaction integration. |
| `test_metropolis.py` | Metropolis acceptance criterion. |
| `test_reaction_coordinate.py` | Reaction coordinate (σ) evaluation. |
| `test_template_pf5.py` | PF₅ template generation. |
| `test_topology_change.py` | Topology modification after accepted reactions. |
If you use this software in academic work, please cite:
- JAX-MD
@inproceedings{jaxmd2020,
author = {Schoenholz, Samuel S. and Cubuk, Ekin D.},
booktitle = {Advances in Neural Information Processing Systems},
publisher = {Curran Associates, Inc.},
title = {JAX M.D. A Framework for Differentiable Physics},
url = {https://papers.nips.cc/paper/2020/file/83d3d4b6c9579515e1679aca8cbc8033-Paper.pdf},
volume = {33},
year = {2020}
}
- MACE / MACE-JAX, if used https://github.com/ACEsuit/mace-jax
@inproceedings{Batatia2022mace,
title={{MACE}: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields},
author={Ilyes Batatia and David Peter Kovacs and Gregor N. C. Simm and Christoph Ortner and Gabor Csanyi},
booktitle={Advances in Neural Information Processing Systems},
editor={Alice H. Oh and Alekh Agarwal and Danielle Belgrave and Kyunghyun Cho},
year={2022},
url={https://openreview.net/forum?id=YPpSngE-ZU}
}
@misc{Batatia2022Design,
title = {The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials},
author = {Batatia, Ilyes and Batzner, Simon and Kov{\'a}cs, D{\'a}vid P{\'e}ter and Musaelian, Albert and Simm, Gregor N. C. and Drautz, Ralf and Ortner, Christoph and Kozinsky, Boris and Cs{\'a}nyi, G{\'a}bor},
year = {2022},
number = {arXiv:2205.06643},
eprint = {2205.06643},
eprinttype = {arxiv},
doi = {10.48550/arXiv.2205.06643},
archiveprefix = {arXiv}
} - The original rs@MD methodology by Heuer and co-workers, J. Chem. Theory Comput. 2021, 17, 1074-1085
- The ab initio work motivating the LiPF6 reaction coordinate, Journal of The Electrochemical Society, 2024 171 080505
See LICENSE file.