Langevin Dynamics simulating methane using ANI1x and plotting H-C-H angles

langevin_dyamics_methane.py

from openmmtools.constants import kB
from simtk import unit
import numpy as np
from tqdm import tqdm
from IPython.core.display import display, HTML
from IPython.display import SVG
from rdkit.Chem.Draw import IPythonConsole
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import rdFMCS
import copy
import random, math
import os
import mdtraj as md
import nglview
import matplotlib.pyplot as plt


# temperature, mass, and related constants
temperature = 300 * unit.kelvin
kT = kB * temperature
mass = (12.0 * unit.dalton)
sigma_v = np.sqrt(kB * temperature / mass)

# openmm units
mass_unit = unit.dalton
distance_unit = unit.nanometer
time_unit = unit.femtosecond
energy_unit = unit.kilojoule_per_mole
speed_unit = distance_unit / time_unit
force_unit = unit.kilojoule_per_mole / unit.nanometer

# ANI-1 units and conversion factors
ani_distance_unit = unit.angstrom
hartree_to_kJ_mol = 2625.499638
ani_energy_unit = hartree_to_kJ_mol * unit.kilojoule_per_mole # simtk.unit doesn't have hartree?
nm_to_angstroms = (1.0 * distance_unit) / (1.0 * ani_distance_unit) 

# use torchani to sample methane
# TODO: modify this to be an instance method of a Molecule class instead of a function sitting in main...
import torchani
import torch
device = torch.device('cpu')
model = torchani.models.ANI1x()


def ANI1ccx_force(x):
    """
    Parameters
    ----------
    x : numpy array, unit'd
        coordinates
    
    Returns
    -------
    F : numpy array, unit'd
        force, with units of kJ/mol/nm attached
    """
    assert(type(x) == unit.Quantity)
    
    x_in_nm = x.value_in_unit(unit.nanometer)

    coordinates = torch.tensor([x_in_nm],
                               requires_grad=True, device=device, dtype=torch.float32)
    
    # convert from nm to angstroms
    coordinates_in_angstroms = coordinates * nm_to_angstroms
    _, energy_in_hartree = model((species, coordinates_in_angstroms))
    
    # convert energy from hartrees to kJ/mol
    energy_in_kJ_mol = energy_in_hartree * hartree_to_kJ_mol
    
    # derivative of E (in kJ/mol) w.r.t. coordinates (in nm)
    derivative = torch.autograd.grad((energy_in_kJ_mol).sum(), coordinates)[0]
    
    F_in_openmm_unit = - np.array(derivative)[0]
    return F_in_openmm_unit * (unit.kilojoule_per_mole / unit.nanometer)

  
def langevin(x0,
             force,
             n_steps=100,
             stepsize=1 * unit.femtosecond,
             collision_rate=10/unit.picoseconds,
             temperature=300 * unit.kelvin,
             progress_bar=False
            ):
    """Unadjusted Langevin dynamics.

    Parameters
    ----------
    x0 : array of floats, unit'd (distance unit)
        initial configuration
    force : callable, accepts a unit'd array and returns a unit'd array
        assumes input is in units of distance
        output is in units of energy / distance
    n_steps : integer
        number of Langevin steps
    stepsize : float > 0, in units of time
        finite timestep parameter
    collision_rate : float > 0, in units of 1/time
        controls the rate of interaction with the heat bath
    temperature : float > 0, in units of temperature
    progress_bar : bool
        use tqdm to show progress bar

    Returns
    -------
    traj : [n_steps + 1 x dim] array of floats, unit'd
        trajectory of samples generated by Langevin dynamics
        
    Note
    ----
        Assumes all particles have equal mass for convenience
        (Hardcoded to 12 daltons, ignores the global variable `mass`!)
    """
    
    # x,v are input using units
    v0 = np.random.randn(len(x0),3) * sigma_v

    
    assert(type(x0) == unit.Quantity)
    assert(type(v0) == unit.Quantity)
    assert(type(stepsize) == unit.Quantity)
    assert(type(collision_rate) == unit.Quantity)
    assert(type(temperature) == unit.Quantity)
    
    # convert initial state numpy arrays with correct attached units
    x = np.array(x0.value_in_unit(distance_unit)) * distance_unit
    v = np.array(v0.value_in_unit(speed_unit)) * speed_unit
    
    # traj is accumulated as a list of arrays with attached units
    traj = [x]
    
    # defining particles to have uniform masses
    mass = 12.0 * unit.dalton # TODO: allow non-uniform masses?

    # dimensionless scalars
    a = np.exp(- collision_rate * stepsize)
    b = np.sqrt(1 - np.exp(-2 * collision_rate * stepsize))
       
    # compute force on initial configuration
    F = force(x)
    trange = range(n_steps)
    trange = tqdm(trange)

    for _ in trange:
            
        # v
        v += (stepsize * 0.5) * F / mass
        # r
        x += (stepsize * 0.5) * v
        # o
        v = (a * v) + (b * sigma_v * np.random.randn(*x.shape))
        # r
        x += (stepsize * 0.5) * v

        F = force(x)
        # v
        v += (stepsize * 0.5) * F / mass
               
        norm_F = np.linalg.norm(F)
        
        # report gradient norm
        trange.set_postfix({'|force|': norm_F})
        
        # check positions and forces are finite
        if (not np.isfinite(x).all()) or (not np.isfinite(norm_F)):
            print("Numerical instability encountered!")
            return traj

        traj.append(x)
       
    return traj, x



