ANI-1xnr experimental NPT driver

An experimental direct-PyTorch molecular-dynamics driver for running the ANI-1xnr potential with an isotropic Berendsen-style NPT controller. It includes checkpoint/restart support, HDF5 trajectories, pressure and stress logging, and an optional TorchANI CUAev patch for strain-gradient stress.

Warning

This is research code. The barostat is useful for exploratory simulations and screening, but it is not a rigorous MTK or Parrinello-Rahman NPT integrator. Validate scientific results against an established MD implementation.

What is and is not included

This repository contains the NPT driver, launchers, examples, and the attributed TorchANI patch. It does not redistribute the ANI-1xnr model weights. The setup script obtains the model from the public upstream repository at a pinned commit.

Cluster setup

Install a PyTorch build compatible with the cluster's CUDA driver/toolkit first. Then:

git clone https://github.com/Austin243/ani-1xnr-npt-driver.git
cd ani-1xnr-npt-driver

# Use the Python executable from the intended conda environment or venv.
PYTHON="$CONDA_PREFIX/bin/python" scripts/bootstrap_cluster.sh

For a venv, substitute PYTHON="$VIRTUAL_ENV/bin/python". The bootstrap script:

1. installs the small Python dependencies in requirements-core.txt;
2. checks out ANI-1xnr under vendor/ani-1xnr;
3. checks out the pinned TorchANI source under vendor/torchani;
4. applies the CUAev stress patch and builds TorchANI's compiled extensions.

Set TORCHANI_BUILD_ARGS, for example TORCHANI_BUILD_ARGS="--sm 8.0", when the cluster needs an explicit GPU compute capability.

Verify the installation

The included water cell is only a loader/stress-path smoke test, not a scientifically meaningful ANI-1xnr NPT system.

PYTHON="$CONDA_PREFIX/bin/python" \
  scripts/run_ani_npt_config.sh examples/npt_config_minimal.json --check-only

Copy examples/npt_config_minimal.json, replace structure_path and output_dir, and then launch a run:

PYTHON="$CONDA_PREFIX/bin/python" \
  scripts/run_ani_npt_config.sh path/to/config.json

scripts/run_ani_npt_config_direct.sh is the original Perlmutter-specific launcher. It loads the original GCC/CUDA/CMake modules and should not be used unchanged on a different cluster.

See NPT_DRIVER.md for the algorithm, configuration keys, limitations, and patch details.

Provenance and licensing

The original code in this repository is released under the MIT License. The patched cuaev.cpp is a derivative of MIT-licensed TorchANI source and retains TorchANI's license and provenance under patches/torchani-cuaev-stress/. ANI-1xnr is downloaded separately and remains governed by its upstream terms. See THIRD_PARTY_NOTICES.md.

If you use ANI-1xnr in research, cite the model paper:

Zhang et al., Exploring the frontiers of condensed-phase chemistry with a general reactive machine learning potential, Nature Chemistry (2024), https://doi.org/10.1038/s41557-023-01427-3