def test_gbsa_handler():
patterns = [['[*:1]', 99., 99.], ['[#1:1]', 99., 99.],
['[#1:1]~[#7]', 99., 99.], ['[#6:1]', 0.1, 0.2],
['[#7:1]', 0.3, 0.4], ['[#8:1]', 0.5, 0.6],
['[#9:1]', 0.7, 0.8], ['[#14:1]', 99., 99.],
['[#15:1]', 99., 99.], ['[#16:1]', 99., 99.],
['[#17:1]', 99., 99.]]
props = {
'solvent_dielectric': 78.3, # matches OBC2,
'solute_dielectric': 1.0,
'probe_radius': 0.14,
'surface_tension': 28.3919551,
'dielectric_offset': 0.009,
# GBOBC1
'alpha': 0.8,
'beta': 0.0,
'gamma': 2.909125
}
smirks = [x[0] for x in patterns]
params = np.array([[x[1], x[2]] for x in patterns])
gbh = nonbonded.GBSAHandler(smirks, params, props)
obj = gbh.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
assert len(all_handlers) == 1
new_gbh = all_handlers[0]
np.testing.assert_equal(new_gbh.smirks, gbh.smirks)
np.testing.assert_equal(new_gbh.params, gbh.params)
assert new_gbh.props == gbh.props
def test_am1ccc():
patterns = [['[#6X4:1]-[#1:2]', 0.46323257920556493],
['[#6X3$(*=[#8,#16]):1]-[#6a:2]', 0.24281402370571598],
['[#6X3$(*=[#8,#16]):1]-[#8X1,#8X2:2]', 1.0620166764992722],
[
'[#6X3$(*=[#8,#16]):1]=[#8X1$(*=[#6X3]-[#8X2]):2]',
2.227759732057297
], ['[#6X3$(*=[#8,#16]):1]=[#8X1,#8X2:2]', 2.8182928673804217],
['[#6a:1]-[#8X1,#8X2:2]', 0.5315976926761063],
['[#6a:1]-[#1:2]', 0.0], ['[#6a:1]:[#6a:2]', 0.0],
['[#6a:1]:[#6a:2]', 0.0],
['[#8X1,#8X2:1]-[#1:2]', -2.3692047944101415],
['[#16:1]-[#8:2]', 99.]]
smirks = [x[0] for x in patterns]
params = np.array([x[1] * np.sqrt(138.935456) for x in patterns])
props = None
am1h = nonbonded.AM1CCCHandler(smirks, params, props)
obj = am1h.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
assert len(all_handlers) == 1
new_am1h = all_handlers[0]
np.testing.assert_equal(new_am1h.smirks, am1h.smirks)
np.testing.assert_equal(new_am1h.params, am1h.params)
assert new_am1h.props == am1h.props
def test_combine_recipe():
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_ccc.py').read())
aspirin = Chem.AddHs(Chem.MolFromSmiles("CC(=O)OC1=CC=CC=C1C(=O)O"))
AllChem.EmbedMolecule(aspirin)
ligand_recipe = md.Recipe.from_rdkit(aspirin, ff_handlers)
fname = 'tests/data/hif2a_nowater_min.pdb'
pdb = open(fname, 'r').read()
openmm_system, openmm_conf, _, _, _, _ = builders.build_protein_system(
'tests/data/hif2a_nowater_min.pdb')
protein_recipe = md.Recipe.from_openmm(openmm_system)
for left_recipe, right_recipe in [[protein_recipe, ligand_recipe],
[ligand_recipe, protein_recipe]]:
combined_recipe = left_recipe.combine(right_recipe)
qlj = np.ones((aspirin.GetNumAtoms() + openmm_conf.shape[0], 3))
left_nonbonded_potential = left_recipe.bound_potentials[-1]
right_nonbonded_potential = right_recipe.bound_potentials[-1]
combined_nonbonded_potential = combined_recipe.bound_potentials[-1]
left_idxs = left_nonbonded_potential.get_exclusion_idxs()
right_idxs = right_nonbonded_potential.get_exclusion_idxs()
combined_idxs = combined_nonbonded_potential.get_exclusion_idxs()
n_left = len(left_recipe.masses)
n_right = len(right_recipe.masses)
np.testing.assert_array_equal(
np.concatenate([left_idxs, right_idxs + n_left]), combined_idxs)
for bp in combined_recipe.bound_potentials:
bp.bound_impl(precision=np.float32)
def test_proper_torsion():
# proper torsions have a variadic number of terms
patterns = [
['[*:1]-[#6X3:2]=[#6X3:3]-[*:4]', [[99., 99., 99.]]],
['[*:1]-[#6X3:2]=[#6X3:3]-[#35:4]', [[99., 99., 99.]]],
['[#9:1]-[#6X3:2]=[#6X3:3]-[#35:4]', [[1., 2., 3.], [4., 5., 6.]]],
[
'[#35:1]-[#6X3:2]=[#6X3:3]-[#35:4]',
[[7., 8., 9.], [1., 3., 5.], [4., 4., 4.]]
],
['[#9:1]-[#6X3:2]=[#6X3:3]-[#9:4]', [[7., 8., 9.]]],
]
smirks = [x[0] for x in patterns]
params = [x[1] for x in patterns]
props = None
ph = bonded.ProperTorsionHandler(smirks, params, None)
obj = ph.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
assert len(all_handlers) == 1
new_ph = all_handlers[0]
np.testing.assert_equal(new_ph.smirks, ph.smirks)
np.testing.assert_equal(new_ph.params, ph.params)
assert new_ph.props == ph.props
def test_improper_torsion():
patterns = [[
'[*:1]~[#6X3:2](~[*:3])~[*:4]', 1.5341333333333333, 3.141592653589793,
2.0
], ['[*:1]~[#6X3:2](~[#8X1:3])~[#8:4]', 99., 99., 99.],
[
'[*:1]~[#7X3$(*~[#15,#16](!-[*])):2](~[*:3])~[*:4]', 99.,
99., 99.
],
[
'[*:1]~[#7X3$(*~[#6X3]):2](~[*:3])~[*:4]',
1.3946666666666667, 3.141592653589793, 2.0
], ['[*:1]~[#7X3$(*~[#7X2]):2](~[*:3])~[*:4]', 99., 99., 99.],
[
'[*:1]~[#7X3$(*@1-[*]=,:[*][*]=,:[*]@1):2](~[*:3])~[*:4]',
99., 99., 99.
], ['[*:1]~[#6X3:2](=[#7X2,#7X3+1:3])~[#7:4]', 99., 99., 99.]]
smirks = [x[0] for x in patterns]
params = np.array([[x[1], x[2], x[3]] for x in patterns])
imph = bonded.ImproperTorsionHandler(smirks, params, None)
obj = imph.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
assert len(all_handlers) == 1
new_imph = all_handlers[0]
np.testing.assert_equal(new_imph.smirks, imph.smirks)
np.testing.assert_equal(new_imph.params, imph.params)
assert new_imph.props == imph.props
def test_recipe_from_rdkit():
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_ccc.py').read())
suppl = Chem.SDMolSupplier('tests/data/ligands_40.sdf', removeHs=False)
for mol_idx, mol in enumerate(suppl):
print(mol_idx, Chem.MolToSmiles(mol))
system = md.Recipe.from_rdkit(mol, ff_handlers)
if mol_idx > 2:
break
def test_simple_charge_handler():
patterns = [
['[#1:1]', 99.],
['[#1:1]-[#6X4]', 99.],
['[#1:1]-[#6X4]-[#7,#8,#9,#16,#17,#35]', 99.],
['[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]', 99.],
[
'[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])(-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]',
99.
],
['[#1:1]-[#6X4]~[*+1,*+2]', 99.],
['[#1:1]-[#6X3]', 99.],
['[#1:1]-[#6X3]~[#7,#8,#9,#16,#17,#35]', 99.],
['[#1:1]-[#6X3](~[#7,#8,#9,#16,#17,#35])~[#7,#8,#9,#16,#17,#35]', 99.],
['[#1:1]-[#6X2]', 99.],
['[#1:1]-[#7]', 99.],
['[#1:1]-[#8]', 99.],
['[#1:1]-[#16]', 99.],
['[#6:1]', 0.7],
['[#6X2:1]', 99.],
['[#6X4:1]', 0.1],
['[#8:1]', 99.],
['[#8X2H0+0:1]', 0.5],
['[#8X2H1+0:1]', 99.],
['[#7:1]', 0.3],
['[#16:1]', 99.],
['[#15:1]', 99.],
['[#9:1]', 1.0],
['[#17:1]', 99.],
['[#35:1]', 99.],
['[#53:1]', 99.],
['[#3+1:1]', 99.],
['[#11+1:1]', 99.],
['[#19+1:1]', 99.],
['[#37+1:1]', 99.],
['[#55+1:1]', 99.],
['[#9X0-1:1]', 99.],
['[#17X0-1:1]', 99.],
['[#35X0-1:1]', 99.],
['[#53X0-1:1]', 99.],
]
smirks = [x[0] for x in patterns]
params = np.array([x[1] for x in patterns])
props = None
sch = nonbonded.SimpleChargeHandler(smirks, params, props)
obj = sch.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
assert len(all_handlers) == 1
new_sch = all_handlers[0]
np.testing.assert_equal(new_sch.smirks, sch.smirks)
np.testing.assert_equal(new_sch.params, sch.params)
assert new_sch.props == sch.props
def test_am1bcc():
smirks = []
params = []
props = None
am1 = nonbonded.AM1BCCHandler(smirks, params, props)
obj = am1.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
am1 = all_handlers[0]
np.testing.assert_equal(am1.smirks, am1.smirks)
np.testing.assert_equal(am1.params, am1.params)
assert am1.props == am1.props
def setUp(self, *args, **kwargs):
suppl = Chem.SDMolSupplier('tests/data/benzene_phenol_sparse.sdf',
removeHs=False)
all_mols = [x for x in suppl]
self.mol_a = all_mols[0]
self.mol_b = all_mols[1]
# atom type free
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_recharge.py').read())
self.ff = Forcefield(ff_handlers)
super(BenzenePhenolSparseTest, self).__init__(*args, **kwargs)
def test_bad_factor(self):
# test a bad mapping that results in a non-cancellable endpoint
suppl = Chem.SDMolSupplier('tests/data/ligands_40.sdf', removeHs=False)
all_mols = [x for x in suppl]
mol_a = all_mols[0]
mol_b = all_mols[1]
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_recharge.py').read())
ff = Forcefield(ff_handlers)
core = np.array([[4, 1], [5, 2], [6, 3], [7, 4], [8, 5], [9,
6], [10, 7],
[11, 8], [12, 9], [13, 10], [15, 11], [16, 12],
[18, 14], [34, 31], [17, 13], [23, 23], [33, 30],
[32, 28], [31, 27], [30, 26], [19, 15], [20, 16],
[21, 17]])
with self.assertRaises(topology.AtomMappingError):
st = topology.SingleTopology(mol_a, mol_b, core, ff)
def test_am1_differences():
ff_raw = open("ff/params/smirnoff_1_1_0_ccc.py").read()
ff_handlers = deserialize_handlers(ff_raw)
for ccc in ff_handlers:
if isinstance(ccc, nonbonded.AM1CCCHandler):
break
suppl = Chem.SDMolSupplier('tests/data/ligands_40.sdf', removeHs=False)
smi = "[H]c1c(OP(=S)(OC([H])([H])C([H])([H])[H])OC([H])([H])C([H])([H])[H])nc(C([H])(C([H])([H])[H])C([H])([H])[H])nc1C([H])([H])[H]"
smi = "Clc1c(Cl)c(Cl)c(-c2c(Cl)c(Cl)c(Cl)c(Cl)c2Cl)c(Cl)c1Cl"
mol = Chem.MolFromSmiles(smi)
mol = Chem.AddHs(mol)
mol.SetProp("_Name", "Debug")
assert AllChem.EmbedMolecule(mol) == 0
suppl = [mol]
am1 = nonbonded.AM1Handler([], [], None)
bcc = nonbonded.AM1BCCHandler([], [], None)
for mol in suppl:
am1_params = am1.parameterize(mol)
ccc_params = ccc.parameterize(mol)
bcc_params = bcc.parameterize(mol)
if np.sum(np.abs(ccc_params - bcc_params)) > 0.1:
print(mol.GetProp("_Name"), Chem.MolToSmiles(mol))
print(" AM1 CCC BCC S ?")
for atom_idx, atom in enumerate(mol.GetAtoms()):
a = am1_params[atom_idx]
b = bcc_params[atom_idx]
c = ccc_params[atom_idx]
print("{:6.2f}".format(a), "{:6.2f}".format(c), "{:6.2f}".format(b), atom.GetSymbol(), end="")
if np.abs(b-c) > 0.1:
print(" *")
else:
print(" ")
assert 0
def test_good_factor(self):
# test a good mapping
suppl = Chem.SDMolSupplier('tests/data/ligands_40.sdf', removeHs=False)
all_mols = [x for x in suppl]
mol_a = all_mols[1]
mol_b = all_mols[4]
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_recharge.py').read())
ff = Forcefield(ff_handlers)
core = np.array([[0, 0], [2, 2], [1, 1], [6, 6], [5, 5], [4, 4],
[3, 3], [15, 16], [16, 17], [17, 18], [18, 19],
[19, 20], [20, 21], [32, 30], [26, 25], [27, 26],
[7, 7], [8, 8], [9, 9], [10, 10], [29, 11], [11, 12],
[12, 13], [14, 15], [31, 29], [13, 14], [23, 24],
[30, 28], [28, 27], [21, 22]])
st = topology.SingleTopology(mol_a, mol_b, core, ff)
# test that the vjps work
_ = jax.vjp(st.parameterize_harmonic_bond,
ff.hb_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_harmonic_angle,
ff.ha_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_proper_torsion,
ff.pt_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_improper_torsion,
ff.it_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_nonbonded,
ff.q_handle.params,
ff.lj_handle.params,
has_aux=True)
def test_harmonic_bond():
patterns = [['[#6X4:1]-[#6X4:2]', 0.1,
0.2], ['[#6X4:1]-[#6X3:2]', 99., 99.],
['[#6X4:1]-[#6X3:2]=[#8X1+0]', 99., 99.],
['[#6X3:1]-[#6X3:2]', 99.,
99.], ['[#6X3:1]:[#6X3:2]', 99., 99.],
['[#6X3:1]=[#6X3:2]', 99., 99.], ['[#6:1]-[#7:2]', 0.1, 0.2],
['[#6X3:1]-[#7X3:2]', 99., 99.],
['[#6X4:1]-[#7X3:2]-[#6X3]=[#8X1+0]', 99., 99.],
['[#6X3:1](=[#8X1+0])-[#7X3:2]', 99., 99.],
['[#6X3:1]-[#7X2:2]', 99., 99.],
['[#6X3:1]:[#7X2,#7X3+1:2]', 99., 99.],
['[#6X3:1]=[#7X2,#7X3+1:2]', 99., 99.],
['[#6:1]-[#8:2]', 99., 99.], ['[#6X3:1]-[#8X1-1:2]', 99., 99.],
['[#6X4:1]-[#8X2H0:2]', 0.3, 0.4],
['[#6X3:1]-[#8X2:2]', 99., 99.],
['[#6X3:1]-[#8X2H1:2]', 99., 99.],
['[#6X3a:1]-[#8X2H0:2]', 99., 99.],
['[#6X3:1](=[#8X1])-[#8X2H0:2]', 99., 99.],
['[#6:1]=[#8X1+0,#8X2+1:2]', 99., 99.],
['[#6X3:1](~[#8X1])~[#8X1:2]', 99., 99.],
['[#6X3:1]~[#8X2+1:2]~[#6X3]', 99., 99.],
['[#6X2:1]-[#6:2]', 99., 99.], ['[#6X2:1]-[#6X4:2]', 99., 99.],
['[#6X2:1]=[#6X3:2]', 99., 99.], ['[#6:1]#[#7:2]', 99., 99.],
['[#6X2:1]#[#6X2:2]', 99.,
99.], ['[#6X2:1]-[#8X2:2]', 99., 99.],
['[#6X2:1]-[#7:2]', 99., 99.], ['[#6X2:1]=[#7:2]', 99., 99.],
['[#16:1]=[#6:2]', 99., 99.], ['[#6X2:1]=[#16:2]', 99., 99.],
['[#7:1]-[#7:2]', 99., 99.], ['[#7X3:1]-[#7X2:2]', 99., 99.],
['[#7X2:1]-[#7X2:2]', 99., 99.], ['[#7:1]:[#7:2]', 99., 99.],
['[#7:1]=[#7:2]', 99., 99.], ['[#7+1:1]=[#7-1:2]', 99., 99.],
['[#7:1]#[#7:2]', 99., 99.], ['[#7:1]-[#8X2:2]', 99., 99.],
['[#7:1]~[#8X1:2]', 99., 99.], ['[#8X2:1]-[#8X2:2]', 99., 99.],
['[#16:1]-[#6:2]', 99., 99.], ['[#16:1]-[#1:2]', 99., 99.],
['[#16:1]-[#16:2]', 99., 99.], ['[#16:1]-[#9:2]', 99., 99.],
['[#16:1]-[#17:2]', 99., 99.], ['[#16:1]-[#35:2]', 99., 99.],
['[#16:1]-[#53:2]', 99., 99.],
['[#16X2,#16X1-1,#16X3+1:1]-[#6X4:2]', 99., 99.],
['[#16X2,#16X1-1,#16X3+1:1]-[#6X3:2]', 99., 99.],
['[#16X2:1]-[#7:2]', 99., 99.],
['[#16X2:1]-[#8X2:2]', 99., 99.],
['[#16X2:1]=[#8X1,#7X2:2]', 99., 99.],
['[#16X4,#16X3!+1:1]-[#6:2]', 99., 99.],
['[#16X4,#16X3:1]~[#7:2]', 99., 99.],
['[#16X4,#16X3:1]-[#8X2:2]', 99., 99.],
['[#16X4,#16X3:1]~[#8X1:2]', 99., 99.],
['[#15:1]-[#1:2]', 99., 99.], ['[#15:1]~[#6:2]', 99., 99.],
['[#15:1]-[#7:2]', 99., 99.], ['[#15:1]=[#7:2]', 99., 99.],
['[#15:1]~[#8X2:2]', 99., 99.], ['[#15:1]~[#8X1:2]', 99., 99.],
['[#16:1]-[#15:2]', 99., 99.], ['[#15:1]=[#16X1:2]', 99., 99.],
['[#6:1]-[#9:2]', 99., 99.], ['[#6X4:1]-[#9:2]', 0.6, 0.7],
['[#6:1]-[#17:2]', 99., 99.], ['[#6X4:1]-[#17:2]', 99., 99.],
['[#6:1]-[#35:2]', 99., 99.], ['[#6X4:1]-[#35:2]', 99., 99.],
['[#6:1]-[#53:2]', 99., 99.], ['[#6X4:1]-[#53:2]', 99., 99.],
['[#7:1]-[#9:2]', 99., 99.], ['[#7:1]-[#17:2]', 99., 99.],
['[#7:1]-[#35:2]', 99., 99.], ['[#7:1]-[#53:2]', 99., 99.],
['[#15:1]-[#9:2]', 99., 99.], ['[#15:1]-[#17:2]', 99., 99.],
['[#15:1]-[#35:2]', 99., 99.], ['[#15:1]-[#53:2]', 99., 99.],
['[#6X4:1]-[#1:2]', 99., 99.], ['[#6X3:1]-[#1:2]', 99., 99.],
['[#6X2:1]-[#1:2]', 99., 99.], ['[#7:1]-[#1:2]', 99., 99.],
['[#8:1]-[#1:2]', 99., 99.1]]
smirks = [x[0] for x in patterns]
params = np.array([[x[1], x[2]] for x in patterns])
props = None
hbh = bonded.HarmonicBondHandler(smirks, params, None)
obj = hbh.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
assert len(all_handlers) == 1
new_hbh = all_handlers[0]
np.testing.assert_equal(new_hbh.smirks, hbh.smirks)
np.testing.assert_equal(new_hbh.params, hbh.params)
assert new_hbh.props == hbh.props
def main(args, stage):
# benzene = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1")) # a
# phenol = Chem.AddHs(Chem.MolFromSmiles("Oc1ccccc1")) # b
#01234567890
benzene = Chem.AddHs(Chem.MolFromSmiles("C1=CC=C2C=CC=CC2=C1")) # a
phenol = Chem.AddHs(Chem.MolFromSmiles("C1=CC=C2C=CC=CC2=C1")) # b
AllChem.EmbedMolecule(benzene)
AllChem.EmbedMolecule(phenol)
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_ccc.py').read())
r_benzene = Recipe.from_rdkit(benzene, ff_handlers)
r_phenol = Recipe.from_rdkit(phenol, ff_handlers)
r_combined = r_benzene.combine(r_phenol)
core_pairs = np.array(
[
[0, 0],
[1, 1],
[2, 2],
[3, 3],
[4, 4],
[5, 5],
[6, 6],
[7, 7],
[8, 8],
[9, 9],
# [10,10]
],
dtype=np.int32)
core_pairs[:, 1] += benzene.GetNumAtoms()
a_idxs = np.arange(benzene.GetNumAtoms())
b_idxs = np.arange(phenol.GetNumAtoms()) + benzene.GetNumAtoms()
core_k = 20.0
if stage == 0:
centroid_k = 200.0
rbfe.stage_0(r_combined, b_idxs, core_pairs, centroid_k, core_k)
# lambda_schedule = np.linspace(0.0, 1.0, 2)
# lambda_schedule = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
lambda_schedule = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
elif stage == 1:
rbfe.stage_1(r_combined, a_idxs, b_idxs, core_pairs, core_k)
lambda_schedule = np.linspace(0.0, 1.2, 60)
else:
assert 0
system, host_coords, box, topology = builders.build_water_system(4.0)
r_host = Recipe.from_openmm(system)
r_final = r_host.combine(r_combined)
# minimize coordinates of host + ligand A
ha_coords = np.concatenate([host_coords, get_romol_conf(benzene)])
pool = Pool(args.num_gpus)
# we need to run this in a subprocess since the cuda runtime
# must not be initialized in the master thread due to lack of
# fork safety
r_minimize = minimize_setup(r_host, r_benzene)
ha_coords = pool.map(
minimize,
[(r_minimize.bound_potentials, r_minimize.masses, ha_coords, box)],
chunksize=1)
# this is a list
ha_coords = ha_coords[0]
pool.close()
pool = Pool(args.num_gpus)
x0 = np.concatenate([ha_coords, get_romol_conf(phenol)])
masses = np.concatenate([r_host.masses, r_benzene.masses, r_phenol.masses])
seed = np.random.randint(np.iinfo(np.int32).max)
intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, masses, seed)
# production run at various values of lambda
for epoch in range(10):
avg_du_dls = []
run_args = []
for lamb_idx, lamb in enumerate(lambda_schedule):
run_args.append(
(lamb, intg, r_final.bound_potentials, r_final.masses, x0, box,
lamb_idx % args.num_gpus, stage))
avg_du_dls = pool.map(run, run_args, chunksize=1)
print("stage", stage, "epoch", epoch, "dG",
np.trapz(avg_du_dls, lambda_schedule))
def test_lennard_jones_handler():
patterns = [
['[#1:1]', 99., 999.],
['[#1:1]-[#6X4]', 99., 999.],
['[#1:1]-[#6X4]-[#7,#8,#9,#16,#17,#35]', 99., 999.],
[
'[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]',
99., 999.
],
[
'[#1:1]-[#6X4](-[#7,#8,#9,#16,#17,#35])(-[#7,#8,#9,#16,#17,#35])-[#7,#8,#9,#16,#17,#35]',
99., 999.
],
['[#1:1]-[#6X4]~[*+1,*+2]', 99., 999.],
['[#1:1]-[#6X3]', 99., 999.],
['[#1:1]-[#6X3]~[#7,#8,#9,#16,#17,#35]', 99., 999.],
[
'[#1:1]-[#6X3](~[#7,#8,#9,#16,#17,#35])~[#7,#8,#9,#16,#17,#35]',
99., 999.
],
['[#1:1]-[#6X2]', 99., 999.],
['[#1:1]-[#7]', 99., 999.],
['[#1:1]-[#8]', 99., 999.],
['[#1:1]-[#16]', 99., 999.],
['[#6:1]', 0.7, 0.8],
['[#6X2:1]', 99., 999.],
['[#6X4:1]', 0.1, 0.2],
['[#8:1]', 99., 999.],
['[#8X2H0+0:1]', 0.5, 0.6],
['[#8X2H1+0:1]', 99., 999.],
['[#7:1]', 0.3, 0.4],
['[#16:1]', 99., 999.],
['[#15:1]', 99., 999.],
['[#9:1]', 1.0, 1.1],
['[#17:1]', 99., 999.],
['[#35:1]', 99., 999.],
['[#53:1]', 99., 999.],
['[#3+1:1]', 99., 999.],
['[#11+1:1]', 99., 999.],
['[#19+1:1]', 99., 999.],
['[#37+1:1]', 99., 999.],
['[#55+1:1]', 99., 999.],
['[#9X0-1:1]', 99., 999.],
['[#17X0-1:1]', 99., 999.],
['[#35X0-1:1]', 99., 999.],
['[#53X0-1:1]', 99., 999.],
]
smirks = [x[0] for x in patterns]
params = np.array([[x[1], x[2]] for x in patterns])
props = None
ljh = nonbonded.LennardJonesHandler(smirks, params, props)
obj = ljh.serialize()
all_handlers = deserialize_handlers(bin_to_str(obj))
ljh = all_handlers[0]
np.testing.assert_equal(ljh.smirks, ljh.smirks)
np.testing.assert_equal(ljh.params, ljh.params)
assert ljh.props == ljh.props
for address in worker_address_list:
print("connecting to", address)
channel = grpc.insecure_channel(address,
options = [
('grpc.max_send_message_length', 500 * 1024 * 1024),
('grpc.max_receive_message_length', 500 * 1024 * 1024)
]
)
stub = service_pb2_grpc.WorkerStub(channel)
stubs.append(stub)
ff_raw = open(forcefield, "r").read()
ff_handlers = deserialize_handlers(ff_raw)
box_width = 3.0
host_system, host_coords, box, _ = water_box.prep_system(box_width)
lambda_schedule = np.array([float(x) for x in general_cfg['lambda_schedule'].split(',')])
num_steps = int(general_cfg['n_steps'])
for epoch in range(100):
print("Starting Epoch", epoch, datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S"))
epoch_dir = os.path.join(general_cfg["out_dir"], "epoch_"+str(epoch))
if not os.path.exists(epoch_dir):
def pose_dock(
guests_sdfile,
host_pdbfile,
transition_type,
n_steps,
transition_steps,
max_lambda,
outdir,
random_rotation=False,
constant_atoms=[],
):
"""Runs short simulations in which the guests phase in or out over time
Parameters
----------
guests_sdfile: path to input sdf with guests to pose/dock
host_pdbfile: path to host pdb file to dock into
transition_type: "insertion" or "deletion"
n_steps: how many total steps of simulation to do (recommended: <= 1000)
transition_steps: how many steps to insert/delete the guest over (recommended: <= 500)
(must be <= n_steps)
max_lambda: lambda value the guest should insert from or delete to
(recommended: 1.0 for work calulation, 0.25 to stay close to original pose)
(must be =1 for work calculation to be applicable)
outdir: where to write output (will be created if it does not already exist)
random_rotation: whether to apply a random rotation to each guest before inserting
constant_atoms: atom numbers from the host_pdbfile to hold mostly fixed across the simulation
(1-indexed, like PDB files)
Output
------
A pdb & sdf file every 100 steps (outdir/<guest_name>_<step>.pdb)
stdout every 100 steps noting the step number, lambda value, and energy
stdout for each guest noting the work of transition
stdout for each guest noting how long it took to run
Note
----
If any norm of force per atom exceeds 20000 kJ/(mol*nm) [MAX_NORM_FORCE defined in docking/report.py],
the simulation for that guest will stop and the work will not be calculated.
"""
assert transition_steps <= n_steps
assert transition_type in ("insertion", "deletion")
if random_rotation:
assert transition_type == "insertion"
if not os.path.exists(outdir):
os.makedirs(outdir)
host_mol = Chem.MolFromPDBFile(host_pdbfile, removeHs=False)
amber_ff = app.ForceField("amber99sbildn.xml", "tip3p.xml")
host_file = PDBFile(host_pdbfile)
host_system = amber_ff.createSystem(
host_file.topology,
nonbondedMethod=app.NoCutoff,
constraints=None,
rigidWater=False,
)
host_conf = []
for x, y, z in host_file.positions:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
final_potentials = []
host_potentials, host_masses = openmm_deserializer.deserialize_system(
host_system, cutoff=1.2)
host_nb_bp = None
for bp in host_potentials:
if isinstance(bp, potentials.Nonbonded):
# (ytz): hack to ensure we only have one nonbonded term
assert host_nb_bp is None
host_nb_bp = bp
else:
final_potentials.append(bp)
# TODO (ytz): we should really fix this later on. This padding was done to
# address the particles that are too close to the boundary.
padding = 0.1
box_lengths = np.amax(host_conf, axis=0) - np.amin(host_conf, axis=0)
box_lengths = box_lengths + padding
box = np.eye(3, dtype=np.float64) * box_lengths
suppl = Chem.SDMolSupplier(guests_sdfile, removeHs=False)
for guest_mol in suppl:
start_time = time.time()
guest_name = guest_mol.GetProp("_Name")
guest_ff_handlers = deserialize_handlers(
open(
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"..",
"ff/params/smirnoff_1_1_0_ccc.py",
)).read())
ff = Forcefield(guest_ff_handlers)
guest_base_topology = topology.BaseTopology(guest_mol, ff)
# combine
hgt = topology.HostGuestTopology(host_nb_bp, guest_base_topology)
# setup the parameter handlers for the ligand
bonded_tuples = [[hgt.parameterize_harmonic_bond, ff.hb_handle],
[hgt.parameterize_harmonic_angle, ff.ha_handle],
[hgt.parameterize_proper_torsion, ff.pt_handle],
[hgt.parameterize_improper_torsion, ff.it_handle]]
these_potentials = list(final_potentials)
# instantiate the vjps while parameterizing (forward pass)
for fn, handle in bonded_tuples:
params, potential = fn(handle.params)
these_potentials.append(potential.bind(params))
nb_params, nb_potential = hgt.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
these_potentials.append(nb_potential.bind(nb_params))
bps = these_potentials
guest_masses = [a.GetMass() for a in guest_mol.GetAtoms()]
masses = np.concatenate([host_masses, guest_masses])
for atom_num in constant_atoms:
masses[atom_num - 1] += 50000
conformer = guest_mol.GetConformer(0)
mol_conf = np.array(conformer.GetPositions(), dtype=np.float64)
mol_conf = mol_conf / 10 # convert to md_units
if random_rotation:
center = np.mean(mol_conf, axis=0)
mol_conf -= center
from scipy.stats import special_ortho_group
mol_conf = np.matmul(mol_conf, special_ortho_group.rvs(3))
mol_conf += center
x0 = np.concatenate([host_conf, mol_conf]) # combined geometry
v0 = np.zeros_like(x0)
seed = 2021
intg = LangevinIntegrator(300, 1.5e-3, 1.0, masses, seed).impl()
impls = []
precision = np.float32
for b in bps:
p_impl = b.bound_impl(precision)
impls.append(p_impl)
ctxt = custom_ops.Context(x0, v0, box, intg, impls)
# collect a du_dl calculation once every other step
subsample_freq = 2
du_dl_obs = custom_ops.FullPartialUPartialLambda(impls, subsample_freq)
ctxt.add_observable(du_dl_obs)
if transition_type == "insertion":
new_lambda_schedule = np.concatenate([
np.linspace(max_lambda, 0.0, transition_steps),
np.zeros(n_steps - transition_steps),
])
elif transition_type == "deletion":
new_lambda_schedule = np.concatenate([
np.linspace(0.0, max_lambda, transition_steps),
np.ones(n_steps - transition_steps) * max_lambda,
])
else:
raise (RuntimeError(
'invalid `transition_type` (must be one of ["insertion", "deletion"])'
))
calc_work = True
for step, lamb in enumerate(new_lambda_schedule):
ctxt.step(lamb)
if step % 100 == 0:
report.report_step(ctxt, step, lamb, box, bps, impls,
guest_name, n_steps, 'pose_dock')
host_coords = ctxt.get_x_t()[:len(host_conf)] * 10
guest_coords = ctxt.get_x_t()[len(host_conf):] * 10
report.write_frame(host_coords, host_mol, guest_coords,
guest_mol, guest_name, outdir, step, 'pd')
if step in (0, int(n_steps / 2), n_steps - 1):
if report.too_much_force(ctxt, lamb, box, bps, impls):
calc_work = False
break
# Note: this condition only applies for ABFE, not RBFE
if (abs(du_dl_obs.full_du_dl()[0]) > 0.001
or abs(du_dl_obs.full_du_dl()[-1]) > 0.001):
print("Error: du_dl endpoints are not ~0")
calc_work = False
if calc_work:
work = np.trapz(du_dl_obs.full_du_dl(),
new_lambda_schedule[::subsample_freq])
print(f"guest_name: {guest_name}\twork: {work:.2f}")
end_time = time.time()
print(f"{guest_name} took {(end_time - start_time):.2f} seconds")
type=int,
help="number of absolute lambda windows",
required=True
)
cmd_args = parser.parse_args()
multiprocessing.set_start_method('spawn') # CUDA runtime is not forkable
pool = multiprocessing.Pool(cmd_args.num_gpus)
suppl = Chem.SDMolSupplier('tests/data/benzene_fluorinated.sdf', removeHs=False)
all_mols = [x for x in suppl]
mol_a = all_mols[0]
mol_b = all_mols[1]
ff_handlers = deserialize_handlers(open('ff/params/smirnoff_1_1_0_ccc.py').read())
ff = Forcefield(ff_handlers)
# the water system first.
solvent_system, solvent_coords, solvent_box, omm_topology = builders.build_water_system(4.0)
solvent_box += np.eye(3)*0.1 # BFGS this later
print("Minimizing the host structure to remove clashes.")
minimized_solvent_coords = minimizer.minimize_host_4d(mol_a, solvent_system, solvent_coords, ff, solvent_box)
absolute_lambda_schedule = np.concatenate([
np.linspace(0.0, 0.333, cmd_args.num_absolute_windows - cmd_args.num_absolute_windows//3, endpoint=False),
np.linspace(0.333, 1.0, cmd_args.num_absolute_windows//3),
])
abs_dGs = []
def calculate_rigorous_work(
host_pdbfile, guests_sdfile, outdir, fewer_outfiles=False, no_outfiles=False
):
"""
"""
if not os.path.exists(outdir):
os.makedirs(outdir)
print(
f"""
HOST_PDBFILE = {host_pdbfile}
GUESTS_SDFILE = {guests_sdfile}
OUTDIR = {outdir}
INSERTION_MAX_LAMBDA = {INSERTION_MAX_LAMBDA}
DELETION_MAX_LAMBDA = {DELETION_MAX_LAMBDA}
MIN_LAMBDA = {MIN_LAMBDA}
TRANSITION_STEPS = {TRANSITION_STEPS}
EQ1_STEPS = {EQ1_STEPS}
EQ2_STEPS = {EQ2_STEPS}
"""
)
# Prepare host
# TODO: handle extra (non-transitioning) guests?
print("Solvating host...")
(
solvated_host_system,
solvated_host_coords,
_,
_,
host_box,
solvated_topology,
) = builders.build_protein_system(host_pdbfile)
# sometimes water boxes are sad. Should be minimized first; this is a workaround
host_box += np.eye(3) * 0.1
print("host box", host_box)
solvated_host_pdb = os.path.join(outdir, "solvated_host.pdb")
writer = pdb_writer.PDBWriter([solvated_topology], solvated_host_pdb)
writer.write_frame(solvated_host_coords)
writer.close()
solvated_host_mol = Chem.MolFromPDBFile(solvated_host_pdb, removeHs=False)
if no_outfiles:
os.remove(solvated_host_pdb)
final_host_potentials = []
host_potentials, host_masses = openmm_deserializer.deserialize_system(solvated_host_system, cutoff=1.2)
host_nb_bp = None
for bp in host_potentials:
if isinstance(bp, potentials.Nonbonded):
# (ytz): hack to ensure we only have one nonbonded term
assert host_nb_bp is None
host_nb_bp = bp
else:
final_host_potentials.append(bp)
# Prepare water box
print("Generating water box...")
# TODO: water box probably doesn't need to be this big
box_lengths = host_box[np.diag_indices(3)]
water_box_width = min(box_lengths)
(
water_system,
orig_water_coords,
water_box,
water_topology,
) = builders.build_water_system(water_box_width)
# sometimes water boxes are sad. should be minimized first; this is a workaround
water_box += np.eye(3) * 0.1
print("water box", water_box)
# it's okay if the water box here and the solvated protein box don't align -- they have PBCs
water_pdb = os.path.join(outdir, "water_box.pdb")
writer = pdb_writer.PDBWriter([water_topology], water_pdb)
writer.write_frame(orig_water_coords)
writer.close()
water_mol = Chem.MolFromPDBFile(water_pdb, removeHs=False)
if no_outfiles:
os.remove(water_pdb)
final_water_potentials = []
water_potentials, water_masses = openmm_deserializer.deserialize_system(water_system, cutoff=1.2)
water_nb_bp = None
for bp in water_potentials:
if isinstance(bp, potentials.Nonbonded):
# (ytz): hack to ensure we only have one nonbonded term
assert water_nb_bp is None
water_nb_bp = bp
else:
final_water_potentials.append(bp)
# Run the procedure
print("Getting guests...")
suppl = Chem.SDMolSupplier(guests_sdfile, removeHs=False)
for guest_mol in suppl:
start_time = time.time()
guest_name = guest_mol.GetProp("_Name")
guest_conformer = guest_mol.GetConformer(0)
orig_guest_coords = np.array(guest_conformer.GetPositions(), dtype=np.float64)
orig_guest_coords = orig_guest_coords / 10 # convert to md_units
guest_ff_handlers = deserialize_handlers(
open(
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"..",
"ff/params/smirnoff_1_1_0_ccc.py",
)
).read()
)
ff = Forcefield(guest_ff_handlers)
guest_base_top = topology.BaseTopology(guest_mol, ff)
# combine host & guest
hgt = topology.HostGuestTopology(host_nb_bp, guest_base_top)
# setup the parameter handlers for the ligand
bonded_tuples = [
[hgt.parameterize_harmonic_bond, ff.hb_handle],
[hgt.parameterize_harmonic_angle, ff.ha_handle],
[hgt.parameterize_proper_torsion, ff.pt_handle],
[hgt.parameterize_improper_torsion, ff.it_handle]
]
combined_bps = list(final_host_potentials)
# instantiate the vjps while parameterizing (forward pass)
for fn, handle in bonded_tuples:
params, potential = fn(handle.params)
combined_bps.append(potential.bind(params))
nb_params, nb_potential = hgt.parameterize_nonbonded(ff.q_handle.params, ff.lj_handle.params)
combined_bps.append(nb_potential.bind(nb_params))
guest_masses = [a.GetMass() for a in guest_mol.GetAtoms()]
combined_masses = np.concatenate([host_masses, guest_masses])
run_leg(
solvated_host_coords,
orig_guest_coords,
combined_bps,
combined_masses,
host_box,
guest_name,
"host",
solvated_host_mol,
guest_mol,
outdir,
fewer_outfiles,
no_outfiles,
)
end_time = time.time()
print(
f"{guest_name} host leg time:", "%.2f" % (end_time - start_time), "seconds"
)
# combine water & guest
wgt = topology.HostGuestTopology(water_nb_bp, guest_base_top)
# setup the parameter handlers for the ligand
bonded_tuples = [
[wgt.parameterize_harmonic_bond, ff.hb_handle],
[wgt.parameterize_harmonic_angle, ff.ha_handle],
[wgt.parameterize_proper_torsion, ff.pt_handle],
[wgt.parameterize_improper_torsion, ff.it_handle]
]
combined_bps = list(final_water_potentials)
# instantiate the vjps while parameterizing (forward pass)
for fn, handle in bonded_tuples:
params, potential = fn(handle.params)
combined_bps.append(potential.bind(params))
nb_params, nb_potential = wgt.parameterize_nonbonded(ff.q_handle.params, ff.lj_handle.params)
combined_bps.append(nb_potential.bind(nb_params))
guest_masses = [a.GetMass() for a in guest_mol.GetAtoms()]
combined_masses = np.concatenate([water_masses, guest_masses])
start_time = time.time()
run_leg(
orig_water_coords,
orig_guest_coords,
combined_bps,
combined_masses,
water_box,
guest_name,
"water",
water_mol,
guest_mol,
outdir,
fewer_outfiles,
no_outfiles,
)
end_time = time.time()
print(
f"{guest_name} water leg time:", "%.2f" % (end_time - start_time), "seconds"
)
def dock_and_equilibrate(host_pdbfile,
guests_sdfile,
max_lambda,
insertion_steps,
eq_steps,
outdir,
fewer_outfiles=False,
constant_atoms=[]):
"""Solvates a host, inserts guest(s) into solvated host, equilibrates
Parameters
----------
host_pdbfile: path to host pdb file to dock into
guests_sdfile: path to input sdf with guests to pose/dock
max_lambda: lambda value the guest should insert from or delete to
(recommended: 1.0 for work calulation, 0.25 to stay close to original pose)
(must be =1 for work calculation to be applicable)
insertion_steps: how many steps to insert the guest over (recommended: 501)
eq_steps: how many steps of equilibration to do after insertion (recommended: 15001)
outdir: where to write output (will be created if it does not already exist)
fewer_outfiles: if True, will only write frames for the equilibration, not insertion
constant_atoms: atom numbers from the host_pdbfile to hold mostly fixed across the simulation
(1-indexed, like PDB files)
Output
------
A pdb & sdf file every 100 steps of insertion (outdir/<guest_name>/<guest_name>_<step>.[pdb/sdf])
A pdb & sdf file every 1000 steps of equilibration (outdir/<guest_name>/<guest_name>_<step>.[pdb/sdf])
stdout every 100(0) steps noting the step number, lambda value, and energy
stdout for each guest noting the work of transition
stdout for each guest noting how long it took to run
Note
----
If any norm of force per atom exceeds 20000 kJ/(mol*nm) [MAX_NORM_FORCE defined in docking/report.py],
the simulation for that guest will stop and the work will not be calculated.
"""
if not os.path.exists(outdir):
os.makedirs(outdir)
print(f"""
HOST_PDBFILE = {host_pdbfile}
GUESTS_SDFILE = {guests_sdfile}
OUTDIR = {outdir}
MAX_LAMBDA = {max_lambda}
INSERTION_STEPS = {insertion_steps}
EQ_STEPS = {eq_steps}
""")
# Prepare host
# TODO: handle extra (non-transitioning) guests?
print("Solvating host...")
# TODO: return topology from builders.build_protein_system
(
solvated_host_system,
solvated_host_coords,
_,
_,
host_box,
solvated_topology,
) = builders.build_protein_system(host_pdbfile)
# sometimes water boxes are sad. Should be minimized first; this is a workaround
host_box += np.eye(3) * 0.1
print("host box", host_box)
solvated_host_pdb = os.path.join(outdir, "solvated_host.pdb")
writer = pdb_writer.PDBWriter([solvated_topology], solvated_host_pdb)
writer.write_frame(solvated_host_coords)
writer.close()
solvated_host_mol = Chem.MolFromPDBFile(solvated_host_pdb, removeHs=False)
os.remove(solvated_host_pdb)
final_host_potentials = []
host_potentials, host_masses = openmm_deserializer.deserialize_system(
solvated_host_system, cutoff=1.2)
host_nb_bp = None
for bp in host_potentials:
if isinstance(bp, potentials.Nonbonded):
# (ytz): hack to ensure we only have one nonbonded term
assert host_nb_bp is None
host_nb_bp = bp
else:
final_host_potentials.append(bp)
# Run the procedure
print("Getting guests...")
suppl = Chem.SDMolSupplier(guests_sdfile, removeHs=False)
for guest_mol in suppl:
start_time = time.time()
guest_name = guest_mol.GetProp("_Name")
guest_conformer = guest_mol.GetConformer(0)
orig_guest_coords = np.array(guest_conformer.GetPositions(),
dtype=np.float64)
orig_guest_coords = orig_guest_coords / 10 # convert to md_units
guest_ff_handlers = deserialize_handlers(
open(
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"..",
"ff/params/smirnoff_1_1_0_ccc.py",
)).read())
ff = Forcefield(guest_ff_handlers)
guest_base_top = topology.BaseTopology(guest_mol, ff)
# combine host & guest
hgt = topology.HostGuestTopology(host_nb_bp, guest_base_top)
# setup the parameter handlers for the ligand
bonded_tuples = [[hgt.parameterize_harmonic_bond, ff.hb_handle],
[hgt.parameterize_harmonic_angle, ff.ha_handle],
[hgt.parameterize_proper_torsion, ff.pt_handle],
[hgt.parameterize_improper_torsion, ff.it_handle]]
combined_bps = list(final_host_potentials)
# instantiate the vjps while parameterizing (forward pass)
for fn, handle in bonded_tuples:
params, potential = fn(handle.params)
combined_bps.append(potential.bind(params))
nb_params, nb_potential = hgt.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
combined_bps.append(nb_potential.bind(nb_params))
guest_masses = [a.GetMass() for a in guest_mol.GetAtoms()]
combined_masses = np.concatenate([host_masses, guest_masses])
x0 = np.concatenate([solvated_host_coords, orig_guest_coords])
v0 = np.zeros_like(x0)
print(
f"SYSTEM",
f"guest_name: {guest_name}",
f"num_atoms: {len(x0)}",
)
for atom_num in constant_atoms:
combined_masses[atom_num - 1] += 50000
seed = 2021
intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, combined_masses,
seed).impl()
u_impls = []
for bp in combined_bps:
bp_impl = bp.bound_impl(precision=np.float32)
u_impls.append(bp_impl)
ctxt = custom_ops.Context(x0, v0, host_box, intg, u_impls)
# collect a du_dl calculation once every other step
subsample_freq = 2
du_dl_obs = custom_ops.FullPartialUPartialLambda(
u_impls, subsample_freq)
ctxt.add_observable(du_dl_obs)
# insert guest
insertion_lambda_schedule = np.linspace(max_lambda, 0.0,
insertion_steps)
calc_work = True
for step, lamb in enumerate(insertion_lambda_schedule):
ctxt.step(lamb)
if step % 100 == 0:
report.report_step(ctxt, step, lamb, host_box, combined_bps,
u_impls, guest_name, insertion_steps,
"INSERTION")
if not fewer_outfiles:
host_coords = ctxt.get_x_t()[:len(solvated_host_coords
)] * 10
guest_coords = ctxt.get_x_t()[len(solvated_host_coords
):] * 10
report.write_frame(
host_coords,
solvated_host_mol,
guest_coords,
guest_mol,
guest_name,
outdir,
str(step).zfill(len(str(insertion_steps))),
f"ins",
)
if step in (0, int(insertion_steps / 2), insertion_steps - 1):
if report.too_much_force(ctxt, lamb, host_box, combined_bps,
u_impls):
calc_work = False
break
# Note: this condition only applies for ABFE, not RBFE
if (abs(du_dl_obs.full_du_dl()[0]) > 0.001
or abs(du_dl_obs.full_du_dl()[-1]) > 0.001):
print("Error: du_dl endpoints are not ~0")
calc_work = False
if calc_work:
work = np.trapz(du_dl_obs.full_du_dl(),
insertion_lambda_schedule[::subsample_freq])
print(f"guest_name: {guest_name}\tinsertion_work: {work:.2f}")
# equilibrate
for step in range(eq_steps):
ctxt.step(0.00)
if step % 1000 == 0:
report.report_step(ctxt, step, 0.00, host_box, combined_bps,
u_impls, guest_name, eq_steps,
'EQUILIBRATION')
host_coords = ctxt.get_x_t()[:len(solvated_host_coords)] * 10
guest_coords = ctxt.get_x_t()[len(solvated_host_coords):] * 10
report.write_frame(
host_coords,
solvated_host_mol,
guest_coords,
guest_mol,
guest_name,
outdir,
str(step).zfill(len(str(eq_steps))),
f"eq",
)
if step in (0, int(eq_steps / 2), eq_steps - 1):
if report.too_much_force(ctxt, 0.00, host_box, combined_bps,
u_impls):
break
end_time = time.time()
print(f"{guest_name} took {(end_time - start_time):.2f} seconds")
def convergence(args):
epoch, lamb, lamb_idx = args
suppl = Chem.SDMolSupplier("tests/data/ligands_40.sdf", removeHs=False)
ligands = []
for mol in suppl:
ligands.append(mol)
ligand_a = ligands[0]
ligand_b = ligands[1]
# print(ligand_a.GetNumAtoms())
# print(ligand_b.GetNumAtoms())
# ligand_a = Chem.AddHs(Chem.MolFromSmiles("CCCC1=NN(C2=C1N=C(NC2=O)C3=C(C=CC(=C3)S(=O)(=O)N4CCN(CC4)C)OCC)C"))
# ligand_b = Chem.AddHs(Chem.MolFromSmiles("CCCC1=NN(C2=C1N=C(NC2=O)C3=C(C=CC(=C3)S(=O)(=O)N4CCN(CC4)C)OCC)C"))
# ligand_a = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1CC"))
# ligand_b = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1CC"))
# AllChem.EmbedMolecule(ligand_a, randomSeed=2020)
# AllChem.EmbedMolecule(ligand_b, randomSeed=2020)
coords_a = get_conf(ligand_a, idx=0)
coords_b = get_conf(ligand_b, idx=0)
# coords_b = np.matmul(coords_b, special_ortho_group.rvs(3))
coords_a = recenter(coords_a)
coords_b = recenter(coords_b)
coords = np.concatenate([coords_a, coords_b])
a_idxs = get_heavy_atom_idxs(ligand_a)
b_idxs = get_heavy_atom_idxs(ligand_b)
a_full_idxs = np.arange(0, ligand_a.GetNumAtoms())
b_full_idxs = np.arange(0, ligand_b.GetNumAtoms())
b_idxs += ligand_a.GetNumAtoms()
b_full_idxs += ligand_a.GetNumAtoms()
nrg_fns = []
forcefield = 'ff/params/smirnoff_1_1_0_ccc.py'
ff_raw = open(forcefield, "r").read()
ff_handlers = deserialize_handlers(ff_raw)
combined_mol = Chem.CombineMols(ligand_a, ligand_b)
for handler in ff_handlers:
if isinstance(handler, handlers.HarmonicBondHandler):
bond_idxs, (bond_params, _) = handler.parameterize(combined_mol)
nrg_fns.append(
functools.partial(bonded.harmonic_bond,
params=bond_params,
box=None,
bond_idxs=bond_idxs
)
)
elif isinstance(handler, handlers.HarmonicAngleHandler):
angle_idxs, (angle_params, _) = handler.parameterize(combined_mol)
nrg_fns.append(
functools.partial(bonded.harmonic_angle,
params=angle_params,
box=None,
angle_idxs=angle_idxs
)
)
# elif isinstance(handler, handlers.ImproperTorsionHandler):
# torsion_idxs, (torsion_params, _) = handler.parameterize(combined_mol)
# print(torsion_idxs)
# assert 0
# nrg_fns.append(
# functools.partial(bonded.periodic_torsion,
# params=torsion_params,
# box=None,
# lamb=None,
# torsion_idxs=torsion_idxs
# )
# )
# elif isinstance(handler, handlers.ProperTorsionHandler):
# torsion_idxs, (torsion_params, _) = handler.parameterize(combined_mol)
# # print(torsion_idxs)
# nrg_fns.append(
# functools.partial(bonded.periodic_torsion,
# params=torsion_params,
# box=None,
# lamb=None,
# torsion_idxs=torsion_idxs
# )
# )
masses_a = onp.array([a.GetMass() for a in ligand_a.GetAtoms()]) * 10000
masses_b = onp.array([a.GetMass() for a in ligand_b.GetAtoms()])
combined_masses = np.concatenate([masses_a, masses_b])
# com_restraint_fn = functools.partial(bonded.centroid_restraint,
# params=None,
# box=None,
# lamb=None,
# # masses=combined_masses, # try making this ones-like
# masses=np.ones_like(combined_masses),
# group_a_idxs=a_idxs,
# group_b_idxs=b_idxs,
# kb=50.0,
# b0=0.0)
pmi_restraint_fn = functools.partial(pmi_restraints_new,
params=None,
box=None,
lamb=None,
# masses=np.ones_like(combined_masses),
masses=combined_masses,
# a_idxs=a_full_idxs,
# b_idxs=b_full_idxs,
a_idxs=a_idxs,
b_idxs=b_idxs,
angle_force=100.0,
com_force=100.0
)
prefactor = 2.7 # unitless
shape_lamb = (4*np.pi)/(3*prefactor) # unitless
kappa = np.pi/(np.power(shape_lamb, 2/3)) # unitless
sigma = 0.15 # 1 angstrom std, 95% coverage by 2 angstroms
alpha = kappa/(sigma*sigma)
alphas = np.zeros(combined_mol.GetNumAtoms())+alpha
weights = np.zeros(combined_mol.GetNumAtoms())+prefactor
shape_restraint_fn = functools.partial(
shape.harmonic_overlap,
box=None,
lamb=None,
params=None,
a_idxs=a_idxs,
b_idxs=b_idxs,
alphas=alphas,
weights=weights,
k=150.0
)
# shape_restraint_4d_fn = functools.partial(
# shape.harmonic_4d_overlap,
# box=None,
# params=None,
# a_idxs=a_idxs,
# b_idxs=b_idxs,
# alphas=alphas,
# weights=weights,
# k=200.0
# )
def restraint_fn(conf, lamb):
return pmi_restraint_fn(conf) + lamb*shape_restraint_fn(conf)
# return (1-lamb)*pmi_restraint_fn(conf) + lamb*shape_restraint_fn(conf)
nrg_fns.append(restraint_fn)
def nrg_fn(conf, lamb):
s = []
for u in nrg_fns:
s.append(u(conf, lamb=lamb))
return np.sum(s)
grad_fn = jax.grad(nrg_fn, argnums=(0,1))
grad_fn = jax.jit(grad_fn)
du_dx_fn = jax.grad(nrg_fn, argnums=(0))
du_dx_fn = jax.jit(du_dx_fn)
x_t = coords
v_t = np.zeros_like(x_t)
w = Chem.SDWriter('frames_heavy_'+str(epoch)+'_'+str(lamb_idx)+'.sdf')
dt = 1.5e-3
ca, cb, cc = langevin_coefficients(300.0, dt, 1.0, combined_masses)
cb = -1*onp.expand_dims(cb, axis=-1)
cc = onp.expand_dims(cc, axis=-1)
du_dls = []
# re-seed since forking
onp.random.seed(int.from_bytes(os.urandom(4), byteorder='little'))
# for step in range(100000):
for step in range(100000):
# if step % 1000 == 0:
# u = nrg_fn(x_t, lamb)
# print("step", step, "nrg", onp.asarray(u), "avg_du_dl", onp.mean(du_dls))
# mol = make_conformer(combined_mol, x_t[:ligand_a.GetNumAtoms()], x_t[ligand_a.GetNumAtoms():])
# w.write(mol)
# w.flush()
if step % 5 == 0 and step > 10000:
du_dx, du_dl = grad_fn(x_t, lamb)
du_dls.append(du_dl)
else:
du_dx = du_dx_fn(x_t, lamb)
v_t = ca*v_t + cb*du_dx + cc*onp.random.normal(size=x_t.shape)
x_t = x_t + v_t*dt
return np.mean(onp.mean(du_dls))