# run langevin dynamics using ANI1x and methane
atom_list = 'CHHHH'
    
x = np.array([[0.03192167, 0.00638559, 0.01301679],
                             [-0.83140486, 0.39370209, -0.26395324],
                             [-0.66518241, -0.84461308, 0.20759389],
                             [0.45554739, 0.54289633, 0.81170881],
                             [0.66091919, -0.16799635, -0.91037834]])

x_in_angstroms = x * unit.angstrom
# initial conditions: coordinates from example were given in units of angstroms
species = model.species_to_tensor(atom_list).to(device).unsqueeze(0)
x0 = x_in_angstroms
    
trajectory, x_in_nm = langevin(x0, ANI1ccx_force, n_steps=50000, stepsize=2*unit.femtosecond, temperature=temperature, progress_bar=True)

# generating mdtraj traj object
pdb_file = open('methane.pdb', 'w+')
s = ('''
COMPND      METHANE
AUTHOR      DAVE WOODCOCK  95 12 18
ATOM      1  C           1       0.257  -0.363   0.000  1.00  0.00
ATOM      2  H           1       0.257   0.727   0.000  1.00  0.00
ATOM      3  H           1       0.771  -0.727   0.890  1.00  0.00
ATOM      4  H           1       0.771  -0.727  -0.890  1.00  0.00
ATOM      5  H           1      -0.771  -0.727   0.000  1.00  0.00 
TER       6              1
END
''')

pdb_file.write(s)
pdb_file.close()

topology = md.load('methane.pdb').topology
traj_in_nm = [(x /10)/ unit.angstrom for x in trajectory]

ani_traj = md.Trajectory(traj_in_nm, topology)
ani_traj = ani_traj.superpose(ani_traj[0]) # RMSD align onto first frame
view = nglview.show_mdtraj(ani_traj)

# compute angles

l = md.compute_angles(ani_traj, np.array([[1, 0, 2]], np.int32))
plt.hist(l[1::] * 180/np.pi,linewidth=0.5, bins=np.linspace(90, 130, 60), alpha=0.5)
print('Mean: {}'.format(np.mean(l)* 180/np.pi))

l = md.compute_angles(ani_traj, np.array([[1, 0, 3]], np.int32))
plt.hist(l[1::] * 180/np.pi,linewidth=0.5, bins=np.linspace(90, 130, 60), alpha=0.5)
print('Mean: {}'.format(np.mean(l)* 180/np.pi))

l = md.compute_angles(ani_traj, np.array([[1, 0, 4]], np.int32))
plt.hist(l[1::] * 180/np.pi,linewidth=0.5, bins=np.linspace(90, 130, 60), alpha=0.5)
print('Mean: {}'.format(np.mean(l)* 180/np.pi))

ax = plt.subplot(1,1,1)

plt.xlabel('iteration')
plt.ylabel('degree')
plt.title('H-C-H angle in degree')
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.savefig('angles_vs_iterations.png', dpi=300, bbox_inches='tight')
plt.show()