Training molecular dynamics potentials.
Project description
Training Molecular Dynamics Potentials in JAX
chemtrain is a library for training MD potentials and integrates with JAX, M.D. as its differentiable MD engine. Note that this is the first alpha release of chemtrain, expect breaking changes. Over the course of the next weeks, several updates will extend its functionalities and documentation.
Features
- Differentiable Trajectory Reweighting (DiffTRe) to train on experimental data
- Force Matching and Relative Entropy Minimization for coarse-grained systems
Incoming features:
- Uncertainty Quantification via Deep Ensembles and Stochastic gradient MCMC
- Active Learning
- Hybrid Trainers to combine different training methods
- Molecular Property Prediction
Getting started
To get started with training MD potentials, using existing trainers that implement standard training schemes is the most straightforward apprach. Please refer to the provided examples as a refernece on how to set up the corresponding training environment and dataset.
trainer = trainers.Difftre(init_params,
optimizer)
trainer.add_statepoint(energy_fn_template, simulator_template,
neighbor_fn, timings, kbt, compute_fns, reference_state,
targets)
trainer.train(num_updates)
More advanced users may want to extend existing trainers or combine different trainers to implement custom training pipelines.
Installation
chemtrain can be installed via pip:
pip install jax-dimenet
Requirements
The repository uses the following packages:
'jax',
'jax-md',
'jax-sgmc',
'optax',
'dm-haiku',
'sympy',
'tree_math',
'cloudpickle',
'chex',
The code runs with Python >=3.8.
Contribution
Contributions are always welcome! Please open a pull request to discuss the code additions.
Since this is a very early alpha release, do not hesitate to reach out if some feature or example is failing. If the specific feature is a priority for you, its support can be accelerated.
Contact
For questions, please contact stephan.thaler@tum.de or open an Issue on GitHub.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for chemtrain-0.0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 8a9beeedbd20466932e1cff9fb267282d44792fa680615294792df33223a5f70 |
|
MD5 | 747e877c1411d424518bfb6d0f95cfb1 |
|
BLAKE2b-256 | 25a2cab8d02a613fd71f37a80b11a62365d0e0ead4edca17c54a9caedb31977d |