Pytorch differentiable molecular dynamics
This is not just a regular simulator but a DIFFERENTIABLE simulator!
PyTorch code for End-to-end differetiable molecular simulations. More docs and tutorials are comings. This repo is under heavy development, your contribution is very much welcomed.
I highly recommend creating a dedicated conda environment via:
conda create -n mdgrad python=3.8
Download and install
git clone https://github.com/torchmd/mdgrad.git
cd mdgrad
conda activate mdgrad
pip install -r requirements.txt # I have tested this, it should work
pip install -e . # -e is useful if you want to edit the source code
Wang, W., Axelrod, S., & Gómez-Bombarelli, R. (2020). Differentiable Molecular Simulations for Control and Learning. https://arxiv.org/abs/2003.00868
This repo features the following demos:
Differentiable folding of a polymer
Learning interactions from observables (pair correlation function, velocity auto-correlations)
Quantum isomerization of a minimal retinal model
# Define a box of particles
L = 1.6
atoms = FaceCenteredCubic(symbol='H', size=(3, 3, 3), latticeconstant=L, pbc=True)
# use System to wrap ase.atoms
from torchmd.system import System
device = 'cuda:0'
system = System(atoms, device=device)
system.set_temperature(1.0)
# Define interactions
from torchmd.potentials import ExcludedVolume
pair = PairPotentials(system, ExcludedVolume, **{'epsilon': 1.0, 'sigma': 1.0,"power": 12}, cutoff=2.5).to(device)
# Define simulation
from torchmd.md import NoseHooverChain
integrator = NoseHooverChain(model,
system,
Q=50.0,
T=1.0,
num_chains=5,
adjoint=True).to(device)
sim = Simulations(system, integrator)
# Simulate
v_t, q_t, pv_t = sim.simulate(steps=50, frequency=50, dt=0.01) #v_t: velocity q_t: position pv_t bath: variables
# Compute observable
obs = rdf(system, nbins=100, r_range=(0.75, 2.5))
_, _, g = obs(q_t)
g.sum().backward()
# You will find out g can be backpropagated for gradient cumulation, give it a try!
Backpropagating through the trajectory to train a GNN that reproduces a target pair distribution function. We demonstrated the fitting of water rdf (Oxygen-Oxygen) at 298k with differentiable simulations
Folding a polymer with Graph Neural Networks
We fit electric field to optimize efficiency of a quantum isomerization process for retinal molecule