def __init__(self, graphs=[], first=None):
super(GraphDataset, self).__init__()
from openforcefield.topology import Molecule
if all(
isinstance(graph, Molecule) or isinstance(graph, str)
for graph in graphs):
if first is None or first == -1:
graphs = [esp.Graph(graph) for graph in graphs]
else:
graphs = [esp.Graph(graph) for graph in graphs[:first]]
self.graphs = graphs
def breakdown_along_time_axis(g, batch_size=32):
n_snapshots = g.nodes["g"].data["u_ref"].flatten().shape[0]
idxs = list(range(n_snapshots))
from random import shuffle
shuffle(idxs)
chunks = [
idxs[_idx * batch_size:(_idx + 1) * batch_size]
for _idx in range(n_snapshots // batch_size)
]
_gs = []
for chunk in chunks:
_g = esp.Graph(g.mol)
_g.nodes["g"].data["u_ref"] = (
g.nodes["g"].data["u_ref"][:, chunk].detach().clone())
_g.nodes["n1"].data["xyz"] = (
g.nodes["n1"].data["xyz"][:, chunk, :].detach().clone())
_g.nodes["n1"].data["u_ref_prime"] = (
g.nodes["n1"].data["u_ref_prime"][:, chunk, :].detach().clone())
_g.nodes["n1"].data["xyz"].requires_grad = True
_gs.append(_g)
return _gs
def test_butane():
"""check that esp.graphs.deploy.openmm_system_from_graph runs without error on butane"""
ff = esp.graphs.legacy_force_field.LegacyForceField("smirnoff99Frosst")
g = esp.Graph("CCCC")
g = ff.parametrize(g)
print(g.mol)
esp.graphs.deploy.openmm_system_from_graph(g, suffix="_ref")
def run():
import os
paths = os.listdir('merged_data')
gs = [esp.Graph().load('merged_data/' + path) for path in paths]
ds = esp.data.dataset.GraphDataset(gs)
ds.save('ds.th')
def g():
g = esp.Graph(smiles)
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=n_samples,
n_steps_per_sample=1)
g = simulation.run(g, in_place=True)
return g
def get_graph(collection, record_name):
# get record and trajectory
record = collection.get_record(record_name, specification="default")
entry = collection.get_entry(record_name)
from openff.toolkit.topology import Molecule
mol = Molecule.from_qcschema(entry)
try:
trajectory = record.get_trajectory()
except:
return None
if trajectory is None:
return None
g = esp.Graph(mol)
# energy is already hartree
g.nodes["g"].data["u_ref"] = torch.tensor(
[
Quantity(
snapshot.properties.scf_total_energy,
esp.units.HARTREE_PER_PARTICLE,
).value_in_unit(esp.units.ENERGY_UNIT) for snapshot in trajectory
],
dtype=torch.get_default_dtype(),
)[None, :]
g.nodes["n1"].data["xyz"] = torch.tensor(
np.stack(
[
Quantity(
snapshot.get_molecule().geometry,
unit.bohr,
).value_in_unit(esp.units.DISTANCE_UNIT)
for snapshot in trajectory
],
axis=1,
),
requires_grad=True,
dtype=torch.get_default_dtype(),
)
g.nodes["n1"].data["u_ref_prime"] = torch.stack(
[
torch.tensor(
Quantity(
snapshot.dict()["return_result"],
esp.units.HARTREE_PER_PARTICLE / unit.bohr,
).value_in_unit(esp.units.FORCE_UNIT),
dtype=torch.get_default_dtype(),
) for snapshot in trajectory
],
dim=1,
)
return g
def _create_impropers_only_system(
smiles: str = "CC1=C(C(=O)C2=C(C1=O)N3CC4C(C3(C2COC(=O)N)OC)N4)N",
) -> mm.System:
"""Create a simulation that contains only improper torsion terms,
by parameterizing with openff-1.2.0 and deleting all terms but impropers
"""
molecule = Molecule.from_smiles(smiles, allow_undefined_stereo=True)
g = esp.Graph(molecule)
topology = Topology.from_molecules(molecule)
forcefield = ForceField("openff-1.2.0.offxml")
openmm_system = forcefield.create_openmm_system(topology)
# delete all forces except PeriodicTorsionForce
is_torsion = (
lambda force: "PeriodicTorsionForce" in force.__class__.__name__
)
for i in range(openmm_system.getNumForces())[::-1]:
if not is_torsion(openmm_system.getForce(i)):
openmm_system.removeForce(i)
assert openmm_system.getNumForces() == 1
torsion_force = openmm_system.getForce(0)
assert is_torsion(torsion_force)
# set k = 0 for any torsion that's not an improper
indices = set(
map(
tuple,
esp.graphs.utils.offmol_indices.improper_torsion_indices(
molecule
),
)
)
num_impropers_retained = 0
for i in range(torsion_force.getNumTorsions()):
(
p1,
p2,
p3,
p4,
periodicity,
phase,
k,
) = torsion_force.getTorsionParameters(i)
if (p1, p2, p3, p4) in indices:
num_impropers_retained += 1
else:
torsion_force.setTorsionParameters(
i, p1, p2, p3, p4, periodicity, phase, 0.0
)
assert (
num_impropers_retained > 0
) # otherwise this molecule is not a useful test case!
return openmm_system, topology, g
def test_gaff_parametrize():
ff = esp.graphs.legacy_force_field.LegacyForceField("gaff-1.81")
g = esp.Graph(
"CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
)
ff.parametrize(g)
print(g.nodes["n2"].data)
print(g.nodes["n3"].data)
print(g.nodes["n4"].data)
print(g.nodes["n4_improper"].data)
def test_save_and_load(graph):
import tempfile
import espaloma as esp
with tempfile.TemporaryDirectory() as tempdir:
graph.save(tempdir + "/g.esp")
new_graph = esp.Graph()
new_graph.load(tempdir + "/g.esp")
assert graph.homograph.number_of_nodes == graph.homograph.number_of_nodes
assert graph.homograph.number_of_edges == graph.homograph.number_of_edges
def h5_to_dataset(df):
def get_smiles(x):
try:
return x["offmol"].to_smiles()
except:
return np.nan
df["smiles"] = df.apply(get_smiles, axis=1)
df = df.dropna()
groups = df.groupby("smiles")
gs = []
for name, group in groups:
mol_ref = group["offmol"][0]
assert all(mol_ref == entry for entry in group["offmol"])
g = esp.Graph(mol_ref)
u_ref = np.concatenate(group["energies"].values)
u_ref_prime = np.concatenate(group["gradients"].values,
axis=0).transpose(1, 0, 2)
xyz = np.concatenate(group["xyz"].values, axis=0).transpose(1, 0, 2)
assert u_ref_prime.shape[0] == xyz.shape[0] == mol_ref.n_atoms
assert u_ref.shape[0] == u_ref_prime.shape[1] == xyz.shape[1]
# energy is already hartree
g.nodes["g"].data["u_ref"] = torch.tensor(
Quantity(u_ref, esp.units.HARTREE_PER_PARTICLE).value_in_unit(
esp.units.ENERGY_UNIT),
dtype=torch.get_default_dtype(),
)[None, :]
g.nodes["n1"].data["xyz"] = torch.tensor(
Quantity(
xyz,
unit.bohr,
).value_in_unit(esp.units.DISTANCE_UNIT),
requires_grad=True,
dtype=torch.get_default_dtype(),
)
g.nodes["n1"].data["u_ref_prime"] = torch.tensor(
Quantity(
u_ref_prime,
esp.units.HARTREE_PER_PARTICLE / unit.bohr,
).value_in_unit(esp.units.FORCE_UNIT),
dtype=torch.get_default_dtype(),
)
gs.append(g)
return esp.data.dataset.GraphDataset(gs)
def h5_to_dataset(df):
df['smiles'] = df.apply(lambda x: x['offmol'].to_smiles(), axis=1)
groups = df.groupby("smiles")
gs = []
for name, group in groups:
mol_ref = group['offmol'][0]
assert all(mol_ref == entry for entry in group['offmol'])
g = esp.Graph(mol_ref)
u_ref = np.concatenate(group['energies'].values)
u_ref_prime = np.concatenate(group['gradients'].values, axis=0).transpose(1, 0, 2)
xyz = np.concatenate(group['xyz'].values, axis=0).transpose(1, 0, 2)
assert u_ref_prime.shape[0] == xyz.shape[0] == mol_ref.n_atoms
assert u_ref.shape[0] == u_ref_prime.shape[1] == xyz.shape[1]
# energy is already hartree
g.nodes['g'].data['u_ref'] = torch.tensor(
Quantity(
u_ref,
esp.units.HARTREE_PER_PARTICLE
).value_in_unit(
esp.units.ENERGY_UNIT
),
dtype=torch.get_default_dtype(),
)[None, :]
g.nodes['n1'].data['xyz'] = torch.tensor(
Quantity(
xyz,
unit.bohr,
).value_in_unit(
esp.units.DISTANCE_UNIT
),
requires_grad=True,
dtype=torch.get_default_dtype(),
)
g.nodes['n1'].data['u_ref_prime'] = torch.tensor(
Quantity(
u_ref_prime,
esp.units.HARTREE_PER_PARTICLE / unit.bohr,
).value_in_unit(
esp.units.FORCE_UNIT
),
dtype=torch.get_default_dtype(),
)
gs.append(g)
return esp.data.dataset.GraphDataset(gs)
def all_g():
from espaloma.data.md import MoleculeVacuumSimulation
all_g = {}
for improper_def in expected_n_terms.keys():
g = esp.Graph(smiles)
if improper_def != 'none':
regenerate_impropers(g, improper_def)
simulation = MoleculeVacuumSimulation(n_samples=n_samples,
n_steps_per_sample=1)
g = simulation.run(g, in_place=True)
all_g[improper_def] = g
return all_g
def _type_gaff(self, g):
""" Type a molecular graph using gaff force fields.
"""
# assert the forcefield is indeed of gaff family
assert "gaff" in self.forcefield
# make sure mol is in OpenForceField format `
mol = g.mol
# import template generator
from openmmforcefields.generators import GAFFTemplateGenerator
gaff = GAFFTemplateGenerator(molecules=mol, forcefield=self.forcefield)
# create temporary directory for running antechamber
import os
import shutil
import tempfile
tempdir = tempfile.mkdtemp()
prefix = "molecule"
input_sdf_filename = os.path.join(tempdir, prefix + ".sdf")
gaff_mol2_filename = os.path.join(tempdir, prefix + ".gaff.mol2")
frcmod_filename = os.path.join(tempdir, prefix + ".frcmod")
# write sdf for input
mol.to_file(input_sdf_filename, file_format="sdf")
# run antechamber
gaff._run_antechamber(
molecule_filename=input_sdf_filename,
input_format="mdl",
gaff_mol2_filename=gaff_mol2_filename,
frcmod_filename=frcmod_filename,
)
gaff._read_gaff_atom_types_from_mol2(gaff_mol2_filename, mol)
gaff_types = [atom.gaff_type for atom in mol.atoms]
shutil.rmtree(tempdir)
# put types into graph object
if g is None:
g = esp.Graph(mol)
g.nodes["n1"].data["legacy_typing"] = torch.tensor(
[self._str_2_idx[atom] for atom in gaff_types])
return g
def test_save_and_load():
import espaloma as esp
g = esp.Graph("C")
ds = esp.data.dataset.GraphDataset([g])
# Temporary directory will be automatically cleaned up
from espaloma.data.utils import make_temp_directory
with make_temp_directory() as tmpdir:
import os
filename = os.path.join(tmpdir, "ds")
ds.save(filename)
new_ds = esp.data.dataset.GraphDataset.load(filename)
def breakdown_along_time_axis(g, batch_size=32):
n_snapshots = g.nodes['g'].data['u_ref'].flatten().shape[0]
idxs = list(range(n_snapshots))
from random import shuffle
shuffle(idxs)
chunks = [idxs[_idx * batch_size : (_idx + 1) * batch_size] for _idx in range(n_snapshots // batch_size)]
_gs = []
for chunk in chunks:
_g = esp.Graph(g.mol)
_g.nodes['g'].data['u_ref'] = g.nodes['g'].data['u_ref'][:, chunk]
_g.nodes['n1'].data['xyz'] = g.nodes['n1'].data['xyz'][:, chunk, :]
_g.nodes['n1'].data['u_ref_prime'] = g.nodes['n1'].data['u_ref_prime'][:, chunk, :]
_gs.append(_g)
return _gs
def baseline():
import espaloma as esp
g = esp.Graph("c1ccccc1")
# get force field
forcefield = esp.graphs.legacy_force_field.LegacyForceField(
"smirnoff99Frosst-1.1.0")
# param / typing
operation = forcefield.parametrize
operation(g)
baseline = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)
return baseline
def test_energy():
g = esp.Graph("c1ccccc1")
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=10, n_steps_per_sample=10)
g = simulation.run(g, in_place=True)
param = esp.graphs.legacy_force_field.LegacyForceField(
"smirnoff99Frosst-1.1.0"
).parametrize
g = param(g)
# parametrize
layer = esp.nn.dgl_legacy.gn()
net = torch.nn.Sequential(
esp.nn.Sequential(layer, [32, "tanh", 32, "tanh", 32, "tanh"]),
esp.nn.readout.janossy.JanossyPooling(
in_features=32,
config=[32, "tanh"],
out_features={
1: ["epsilon", "sigma"],
2: ["k", "eq"],
3: ["k", "eq"],
4: ["k"],
},
),
esp.nn.readout.janossy.JanossyPoolingImproper(
in_features=32,
config=[32, "tanh"],
out_features={
"k": 6,
},
),
)
g = net(g.heterograph)
# print(g.nodes['n2'].data)
esp.mm.geometry.geometry_in_graph(g)
# esp.mm.energy.energy_in_graph(g)
esp.mm.energy.energy_in_graph(g, terms=["n2", "n3", "n4", "n4_improper"])
def zinc(first=-1, *args, **kwargs):
""" ZINC collection.
..[1] Irwin, John J, and Brian K Shoichet.
“ZINC
--a free database of commercially available compounds for virtual screening.”
Journal of chemical information and modeling
vol. 45,1 (2005): 177-82. doi:10.1021/ci049714+
"""
import tarfile
from os.path import exists
from openff.toolkit.topology import Molecule
from rdkit import Chem
fname = "parm_at_Frosst.tgz"
url = "http://www.ccl.net/cca/data/parm_at_Frosst/parm_at_Frosst.tgz"
if not exists(fname):
import urllib.request
urllib.request.urlretrieve(url, fname)
archive = tarfile.open(fname)
zinc_file = archive.extractfile("parm_at_Frosst/zinc.sdf")
_mols = Chem.ForwardSDMolSupplier(zinc_file, removeHs=False)
count = 0
gs = []
for mol in _mols:
try:
gs.append(
esp.Graph(Molecule.from_rdkit(mol,
allow_undefined_stereo=True)))
count += 1
except:
pass
if first != -1 and count >= first:
break
return esp.data.dataset.GraphDataset(gs, *args, **kwargs)
def test_small_net():
import torch
import espaloma as esp
# define a layer
layer = esp.nn.layers.dgl_legacy.gn("GraphConv")
# define a representation
representation = esp.nn.Sequential(layer,
[32, "tanh", 32, "tanh", 32, "tanh"])
# define a readout
readout = esp.nn.readout.janossy.JanossyPooling(config=[32, "tanh"],
in_features=32)
net = torch.nn.Sequential(representation, readout)
g = esp.Graph("c1ccccc1")
def test_multiple_conformation():
import espaloma as esp
g = esp.Graph("c1ccccc1")
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=10, n_steps_per_sample=10)
g = simulation.run(g, in_place=True)
param = esp.graphs.legacy_force_field.LegacyForceField(
"smirnoff99Frosst-1.1.0"
).parametrize
g = param(g)
esp.mm.geometry.geometry_in_graph(g.heterograph)
esp.mm.energy.energy_in_graph(g.heterograph, suffix="_ref")
def test_parameter_consistent_caffeine():
ff = esp.graphs.legacy_force_field.LegacyForceField("openff-1.2.0")
g = esp.Graph("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
g = ff.parametrize(g)
system = esp.graphs.deploy.openmm_system_from_graph(g, suffix="_ref")
forces = list(system.getForces())
openff_forces = ff.FF.label_molecules(g.mol.to_topology())[0]
for idx, force in enumerate(forces):
force.setForceGroup(idx)
name = force.__class__.__name__
if "HarmonicBondForce" in name:
for _idx in range(force.getNumBonds()):
start, end, eq, k_openmm = force.getBondParameters(_idx)
k_openff = openff_forces["Bonds"][(start, end)].k
npt.assert_almost_equal(
k_openmm / k_openff,
2.0,
decimal=3,
)
def zinc(first=-1, *args, **kwargs):
import tarfile
from os.path import exists
from openforcefield.topology import Molecule
from rdkit import Chem
fname = 'parm_at_Frosst.tgz'
url = 'http://www.ccl.net/cca/data/parm_at_Frosst/parm_at_Frosst.tgz'
if not exists(fname):
import urllib.request
urllib.request.urlretrieve(url, fname)
archive = tarfile.open(fname)
zinc_file = archive.extractfile('parm_at_Frosst/zinc.sdf')
_mols = Chem.ForwardSDMolSupplier(zinc_file, removeHs=False)
count = 0
gs = []
for mol in _mols:
try:
gs.append(
esp.Graph(
Molecule.from_rdkit(mol, allow_undefined_stereo=True)
)
)
count += 1
except:
pass
if first != -1 and count >= first:
break
return esp.data.dataset.GraphDataset(gs, *args, **kwargs)
def run():
_dict = get_dict()
for record_idx, mol_idxs in _dict.items():
import os
paths = ['data/%s.th' % mol_idx for mol_idx in mol_idxs]
paths = [path for path in paths if os.path.exists(path)]
if len(paths) == 0:
continue
else:
print(record_idx)
import espaloma as esp
gs = [esp.Graph().load(path) for path in paths]
g_ref = gs[0]
g_ref.nodes['g'].data['u_ref'] = torch.cat(
[g.nodes['g'].data['u_ref'] for g in gs],
dim=1,
)
g_ref.nodes['n1'].data['xyz'] = torch.cat(
[g.nodes['n1'].data['xyz'] for g in gs],
dim=1,
)
g_ref.nodes['n1'].data['u_ref_prime'] = torch.cat(
[g.nodes['n1'].data['u_ref_prime'] for g in gs],
dim=1,
)
g_ref.save('merged_data/%s.th' % record_idx)
def test_caffeine():
ff = esp.graphs.legacy_force_field.LegacyForceField("openff-1.2.0")
g = esp.Graph("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
g = ff.parametrize(g)
esp.graphs.deploy.openmm_system_from_graph(g, suffix="_ref")
def test_energy():
g = esp.Graph("c1ccccc1")
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=10, n_steps_per_sample=10)
g = simulation.run(g, in_place=True)
param = esp.graphs.legacy_force_field.LegacyForceField(
"gaff-1.81").parametrize
g = param(g)
# parametrize
# layer
layer = esp.nn.layers.dgl_legacy.gn()
# representation
representation = esp.nn.Sequential(
layer, config=[32, "relu", 32, "relu", 32, "relu"])
# get the last bit of units
units = 32
janossy_config = [32, "relu"]
readout = esp.nn.readout.janossy.JanossyPooling(
in_features=units,
config=janossy_config,
out_features={
2: {
"log_coefficients": 2
},
3: {
"log_coefficients": 2,
"coefficients_urey_bradley": 2,
"k_bond_bond": 1,
"k_bond_angle": 1,
"k_bond_angle": 1,
},
4: {
"k": 6,
"k_angle_angle": 1,
"k_angle_angle_torsion": 1,
"k_angle_torsion": 1,
"k_side_torsion": 1,
"k_center_torsion": 1,
},
},
)
readout_improper = esp.nn.readout.janossy.JanossyPoolingImproper(
in_features=units, config=janossy_config)
class ExpCoeff(torch.nn.Module):
def forward(self, g):
g.nodes["n2"].data["coefficients"] = (
g.nodes["n2"].data["log_coefficients"].exp())
g.nodes["n3"].data["coefficients"] = (
g.nodes["n3"].data["log_coefficients"].exp())
return g
class CarryII(torch.nn.Module):
def forward(self, g):
import math
g.multi_update_all(
{
"n2_as_0_in_n3": (
dgl.function.copy_src("u", "m_u_0"),
dgl.function.sum("m_u_0", "u_left"),
),
"n2_as_1_in_n3": (
dgl.function.copy_src("u", "m_u_1"),
dgl.function.sum("m_u_1", "u_right"),
),
"n2_as_0_in_n4": (
dgl.function.copy_src("u", "m_u_0"),
dgl.function.sum("m_u_0", "u_bond_left"),
),
"n2_as_1_in_n4": (
dgl.function.copy_src("u", "m_u_1"),
dgl.function.sum("m_u_1", "u_bond_center"),
),
"n2_as_2_in_n4": (
dgl.function.copy_src("u", "m_u_2"),
dgl.function.sum("m_u_2", "u_bond_right"),
),
"n3_as_0_in_n4": (
dgl.function.copy_src("u", "m3_u_0"),
dgl.function.sum("m3_u_0", "u_angle_left"),
),
"n3_as_1_in_n4": (
dgl.function.copy_src("u", "m3_u_1"),
dgl.function.sum("m3_u_1", "u_angle_right"),
),
},
cross_reducer="sum",
)
return g
net = torch.nn.Sequential(
representation,
readout,
readout_improper,
ExpCoeff(),
esp.mm.geometry.GeometryInGraph(),
esp.mm.energy.EnergyInGraph(terms=["n2", "n3", "n4", "n4_improper"]),
CarryII(),
esp.mm.energy.EnergyInGraphII(),
)
torch.nn.init.normal_(
net[1].f_out_2_to_log_coefficients.bias,
mean=-5,
)
torch.nn.init.normal_(
net[1].f_out_3_to_log_coefficients.bias,
mean=-5,
)
for name, module in net[1].named_modules():
if "k" in name:
torch.nn.init.normal(module.bias, mean=0.0, std=1e-4)
torch.nn.init.normal(module.weight, mean=0.0, std=1e-4)
g = net(g.heterograph)
print(g.nodes["n3"].data)
print(g.nodes["n4"].data)
# print(g.nodes['n2'].data)
esp.mm.geometry.geometry_in_graph(g)
esp.mm.energy.energy_in_graph(g)
# In[0]
import dgl
import numpy as np
import torch
import espaloma as esp
# In[74]:
g = esp.Graph('C')
# In[75]:
forcefield = esp.graphs.legacy_force_field.LegacyForceField(
"smirnoff99Frosst"
)
forcefield.parametrize(g)
# In[76]:
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
n_samples=100,
n_steps_per_sample=10,
def test_graph():
import espaloma as esp
g = esp.Graph("c1ccccc1")
print(g.heterograph)
def test_butane():
"""check that esp.graphs.deploy.openmm_system_from_graph runs without error on butane"""
ff = esp.graphs.legacy_force_field.LegacyForceField("openff-1.2.0")
g = esp.Graph("CCCC")
g = ff.parametrize(g)
esp.graphs.deploy.openmm_system_from_graph(g, suffix="_ref")
def graph():
import espaloma as esp
return esp.Graph("c1ccccc1")
def test_energy_consistent_caffeine():
"""Deploy a caffeine molecule parametrized by a traditional force field
and deployed by espaloma, make sure the energies computed using espaloma
and OpenMM are same or close.
"""
# grab a force field
ff = esp.graphs.legacy_force_field.LegacyForceField("openff-1.2.0")
# parametrize caffeine molecule using the parametrization
## Should there be a second test for SMIRNOFF impropers?
g = esp.Graph("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
g = ff.parametrize(g)
system = esp.graphs.deploy.openmm_system_from_graph(g, suffix="_ref")
# compute energies using espaloma
import torch
g.nodes["n1"].data["xyz"] = torch.randn(
g.heterograph.number_of_nodes("n1"), 1, 3)
esp.mm.geometry.geometry_in_graph(g.heterograph)
esp.mm.energy.energy_in_graph(g.heterograph,
terms=["n2", "n3", "n4", "n4_improper"],
suffix="_ref")
# compute energies using OpenMM with bond, angle, and torsion breakdown
forces = list(system.getForces())
energies = {}
for idx, force in enumerate(forces):
force.setForceGroup(idx)
name = force.__class__.__name__
if "Nonbonded" in name:
force.setNonbondedMethod(openmm.NonbondedForce.NoCutoff)
# epsilons = {}
# sigmas = {}
# for _idx in range(force.getNumParticles()):
# q, sigma, epsilon = force.getParticleParameters(_idx)
# # record parameters
# epsilons[_idx] = epsilon
# sigmas[_idx] = sigma
# force.setParticleParameters(_idx, 0., sigma, epsilon)
# def sigma_combining_rule(sig1, sig2):
# return (sig1 + sig2) / 2
# def eps_combining_rule(eps1, eps2):
# return np.sqrt(np.abs(eps1 * eps2))
# for _idx in range(force.getNumExceptions()):
# idx0, idx1, q, sigma, epsilon = force.getExceptionParameters(
# _idx)
# force.setExceptionParameters(
# _idx,
# idx0,
# idx1,
# 0.0,
# sigma_combining_rule(sigmas[idx0], sigmas[idx1]),
# eps_combining_rule(epsilons[idx0], epsilons[idx1])
# )
# force.updateParametersInContext(_simulation.context)
# create new simulation
_simulation = openmm.app.Simulation(
g.mol.to_topology().to_openmm(),
system,
openmm.VerletIntegrator(0.0),
)
_simulation.context.setPositions(
Quantity(
g.nodes["n1"].data["xyz"][:, 0, :].numpy(),
unit=esp.units.DISTANCE_UNIT,
).value_in_unit(unit.nanometer))
for idx, force in enumerate(forces):
name = force.__class__.__name__
state = _simulation.context.getState(
getEnergy=True,
getParameters=True,
groups=2**idx,
)
energy = state.getPotentialEnergy().value_in_unit(
esp.units.ENERGY_UNIT)
energies[name] = energy
# test if bond energies are equal
npt.assert_almost_equal(
g.nodes["g"].data["u_n2_ref"].numpy(),
energies["HarmonicBondForce"],
decimal=3,
)
# test if angle energies are equal
npt.assert_almost_equal(
g.nodes["g"].data["u_n3_ref"].numpy(),
energies["HarmonicAngleForce"],
decimal=3,
)
# test if torsion energies are equal
npt.assert_almost_equal(
g.nodes["g"].data["u_n4_ref"].numpy() +
g.nodes["g"].data["u_n4_improper_ref"].numpy(),
energies["PeriodicTorsionForce"],
decimal=3,
)
