def test_run_in_place(graph):
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=10, n_steps_per_sample=10)
graph = simulation.run(graph, in_place=True)
assert graph.nodes["n1"].data["xyz"].shape == torch.Size([12, 10, 3])
def test_run(graph):
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=10, n_steps_per_sample=10)
samples = simulation.run(graph, in_place=False)
assert samples.shape == torch.Size([10, 12, 3])
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 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 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 test_coulomb_energy_consistency(g):
""" We use both `esp.mm` and OpenMM to compute the Coulomb energy of
some molecules with generated geometries and see if the resulting Columb
energy matches.
"""
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
print(g.mol)
# get simulation
esp_simulation = MoleculeVacuumSimulation(
n_samples=10,
n_steps_per_sample=10,
forcefield="gaff-1.81",
charge_method="gasteiger",
)
simulation = esp_simulation.simulation_from_graph(g)
charges = g.mol.partial_charges.flatten()
system = simulation.system
esp_simulation.run(g, in_place=True)
# if MD blows up, forget about it
if g.nodes["n1"].data["xyz"].abs().max() > 100:
return True
g.nodes['n1'].data['q'] = torch.tensor(charges).unsqueeze(-1)
esp.mm.nonbonded.multiply_charges(g.heterograph)
esp.mm.geometry.geometry_in_graph(g.heterograph)
esp.mm.energy.energy_in_graph(g.heterograph, terms=["nonbonded", "onefour"])
print(g.nodes['g'].data['u'].detach())
print(esp.data.md.get_coulomb_force(g)[0])
npt.assert_almost_equal(
g.nodes['g'].data['u'].detach().numpy(),
esp.data.md.get_coulomb_force(g)[0].numpy(),
decimal=3,
)
def test_apply(ds):
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
n_samples=1, n_steps_per_sample=1
).run
ds.apply(simulation, in_place=True)
assert ds.graphs[0].nodes["n1"].data["xyz"].shape[-1] == 3
assert ds.graphs[0].nodes["n1"].data["xyz"].shape[-2] == 1
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_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)
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,
)
simulation.run(g)
# In[77]:
g.heterograph.nodes['n1'].data['xyz'].mean(dim=1)
# In[78]:
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g.heterograph)
def test_energy_angle_and_bond(g):
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
# get simulation
esp_simulation = MoleculeVacuumSimulation(
n_samples=1,
n_steps_per_sample=1000,
forcefield="gaff-1.81",
charge_method="gasteiger",
)
simulation = esp_simulation.simulation_from_graph(g)
system = simulation.system
esp_simulation.run(g, in_place=True)
# if MD blows up, forget about it
if g.nodes["n1"].data["xyz"].abs().max() > 100:
return True
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(
simulation.topology,
system,
openmm.VerletIntegrator(0.0),
)
_simulation.context.setPositions(
g.nodes["n1"].data["xyz"][:, 0, :].detach().numpy() * unit.bohr)
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
# parametrize
ff = esp.graphs.legacy_force_field.LegacyForceField("gaff-1.81")
g = ff.parametrize(g)
# n2 : bond, n3: angle, n1: nonbonded?
# n1 : sigma (k), epsilon (eq), and charge (not included yet)
for term in ["n2", "n3"]:
g.nodes[term].data["k"] = g.nodes[term].data["k_ref"]
g.nodes[term].data["eq"] = g.nodes[term].data["eq_ref"]
"""
for term in ["n1"]:
g.nodes[term].data["sigma"] = g.nodes[term].data["sigma_ref"]
g.nodes[term].data["epsilon"] = g.nodes[term].data["epsilon_ref"]
# g.nodes[term].data['q'] = g.nodes[term].data['q_ref']
"""
for term in ["n4"]:
g.nodes[term].data["phases"] = g.nodes[term].data["phases_ref"]
g.nodes[term].data["periodicity"] = g.nodes[term].data[
"periodicity_ref"]
g.nodes[term].data["k"] = g.nodes[term].data["k_ref"]
# for each atom, store n_snapshots x 3
# g.nodes["n1"].data["xyz"] = torch.tensor(
# simulation.context.getState(getPositions=True)
# .getPositions(asNumpy=True)
# .value_in_unit(esp.units.DISTANCE_UNIT),
# dtype=torch.float32,
# )[None, :, :].permute(1, 0, 2)
# print(g.nodes['n2'].data)
esp.mm.geometry.geometry_in_graph(g.heterograph)
esp.mm.energy.energy_in_graph(g.heterograph, terms=["n2", "n3", "n4"])
# writes into nodes
# .data['u_nonbonded'], .data['u_onefour'], .data['u2'], .data['u3'],
# TODO: consider more carefully how many decimals of precision are needed
n_decimals = 3
# test bonds
npt.assert_almost_equal(
g.nodes["g"].data["u_n2"].detach().numpy(),
energies["HarmonicBondForce"],
decimal=n_decimals,
)
# test angles
npt.assert_almost_equal(
g.nodes["g"].data["u_n3"].detach().numpy(),
energies["HarmonicAngleForce"],
decimal=n_decimals,
)
def test_consistency():
import torch
import espaloma as esp
g = esp.Graph("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=10, n_steps_per_sample=10)
g = simulation.run(g, in_place=True)
g.nodes["n2"].data["coefficients"] = torch.randn(
g.heterograph.number_of_nodes("n2"), 2).exp()
g.nodes["n3"].data["coefficients"] = torch.randn(
g.heterograph.number_of_nodes("n3"), 2).exp()
esp.mm.geometry.geometry_in_graph(g.heterograph)
esp.mm.energy.energy_in_graph(g.heterograph, terms=["n2", "n3"])
u0_2 = g.nodes["n2"].data["u"] - g.nodes["n2"].data["u"].mean(
dim=1, keepdims=True)
u0_3 = g.nodes["n3"].data["u"] - g.nodes["n3"].data["u"].mean(
dim=1, keepdims=True)
u0 = g.nodes["g"].data["u"] - g.nodes["g"].data["u"].mean(dim=1,
keepdims=True)
(
g.nodes["n2"].data["k"],
g.nodes["n2"].data["eq"],
) = esp.mm.functional.linear_mixture_to_original(
g.nodes["n2"].data["coefficients"][:, 0][:, None],
g.nodes["n2"].data["coefficients"][:, 1][:, None],
1.5,
6.0,
)
import math
(
g.nodes["n3"].data["k"],
g.nodes["n3"].data["eq"],
) = esp.mm.functional.linear_mixture_to_original(
g.nodes["n3"].data["coefficients"][:, 0][:, None],
g.nodes["n3"].data["coefficients"][:, 1][:, None],
0.0,
math.pi,
)
g.nodes["n2"].data.pop("coefficients")
g.nodes["n3"].data.pop("coefficients")
esp.mm.energy.energy_in_graph(g.heterograph, terms=["n2", "n3"])
u1_2 = g.nodes["n2"].data["u"] - g.nodes["n2"].data["u"].mean(
dim=1, keepdims=True)
u1_3 = g.nodes["n3"].data["u"] - g.nodes["n3"].data["u"].mean(
dim=1, keepdims=True)
u1 = g.nodes["g"].data["u"] - g.nodes["g"].data["u"].mean(dim=1,
keepdims=True)
import numpy.testing as npt
npt.assert_almost_equal(
u0_2.detach().numpy(),
u1_2.detach().numpy(),
decimal=3,
)
npt.assert_almost_equal(
u0_3.detach().numpy(),
u1_3.detach().numpy(),
decimal=3,
)
npt.assert_almost_equal(
u0.detach().numpy(),
u1.detach().numpy(),
decimal=3,
)
def test_coulomb_energy_consistency(g):
""" We use both `esp.mm` and OpenMM to compute the Coulomb energy of
some molecules with generated geometries and see if the resulting Columb
energy matches.
"""
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
# get simulation
esp_simulation = MoleculeVacuumSimulation(
n_samples=1,
n_steps_per_sample=10,
forcefield="gaff-1.81",
charge_method="gasteiger",
)
simulation = esp_simulation.simulation_from_graph(g)
charges = g.mol.partial_charges.flatten()
system = simulation.system
esp_simulation.run(g, in_place=True)
# if MD blows up, forget about it
if g.nodes["n1"].data["xyz"].abs().max() > 100:
return True
_simulation = openmm.app.Simulation(
simulation.topology,
system,
openmm.VerletIntegrator(0.0),
)
forces = list(system.getForces())
for force in forces:
name = force.__class__.__name__
if "Nonbonded" in name:
force.setNonbondedMethod(openmm.NonbondedForce.NoCutoff)
force.updateParametersInContext(_simulation.context)
_simulation.context.setPositions(
g.nodes["n1"].data["xyz"][:, 0, :].detach().numpy() * unit.bohr)
state = _simulation.context.getState(
getEnergy=True,
getParameters=True,
)
energy_old = state.getPotentialEnergy().value_in_unit(
esp.units.ENERGY_UNIT)
forces = list(system.getForces())
print(forces)
for force in forces:
name = force.__class__.__name__
print(name)
if name == "NonbondedForce":
force.setNonbondedMethod(openmm.NonbondedForce.NoCutoff)
print(force.getNumExceptions())
for idx in range(force.getNumParticles()):
q, sigma, epsilon = force.getParticleParameters(idx)
force.setParticleParameters(idx, 0.0, sigma, epsilon)
for idx in range(force.getNumExceptions()):
idx0, idx1, q, sigma, epsilon = force.getExceptionParameters(
idx)
force.setExceptionParameters(idx, idx0, idx1, 0.0, sigma,
epsilon)
force.updateParametersInContext(_simulation.context)
state = _simulation.context.getState(
getEnergy=True,
getParameters=True,
)
energy_new = state.getPotentialEnergy().value_in_unit(
esp.units.ENERGY_UNIT)
g.nodes['n1'].data['q'] = torch.tensor(charges).unsqueeze(-1)
esp.mm.nonbonded.multiply_charges(g.heterograph)
esp.mm.geometry.geometry_in_graph(g.heterograph)
esp.mm.energy.energy_in_graph(g.heterograph,
terms=["nonbonded", "onefour"])
npt.assert_almost_equal(
g.nodes['g'].data['u'].item(),
energy_old - energy_new,
decimal=3,
)
def run(args):
# define data
data = getattr(esp.data, args.data)(first=1)
# get force field
forcefield = esp.graphs.legacy_force_field.LegacyForceField(
args.forcefield)
# param / typing
operation = forcefield.parametrize
# apply to dataset
data = data.apply(operation, in_place=True)
# apply simulation
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(n_samples=1000,
n_steps_per_sample=10)
data = data.apply(simulation.run, in_place=True)
# only one bit of data
ds = data.view("graph", batch_size=1)
ds_te = ds_tr = ds
for g in ds:
pass
# representation
representation = esp.nn.baselines.FreeParameterBaseline(g_ref=g)
net = torch.nn.Sequential(
representation,
esp.mm.geometry.GeometryInGraph(),
esp.mm.energy.EnergyInGraph(),
esp.mm.energy.EnergyInGraph(suffix='_ref'),
)
optimizer = torch.optim.LBFGS(net.parameters(),
0.01,
line_search_fn='strong_wolfe')
metrics_tr = [
esp.metrics.GraphMetric(base_metric=torch.nn.MSELoss(),
between=['u', 'u_ref'],
level='g')
]
metrics_te = [
esp.metrics.GraphMetric(base_metric=base_metric,
between=[param, param + '_ref'],
level=term) for param in ['u']
for term in ['g']
for base_metric in [esp.metrics.rmse, esp.metrics.r2]
]
exp = esp.TrainAndTest(
ds_tr=ds_tr,
ds_te=ds_te,
net=net,
metrics_tr=metrics_tr,
metrics_te=metrics_te,
n_epochs=args.n_epochs,
normalize=esp.data.normalize.PositiveNotNormalize,
optimizer=optimizer,
)
results = exp.run()
print(esp.app.report.markdown(results))
import os
os.mkdir(args.out)
with open(args.out + "/architecture.txt", "w") as f_handle:
f_handle.write(str(exp))
with open(args.out + "/result_table.md", "w") as f_handle:
f_handle.write(esp.app.report.markdown(results))
curves = esp.app.report.curve(results)
for spec, curve in curves.items():
np.save(args.out + "/" + "_".join(spec) + ".npy", curve)
import pickle
with open(args.out + "/ref_g_test.th", "wb") as f_handle:
pickle.dump(exp.ref_g_test, f_handle)
with open(args.out + "/ref_g_training.th", "wb") as f_handle:
pickle.dump(exp.ref_g_training, f_handle)
print(esp.app.report.markdown(results))
import pickle
with open(args.out + "/ref_g_test.th", "wb") as f_handle:
pickle.dump(exp.ref_g_test, f_handle)
with open(args.out + "/ref_g_training.th", "wb") as f_handle:
pickle.dump(exp.ref_g_training, f_handle)
def test_energy_angle_and_bond():
g = esp.Graph("C")
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
# get simulation
esp_simulation = MoleculeVacuumSimulation(
n_samples=1, n_steps_per_sample=10, forcefield="gaff-1.81"
)
simulation = esp_simulation.simulation_from_graph(g)
system = simulation.system
esp_simulation.run(g)
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)
# create new simulation
_simulation = openmm.app.Simulation(
simulation.topology,
system,
openmm.VerletIntegrator(0.0),
)
_simulation.context.setPositions(
g.nodes["n1"].data["xyz"][:, 0, :].detach().numpy() * 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
for idx, force in enumerate(forces):
name = force.__class__.__name__
if "HarmonicAngleForce" in name:
print("openmm thinks there are %s angles" % force.getNumAngles())
for _idx in range(force.getNumAngles()):
_, __, ___, eq, k = force.getAngleParameters(_idx)
eq = eq.value_in_unit(esp.units.ANGLE_UNIT)
k = k.value_in_unit(esp.units.ANGLE_FORCE_CONSTANT_UNIT)
print(eq, k)
# parametrize
ff = esp.graphs.legacy_force_field.LegacyForceField("gaff-1.81")
g = ff.parametrize(g)
# n2 : bond, n3: angle, n1: nonbonded?
# n1 : sigma (k), epsilon (eq), and charge (not included yet)
for term in ["n2", "n3"]:
g.nodes[term].data["k"] = g.nodes[term].data["k_ref"]
g.nodes[term].data["eq"] = g.nodes[term].data["eq_ref"]
print(
"espaloma thinks there are %s angles"
% g.heterograph.number_of_nodes("n3")
)
print(g.nodes["n3"].data["k"])
print(g.nodes["n3"].data["eq"])
# for each atom, store n_snapshots x 3
# g.nodes["n1"].data["xyz"] = torch.tensor(
# simulation.context.getState(getPositions=True)
# .getPositions(asNumpy=True)
# .value_in_unit(esp.units.DISTANCE_UNIT),
# dtype=torch.float32,
# )[None, :, :].permute(1, 0, 2)
# print(g.nodes['n2'].data)
esp.mm.geometry.geometry_in_graph(g.heterograph)
esp.mm.energy.energy_in_graph(g.heterograph, terms=["n2", "n3", "n4"])
n_decimals = 3
# test angles
npt.assert_almost_equal(
g.nodes["g"].data["u_n3"].detach().numpy(),
energies["HarmonicAngleForce"],
decimal=n_decimals,
)
def run(args):
# define data
data = getattr(esp.data, args.data)(first=args.first)
# get force field
forcefield = esp.graphs.legacy_force_field.LegacyForceField(
args.forcefield)
# param / typing
operation = forcefield.parametrize
# apply to dataset
data = data.apply(operation, in_place=True)
# apply simulation
# make simulation
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation(
n_samples=100,
n_steps_per_sample=10,
)
data = data.apply(simulation.run, in_place=True)
# split
partition = [int(x) for x in args.partition.split(":")]
ds_tr, ds_te = data.split(partition)
# batch
ds_tr = ds_tr.view("graph", batch_size=args.batch_size)
ds_te = ds_te.view("graph", batch_size=args.batch_size)
# layer
layer = esp.nn.layers.dgl_legacy.gn(args.layer)
# representation
representation = esp.nn.Sequential(layer, config=args.config)
# get the last bit of units
units = [int(x) for x in args.config if x.isdigit()][-1]
print(args.janossy_config)
janossy_config = []
for x in args.janossy_config:
if isinstance(x, int):
janossy_config.append(int(x))
elif x.isdigit():
janossy_config.append(int(x))
else:
janossy_config.append(x)
print(janossy_config)
readout = esp.nn.readout.janossy.JanossyPooling(
in_features=units,
config=janossy_config,
)
net = torch.nn.Sequential(
representation,
readout,
esp.mm.geometry.GeometryInGraph(),
esp.mm.energy.EnergyInGraph(),
esp.mm.energy.EnergyInGraph(suffix='_ref'),
)
'''
metrics_tr = [
esp.metrics.GraphMetric(
base_metric=torch.nn.L1Loss(),
between=['u', 'u_ref'],
level='g'
)
]
'''
metrics_tr = [
esp.metrics.GraphMetric(
base_metric=torch.nn.MSELoss(),
between=['u', "u_ref"],
level="g",
),
esp.metrics.GraphDerivativeMetric(
base_metric=torch.nn.MSELoss(),
between=["u", "u_ref"],
level="g",
weight=10.0,
),
]
metrics_te = [
esp.metrics.GraphMetric(
base_metric=esp.metrics.r2,
between=['u', 'u_ref'],
level="g",
),
esp.metrics.GraphMetric(
base_metric=esp.metrics.rmse,
between=['u', 'u_ref'],
level="g",
),
]
'''
metrics_te = [
esp.metrics.GraphMetric(
base_metric=base_metric,
between=[param, param + '_ref'],
level=term
) for param in ['u'] for term in ['g']
for base_metric in [
esp.metrics.rmse,
esp.metrics.r2
]
]
'''
exp = esp.TrainAndTest(
ds_tr=ds_tr,
ds_te=ds_te,
net=net,
metrics_tr=metrics_tr,
metrics_te=metrics_te,
n_epochs=args.n_epochs,
normalize=esp.data.normalize.NotNormalize,
optimizer=lambda net: torch.optim.Adam(net.parameters(), 1e-3),
device=torch.device('cuda:0'),
)
results = exp.run()
print(esp.app.report.markdown(results))
import os
os.mkdir(args.out)
with open(args.out + "/architecture.txt", "w") as f_handle:
f_handle.write(str(exp))
with open(args.out + "/result_table.md", "w") as f_handle:
f_handle.write(esp.app.report.markdown(results))
curves = esp.app.report.curve(results)
for spec, curve in curves.items():
np.save(args.out + "/" + "_".join(spec) + ".npy", curve)
import pickle
with open(args.out + "/ref_g_test.th", "wb") as f_handle:
pickle.dump(exp.ref_g_test, f_handle)
with open(args.out + "/ref_g_training.th", "wb") as f_handle:
pickle.dump(exp.ref_g_training, f_handle)
print(esp.app.report.markdown(results))
import pickle
with open(args.out + "/ref_g_test.th", "wb") as f_handle:
pickle.dump(exp.ref_g_test, f_handle)
with open(args.out + "/ref_g_training.th", "wb") as f_handle:
pickle.dump(exp.ref_g_training, f_handle)
def test_system(graph):
from espaloma.data.md import MoleculeVacuumSimulation
simulation = MoleculeVacuumSimulation()