A Python utility library to perform (multimap) targeted free energy perturbation.
The library includes implementations for the following flows:
- Masked Autoregressive Flows [1] with a ready-to-use MADE [2] conditioner and the following transformers:
- Affine
- (Circular) neural spline [3, 4]
- Sum-of-squares polynomial [5]
- Moebius (based on that proposed in [4])
- Continuous normalizing flows [6] with a dynamics layer based on Equivariant graph neural networks [7].
CentroidandOrientedflows, which can be used to fix the reference system before the map.- A
PCAflow to perform the transformation in uncorrelated coordinates.
The library also includes several utilities:
- PyTorch wrappers for psi4 and the Atomistic Simulation Environment to evaluate potential energies and forces with multiple molecular simulation engines.
- A PyTorch
Datasetthat can wrap an MDAnalysisUniverseto train on molecular simulation trajectories. - A PyTorch-accelerated (T)FEP estimator.
- A PyTorch-accelerated bootstrap analysis utility.
- A simple storage utility class to save potential energies and log training information.
The library has the following required dependencies
python >= 3.9
pytorch >= 1.11
mdanalysis >= 2.0
pint
numpy
lightning >= 2.0
and the following optional dependencies
openmm # To evaluate the target potentials using the OpenMM Python library.
tblite-python # To evaluate the target potentials using the tblite Python library.
psi4 # To evaluate the target potentials using the psi4 Python library.
ase # To evaluate the target potentials using the Atomistic Simulation Environment (ASE) Python library.
torchdiffeq # To use continuous normalizing flows.
bgflow # To use the mixed internal-Cartesian coordinates Lightning module
einops # Required by bgflow
The suggested way of installing tfep is by first installing all the dependencies through conda/pip/setuptools,
and then installing tfep from the source (I plan to add a tfep conda package in the near future). Here is an
example that creates a separate conda environment with all the dependencies and installs tfep.
# Required dependencies.
conda create --name tfepenv python">=3.9" pytorch">=1.11" mdanalysis">=2.0" pint numpy lightning">=2.0" -c conda-forge
conda activate tfepenv
# Optional dependency using conda.
conda install einops -c conda-forge
# Optional dependency from source code.
git clone https://github.com/noegroup/bgflow.git
cd bgflow
pip install .
cd ..
# Optional dependency using pip.
pip install ase
# Install the package.
git clone https://github.com/andrrizzi/tfep.git
cd tfep
pip install .
# Or if you want to modify the source code, install it in editable mode.
# pip install -e .If you find this code useful, please cite the following paper:
Andrea Rizzi, Paolo Carloni, Michele Parrinello. Multimap targeted free energy estimation. arXiv preprint arXiv:2302.07683.
- Papamakarios G, Pavlakou T, Murray I. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems (2017).
- Germain M, Gregor K, Murray I, Larochelle H. Made: Masked autoencoder for distribution estimation. In International Conference on Machine Learning (2015).
- Durkan C, Bekasov A, Murray I, Papamakarios G. Neural spline flows. Advances in Neural Information Processing Systems (2019).
- Rezende DJ, Papamakarios G, Racaniere S, Albergo M, Kanwar G, Shanahan P, Cranmer K. Normalizing flows on tori and spheres. In International Conference on Machine Learning (2020).
- Jaini P, Selby KA, Yu Y. Sum-of-squares polynomial flow. In International Conference on Machine Learning (2019).
- Chen RT, Rubanova Y, Bettencourt J, Duvenaud DK. Neural ordinary differential equations. Advances in Neural Information Processing Systems (2018).
- Garcia Satorras V, Hoogeboom E, Fuchs F, Posner I, Welling M. E(n) Equivariant Normalizing Flows. Advances in Neural Information Processing Systems. (2021).
Project based on the Computational Molecular Science Python Cookiecutter version 1.5.