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, )