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#
chemtrain can be installed with pip:
pip install chemtrain --upgrade
The above command installs JAX for CPU. Running chemtrain on the GPU requires the installation of a particular JAX version. Please follow the JAX Installation Instructions.
Note
Chemtrain installs jax == 0.4.30 which is, in principle, incompatible
with jax_md <= 0.1.29 but resolves an
XLA issue which can prevent
training.
By importing chemtrain or the jax_md_mod module
before importing jax_md, the compatibility is restored by a simple
patch.
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: