Documentation Overview#
Neural Networks are promising models for enhancing the accuracy of classical molecular simulations. However, the training of accurate models is challenging. chemtrain is a framework for learning sophisticated Neural Network potential models by combining customizable training routines with advanced training algorithms. This combination enables the inclusion of high-quality reference data from simulations and experiments and the lowering computational demands of training through complementing algorithms with different advantages.
chemtrain is written in JAX, integrating with the differentiable MD engine JAX, M.D. Therefore, chemtrain leverages end-to-end differentiable physics and hardware acceleration through GPUs to provide flexibility at scale.
Installation#
Install chemtrain with pip:
pip install chemtrain --upgrade
The base installation uses JAX for CPU. chemtrain supports JAX versions from 0.5.0 up to, but not including, 0.12. JAX 0.10.2 is excluded. It requires JAX-MD 0.2.29 or newer. No special import order is required.
For NVIDIA CUDA Python packages, choose the CUDA major version installed on the target system:
pip install 'chemtrain[cuda12]' --upgrade
# or
pip install 'chemtrain[cuda13]' --upgrade
Note
These extras install JAX’s CUDA Python dependencies only. Building chemtrain-deploy and its PJRT runtime is a separate step. See Installation.
Advanced Installation#
Additional Packages#
Some parts of chemtrain require additional packages. To install these, provide the all option.
pip install 'chemtrain[all]' --upgrade
Installation from Source#
The lines below install chemtrain from source for development purposes.
git clone git@github.com:tummfm/chemtrain.git
pip install -e '.[all,docs,test]'
This command additionally installs the requirements to run the tests
pytest ./tests
and to build the documentation (e.g., in html)
make -C docs html
Getting Started#
To get started with chemtrain and with the most important algorithms, we provide simple toy examples. These examples are simple to run on the CPU and sufficient to illustrate the basic concepts of the algorithms:
To see the usage of chemtrain in real examples, we implemented the training procedures of some recent papers:
API Documentation#
chemtrain-deploy
Adapter Development