def test_loading_forcefield_from_file():
builtin_ffs = glob("timemachine/ff/params/smirnoff_*.py")
for path in builtin_ffs:
# Use the full path
ff = Forcefield.load_from_file(path)
assert ff is not None
# Use full path as Path object
ff = Forcefield.load_from_file(Path(path))
assert ff is not None
# Load using just file name of the built in
ff = Forcefield.load_from_file(os.path.basename(path))
assert ff is not None
for prefix in ["", "timemachine/ff/params/"]:
path = os.path.join(prefix, "nosuchfile.py")
with pytest.raises(ValueError) as e:
Forcefield.load_from_file(path)
assert path in str(e.value)
with pytest.raises(ValueError):
Forcefield.load_from_file(Path(path))
assert path in str(e.value)
with temporary_working_dir():
# Verify that if a local file shadows a builtin
for path in builtin_ffs:
basename = os.path.basename(path)
with open(basename, "w") as ofs:
ofs.write("junk")
with catch_warnings(record=True) as w:
Forcefield.load_from_file(basename)
assert len(w) == 1
assert basename in str(w[0].message)
def test_base_topology_conversion_ring_torsion():
# test that the conversion protocol behaves as intended on a
# simple linked cycle.
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
mol = Chem.MolFromSmiles("C1CC1C1CC1")
vanilla_mol_top = topology.BaseTopology(mol, ff)
vanilla_torsion_params, _ = vanilla_mol_top.parameterize_proper_torsion(
ff.pt_handle.params)
mol_top = topology.BaseTopologyConversion(mol, ff)
conversion_torsion_params, torsion_potential = mol_top.parameterize_proper_torsion(
ff.pt_handle.params)
np.testing.assert_array_equal(vanilla_torsion_params,
conversion_torsion_params)
assert torsion_potential.get_lambda_mult() is None
assert torsion_potential.get_lambda_offset() is None
vanilla_qlj_params, _ = vanilla_mol_top.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
assert isinstance(nonbonded_potential, potentials.NonbondedInterpolated)
src_qlj_params = qlj_params[:len(qlj_params) // 2]
dst_qlj_params = qlj_params[len(qlj_params) // 2:]
np.testing.assert_array_equal(vanilla_qlj_params, src_qlj_params)
np.testing.assert_array_equal(topology.standard_qlj_typer(mol),
dst_qlj_params)
def test_parameterize_and_draw_interactions():
"""
This isn't really tested, but is here to verify that drawing code at least runs.
"""
ff_handlers = deserialize_handlers(
open("timemachine/ff/params/smirnoff_1_1_0_sc.py").read())
ff = Forcefield(ff_handlers)
mol = Chem.MolFromSmiles("CC(=O)OC1=CC=CC=C1C(=O)O")
core = [3, 4, 5, 6, 7, 8, 9, 10]
_draw_impl(mol, core, ff, "aspirin")
mol = Chem.MolFromSmiles("C(C1=CC=CC=C1)C1=CC=CC=C1")
core = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
_draw_impl(mol, core, ff, "met_biphenyl")
mol = Chem.MolFromSmiles("F[C@](Cl)(Br)C1=CC=CC=C1")
core = [1, 4, 5, 6, 7, 8, 9]
_draw_impl(mol, core, ff, "toluene")
mol = Chem.MolFromSmiles("C1CC2=CC=CC=C12")
core = [2, 3, 4, 5, 6, 7]
_draw_impl(mol, core, ff, "ring_open")
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 = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
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_periodic_torsion,
ff.pt_handle.params,
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_dual_topology_standard_decoupling():
# this class is used in double decoupling stages of the RABFE protocol. It modifies the
# DualTopology class in one ways:
# 1) the nonbonded terms are interpolated at lambda=0 such that the epsilons and charges are at half strength.
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
mol_a = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1O"))
mol_b = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1F"))
mol_c = Chem.CombineMols(mol_a, mol_b)
mol_top = topology.DualTopologyStandardDecoupling(mol_a, mol_b, ff)
decouple_torsion_params, torsion_potential = mol_top.parameterize_proper_torsion(
ff.pt_handle.params)
combined_decouple_torsion_params, combined_torsion_potential = mol_top.parameterize_periodic_torsion(
ff.pt_handle.params, ff.it_handle.params)
assert len(combined_torsion_potential.get_lambda_mult()) == len(
combined_torsion_potential.get_idxs())
assert len(combined_torsion_potential.get_lambda_mult()) == len(
combined_torsion_potential.get_lambda_offset())
# impropers should always be turned on.
# num_proper_torsions = len(torsion_potential.get_idxs())
assert np.all(combined_torsion_potential.get_lambda_mult() == 0)
assert np.all(combined_torsion_potential.get_lambda_offset() == 1)
qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
assert isinstance(nonbonded_potential, potentials.NonbondedInterpolated)
expected_qlj = topology.standard_qlj_typer(mol_c)
expected_qlj[:, 0] = expected_qlj[:, 0] / 2 # charges should be halved
expected_qlj[:, 2] = expected_qlj[:, 2] / 2 # eps should be halved
src_qlj_params = qlj_params[:len(qlj_params) // 2]
dst_qlj_params = qlj_params[len(qlj_params) // 2:]
np.testing.assert_array_equal(src_qlj_params, expected_qlj)
expected_qlj = topology.standard_qlj_typer(mol_c)
np.testing.assert_array_equal(dst_qlj_params, expected_qlj)
combined_lambda_plane_idxs = nonbonded_potential.get_lambda_plane_idxs()
combined_lambda_offset_idxs = nonbonded_potential.get_lambda_offset_idxs()
A = mol_a.GetNumAtoms()
B = mol_b.GetNumAtoms()
C = mol_c.GetNumAtoms()
np.testing.assert_array_equal(combined_lambda_plane_idxs, np.zeros(C))
np.testing.assert_array_equal(combined_lambda_offset_idxs[:A], np.zeros(A))
np.testing.assert_array_equal(combined_lambda_offset_idxs[A:], np.ones(B))
def test_base_topology_conversion_r_group():
# check that phenol torsions are turned off
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
mol = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1O"))
mol_top = topology.BaseTopologyConversion(mol, ff)
result, potential = mol_top.parameterize_proper_torsion(
ff.pt_handle.params)
# in the conversion phase, torsions that bridge the two rings should be set to
# be alchemically turned off.
assert potential.get_lambda_mult() is None
assert potential.get_lambda_offset() is None
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
self.ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
super(BenzenePhenolSparseTest, self).__init__(*args, **kwargs)
def _setup_hif2a_ligand_pair(ff="timemachine/ff/params/smirnoff_1_1_0_ccc.py"):
"""Manually constructed atom map
TODO: replace this with a testsystem class similar to those used in openmmtools
"""
path_to_ligand = str(root.joinpath("tests/data/ligands_40.sdf"))
forcefield = Forcefield.load_from_file(ff)
suppl = Chem.SDMolSupplier(path_to_ligand, removeHs=False)
all_mols = [x for x in suppl]
mol_a = all_mols[1]
mol_b = all_mols[4]
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],
])
single_topology = topology.SingleTopology(mol_a, mol_b, core, forcefield)
rfe = free_energy.RelativeFreeEnergy(single_topology)
return rfe
def test_base_topology_standard_decoupling():
# this class is typically used in the second step of the RABFE protocol for the solvent leg.
# we expected the charges to be zero, and the lj parameters to be standardized. In addition,
# the torsions should be turned off.
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
mol = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1O"))
vanilla_mol_top = topology.BaseTopology(mol, ff)
vanilla_torsion_params, _ = vanilla_mol_top.parameterize_proper_torsion(
ff.pt_handle.params)
mol_top = topology.BaseTopologyStandardDecoupling(mol, ff)
decouple_torsion_params, torsion_potential = mol_top.parameterize_proper_torsion(
ff.pt_handle.params)
np.testing.assert_array_equal(vanilla_torsion_params,
decouple_torsion_params)
# in the conversion phase, torsions that bridge the two rings should be set to
# be alchemically turned off.
# is_in_ring = [1, 1, 1, 1, 1, 1, 0, 0]
combined_decouple_torsion_params, combined_torsion_potential = mol_top.parameterize_periodic_torsion(
ff.pt_handle.params, ff.it_handle.params)
assert len(combined_torsion_potential.get_lambda_mult()) == len(
combined_torsion_potential.get_idxs())
assert len(combined_torsion_potential.get_lambda_mult()) == len(
combined_torsion_potential.get_lambda_offset())
# impropers should always be turned on.
# num_proper_torsions = len(torsion_potential.get_idxs())
assert np.all(combined_torsion_potential.get_lambda_mult() == 0)
assert np.all(combined_torsion_potential.get_lambda_offset() == 1)
qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
assert not isinstance(nonbonded_potential,
potentials.NonbondedInterpolated)
np.testing.assert_array_equal(topology.standard_qlj_typer(mol), qlj_params)
np.testing.assert_array_equal(nonbonded_potential.get_lambda_plane_idxs(),
np.zeros(mol.GetNumAtoms(), dtype=np.int32))
np.testing.assert_array_equal(nonbonded_potential.get_lambda_offset_idxs(),
np.ones(mol.GetNumAtoms(), dtype=np.int32))
def test_dual_topology_rhfe():
# used in testing the relative hydration protocol. The nonbonded charges and epsilons are reduced
# to half strength
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
mol_a = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1O"))
mol_b = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1F"))
mol_c = Chem.CombineMols(mol_a, mol_b)
mol_top = topology.DualTopologyRHFE(mol_a, mol_b, ff)
C = mol_a.GetNumAtoms() + mol_b.GetNumAtoms()
ref_qlj_params, _ = topology.BaseTopology(mol_c,
ff).parameterize_nonbonded(
ff.q_handle.params,
ff.lj_handle.params)
qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
assert isinstance(nonbonded_potential, potentials.NonbondedInterpolated)
src_qlj_params = qlj_params[:len(qlj_params) // 2]
dst_qlj_params = qlj_params[len(qlj_params) // 2:]
np.testing.assert_array_equal(src_qlj_params[:, 0],
ref_qlj_params[:, 0] / 2)
np.testing.assert_array_equal(src_qlj_params[:, 1], ref_qlj_params[:, 1])
np.testing.assert_array_equal(src_qlj_params[:, 2],
ref_qlj_params[:, 2] / 2)
np.testing.assert_array_equal(dst_qlj_params, ref_qlj_params)
combined_lambda_plane_idxs = nonbonded_potential.get_lambda_plane_idxs()
combined_lambda_offset_idxs = nonbonded_potential.get_lambda_offset_idxs()
A = mol_a.GetNumAtoms()
B = mol_b.GetNumAtoms()
C = mol_c.GetNumAtoms()
np.testing.assert_array_equal(combined_lambda_plane_idxs, np.zeros(C))
np.testing.assert_array_equal(combined_lambda_offset_idxs[:A], np.zeros(A))
np.testing.assert_array_equal(combined_lambda_offset_idxs[A:], np.ones(B))
def test_minimizer():
complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(
"tests/data/hif2a_nowater_min.pdb")
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 = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
# these methods will throw if the minimization failed
minimizer.minimize_host_4d([mol_a, mol_b], complex_system, complex_coords,
ff, complex_box)
minimizer.minimize_host_4d([mol_a], complex_system, complex_coords, ff,
complex_box)
minimizer.minimize_host_4d([mol_b], complex_system, complex_coords, ff,
complex_box)
def test_equilibrate_host():
host_system, host_coords, host_box, _ = builders.build_water_system(4.0)
suppl = Chem.SDMolSupplier("tests/data/ligands_40.sdf", removeHs=False)
mol = next(suppl)
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
coords, box = minimizer.equilibrate_host(mol,
host_system,
host_coords,
300,
1.0,
ff,
host_box,
25,
seed=2022)
assert coords.shape[0] == host_coords.shape[0] + mol.GetNumAtoms()
assert coords.shape[1] == host_coords.shape[1]
assert box.shape == host_box.shape
def test_dual_topology_standard_decoupling_charged():
# special test case for charged molecules, we expect the charges to be rescaled
# based on each individual molecule's charge, as opposed to based on the sum
# of the charges.
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
mol_a = Chem.AddHs(Chem.MolFromSmiles("C1CC1[O-]"))
mol_b = Chem.AddHs(Chem.MolFromSmiles("C1[O+]CCCCC1"))
mol_top = topology.DualTopologyStandardDecoupling(mol_a, mol_b, ff)
qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
assert isinstance(nonbonded_potential, potentials.NonbondedInterpolated)
expected_qlj = np.concatenate([
topology.standard_qlj_typer(mol_a),
topology.standard_qlj_typer(mol_b)
])
# need to set the charges correctly, and manually
N_A = mol_a.GetNumAtoms()
N_B = mol_b.GetNumAtoms()
expected_qlj[:N_A, 0] = -1.0 / N_A
expected_qlj[N_A:, 0] = 1.0 / N_B
expected_qlj[:, 2] = expected_qlj[:, 2] # eps should be halved
src_qlj_params = qlj_params[:len(qlj_params) // 2]
dst_qlj_params = qlj_params[len(qlj_params) // 2:]
np.testing.assert_array_equal(dst_qlj_params, expected_qlj)
expected_qlj[:N_A, 0] /= 2
expected_qlj[N_A:, 0] /= 2
expected_qlj[:, 2] /= 2 # eps should be halved
np.testing.assert_array_equal(src_qlj_params[:, 0], expected_qlj[:, 0])
def test_dual_topology_minimization():
# Identical to the vanilla Dual Topology class, except that both ligands are
# decouple simultaneously
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
mol_a = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1O"))
mol_b = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1F"))
mol_top = topology.DualTopologyMinimization(mol_a, mol_b, ff)
C = mol_a.GetNumAtoms() + mol_b.GetNumAtoms()
_, nonbonded_potential = mol_top.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
assert not isinstance(nonbonded_potential,
potentials.NonbondedInterpolated)
np.testing.assert_array_equal(nonbonded_potential.get_lambda_offset_idxs(),
np.ones(C, dtype=np.int32))
np.testing.assert_array_equal(nonbonded_potential.get_lambda_plane_idxs(),
np.zeros(C, dtype=np.int32))
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 = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
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):
topology.SingleTopology(mol_a, mol_b, core, ff)
def test_predict(self):
"""Just to verify that we can handle the most basic RBFE prediction"""
# Use the Simple Charges to verify determinism of model. Needed as one endpoint uses the ff definition
forcefield = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(
os.path.join(DATA_DIR, "hif2a_nowater_min.pdb"))
# build the water system
solvent_system, solvent_coords, solvent_box, _ = builders.build_water_system(
4.0)
client = CUDAPoolClient(NUM_GPUS)
model = RBFEModel(
client=client,
ff=forcefield,
complex_system=complex_system,
complex_coords=complex_coords,
complex_box=complex_box,
complex_schedule=construct_lambda_schedule(2),
solvent_system=solvent_system,
solvent_coords=solvent_coords,
solvent_box=solvent_box,
solvent_schedule=construct_lambda_schedule(2),
equil_steps=10,
prod_steps=100,
)
ordered_params = forcefield.get_ordered_params()
mol_a = hif2a_ligand_pair.mol_a
mol_b = hif2a_ligand_pair.mol_b
core = hif2a_ligand_pair.core
ddg, results = model.predict(ordered_params, mol_a, mol_b, core)
self.assertEqual(len(results), 2)
self.assertIsInstance(ddg, float)
def get_ff_am1cc():
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
return ff
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 = Forcefield.load_from_file(
"smirnoff_1_1_0_ccc.py").get_ordered_handles()
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 calculate_rigorous_work(
host_pdbfile,
guests_sdfile,
outdir,
num_deletions,
deletion_steps,
insertion_max_lambda=0.5,
insertion_steps=501,
eq1_steps=5001,
fewer_outfiles=False,
no_outfiles=False,
):
"""Runs non-equilibrium deletion jobs:
1. Solvates a protein, inserts guest, equilibrates, equilibrates more & spins off
deletion jobs every 1000th step, calculates work.
2. Does the same thing in solvent instead of protein.
Does num_deletions deletion jobs per leg per compound.
Parameters
----------
host_pdbfile (str): path to host pdb file
guests_sdfile (str): path to guests sdf file
outdir (str): path to directory to which to write output
num_deletions (int): number of deletion trajectories to run per leg per compound
deletion_steps (int): length of each deletion trajectory
insertion_max_lambda (float): how far away to insert from (0.0-1.0)
insertion_steps (int): how long to insert over
eq1_steps (int): how long to equilibrate after insertion and before starting the deletions
fewer_outfiles (bool): only save the starting frame of each deletion trajectory
no_outfiles (bool): don't keep any output files
Returns
-------
{str: {str: float}}: map of compound to leg label to work values
{'guest_1': {'protein': [work values], 'solvent': [work_values]}, ...}
Output
------
A pdb & sdf file for each guest's final insertion step
(outdir/<guest_name>_pd_<step>_host.pdb & outdir/<guest_name>_pd_<step>_guest.sdf)
(unless fewer_outfiles or no_outfiles is True)
A pdb & sdf file for each guest's final eq1 step
(outdir/<guest_name>_pd_<step>_host.pdb & outdir/<guest_name>_pd_<step>_guest.sdf)
(unless fewer_outfiles or no_outfiles is True)
A pdb & sdf file for each deletion job's first step
(outdir/<guest_name>_pd_<step>_host.pdb & outdir/<guest_name>_pd_<step>_guest.sdf)
(unless no_outfiles is True)
stdout corresponding to the files written noting the lambda value and energy
stdout noting the work of deletion, if applicable
stdout noting how long each leg took to run
Note
----
The work will not be calculated if the du_dl endpoints are not close to 0 or if any norm of
force per atom exceeds 20000 kJ/(mol*nm) [MAX_NORM_FORCE defined in docking/report.py]
"""
if not os.path.exists(outdir):
os.makedirs(outdir)
print(f"""
HOST_PDBFILE = {host_pdbfile}
GUESTS_SDFILE = {guests_sdfile}
OUTDIR = {outdir}
DELETION_MAX_LAMBDA = {DELETION_MAX_LAMBDA}
MIN_LAMBDA = {MIN_LAMBDA}
insertion_max_lambda = {insertion_max_lambda}
insertion_steps = {insertion_steps}
eq1_steps = {eq1_steps}
num_deletions = {num_deletions}
deletion_steps = {deletion_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)
_, solvated_host_pdb = tempfile.mkstemp(suffix=".pdb", text=True)
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)
# 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,
water_coords,
water_box,
water_topology,
) = builders.build_water_system(water_box_width)
# it's okay if the water box here and the solvated protein box don't align -- they have PBCs
_, water_pdb = tempfile.mkstemp(suffix=".pdb", text=True)
writer = pdb_writer.PDBWriter([water_topology], water_pdb)
writer.write_frame(water_coords)
writer.close()
water_mol = Chem.MolFromPDBFile(water_pdb, removeHs=False)
os.remove(water_pdb)
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
# Run the procedure
all_works = defaultdict(dict)
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
for system, coords, host_mol, box, label in zip(
[solvated_host_system, water_system],
[solvated_host_coords, water_coords],
[solvated_host_mol, water_mol],
[host_box, water_box],
["protein", "solvent"],
):
minimized_coords = minimizer.minimize_host_4d([guest_mol], system,
coords, ff, box)
afe = free_energy.AbsoluteFreeEnergy(guest_mol, ff)
ups, sys_params, combined_masses, combined_coords = afe.prepare_host_edge(
ff.get_ordered_params(), system, minimized_coords)
combined_bps = []
for up, sp in zip(ups, sys_params):
combined_bps.append(up.bind(sp))
works = run_leg(
minimized_coords,
orig_guest_coords,
combined_bps,
combined_masses,
box,
guest_name,
label,
host_mol,
guest_mol,
outdir,
num_deletions,
deletion_steps,
insertion_max_lambda,
insertion_steps,
eq1_steps,
fewer_outfiles,
no_outfiles,
)
all_works[guest_name][label] = works
end_time = time.time()
print(
f"{guest_name} {label} leg time:",
"%.2f" % (end_time - start_time),
"seconds",
)
return all_works
def test_predict_absolute_conversion(self):
"""Just to verify that we can handle the most basic conversion RABFE prediction"""
# Use the Simple Charges to verify determinism of model. Needed as one endpoint uses the ff definition
forcefield = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
# build the water system
solvent_system, solvent_coords, solvent_box, solvent_topology = builders.build_water_system(
4.0)
temperature = 300.0
pressure = 1.0
dt = 2.5e-3
client = CUDAPoolClient(NUM_GPUS)
model = AbsoluteConversionModel(
client,
forcefield,
solvent_system,
construct_lambda_schedule(2),
solvent_topology,
temperature,
pressure,
dt,
10,
50,
frame_filter=all_frames,
)
mol_a = hif2a_ligand_pair.mol_a
mol_b = hif2a_ligand_pair.mol_b
core_idxs = setup_relative_restraints_by_distance(mol_a, mol_b)
ref_coords = get_romol_conf(mol_a)
mol_coords = get_romol_conf(mol_b) # original coords
# Use core_idxs to generate
R, t = rmsd.get_optimal_rotation_and_translation(
x1=ref_coords[core_idxs[:, 1]], # reference core atoms
x2=mol_coords[core_idxs[:, 0]], # mol core atoms
)
aligned_mol_coords = rmsd.apply_rotation_and_translation(
mol_coords, R, t)
solvent_coords = minimizer.minimize_host_4d([mol_b], solvent_system,
solvent_coords, forcefield,
solvent_box,
[aligned_mol_coords])
solvent_x0 = np.concatenate([solvent_coords, aligned_mol_coords])
ordered_params = forcefield.get_ordered_params()
with temporary_working_dir() as temp_dir:
dG, dG_err = model.predict(ordered_params,
mol_b,
solvent_x0,
solvent_box,
"prefix",
core_idxs=core_idxs[:, 0],
seed=2022)
np.testing.assert_almost_equal(dG, 46.102816, decimal=5)
np.testing.assert_equal(dG_err, 0.0)
created_files = os.listdir(temp_dir)
# 2 npz, 1 pdb and 1 npy per mol due to a->b and b->a
self.assertEqual(len(created_files), 4)
self.assertEqual(
len([x for x in created_files if x.endswith(".pdb")]), 1)
self.assertEqual(
len([x for x in created_files if x.endswith(".npy")]), 1)
self.assertEqual(
len([x for x in created_files if x.endswith(".npz")]), 2)
cmd_args = parser.parse_args()
if not cmd_args.hosts:
num_gpus = cmd_args.num_gpus
# set up multi-GPU client
client = CUDAPoolClient(max_workers=num_gpus)
else:
# Setup GRPC client
print("Connecting to GRPC workers...")
client = GRPCClient(hosts=cmd_args.hosts)
client.verify()
path_to_ligand = "tests/data/ligands_40.sdf"
suppl = Chem.SDMolSupplier(path_to_ligand, removeHs=False)
forcefield = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
mols = [x for x in suppl]
dataset = Dataset(mols)
absolute_solvent_schedule = construct_absolute_lambda_schedule_solvent(
cmd_args.num_windows)
relative_solvent_schedule = construct_relative_lambda_schedule(
cmd_args.num_windows - 1)
solvent_system, solvent_coords, solvent_box, solvent_topology = builders.build_water_system(
4.0)
# pick the largest mol as the blocker
largest_size = 0
ref_mol = None
for mol in mols:
network_generation_methods = {
"manual": manual_hub_selection,
"mcs": mcs_star_map,
}
if __name__ == "__main__":
parser = ArgumentParser(description="Generate FEP edge map")
parser.add_argument("config", help="YAML configuration")
args = parser.parse_args()
config = TimemachineConfig.from_yaml(args.config)
if config.map_generation is None:
print("No map generation configuration provided")
sys.exit(1)
map_config = config.map_generation
forcefield = Forcefield.load_from_file(map_config.forcefield)
mols = []
# Ensure we pickle up mol properties
Chem.SetDefaultPickleProperties(Chem.PropertyPickleOptions.AllProps)
for lig_path in map_config.ligands:
supplier = Chem.SDMolSupplier(lig_path, removeHs=False)
# Prefer the list comp over casting to a list, seems to be inconsistent
# on OSX/linux
mols.extend([mol for mol in supplier])
# In the future hopefully we can programmatically find the cores rather specifying
cores = map_config.cores
core_sets = {}
for i, core in enumerate(cores):
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 for the last step of insertion
(outdir/<guest_name>/<guest_name>_ins_<step>_[host.pdb/guest.sdf])
A pdb & sdf file every 1000 steps of equilibration
(outdir/<guest_name>/<guest_name>_eq_<step>_[host.pdb/guest.sdf])
stdout corresponding to the files written noting the lambda value and energy
stdout for each guest noting the work of transition, if applicable
stdout for each guest noting how long it took to run
Note
----
The work will not be calculated if the du_dl endpoints are not close to 0 or if any norm of
force per atom exceeds 20000 kJ/(mol*nm) [MAX_NORM_FORCE defined in docking/report.py]
"""
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...")
(
solvated_host_system,
solvated_host_coords,
_,
_,
host_box,
solvated_topology,
) = builders.build_protein_system(host_pdbfile)
_, solvated_host_pdb = tempfile.mkstemp(suffix=".pdb", text=True)
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)
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
# 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
minimized_coords = minimizer.minimize_host_4d([guest_mol],
solvated_host_system,
solvated_host_coords, ff,
host_box)
afe = free_energy.AbsoluteFreeEnergy(guest_mol, ff)
ups, sys_params, combined_masses, _ = afe.prepare_host_edge(
ff.get_ordered_params(), solvated_host_system, minimized_coords)
combined_bps = []
for up, sp in zip(ups, sys_params):
combined_bps.append(up.bind(sp))
x0 = np.concatenate([minimized_coords, orig_guest_coords])
v0 = np.zeros_like(x0)
print("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)
# insert guest
insertion_lambda_schedule = np.linspace(max_lambda, 0.0,
insertion_steps)
calc_work = True
# collect a du_dl calculation once every other step
subsample_interval = 1
full_du_dls, _, _ = ctxt.multiple_steps(insertion_lambda_schedule,
subsample_interval)
step = len(insertion_lambda_schedule) - 1
lamb = insertion_lambda_schedule[-1]
ctxt.step(lamb)
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))),
"ins",
)
if report.too_much_force(ctxt, lamb, host_box, combined_bps, u_impls):
print("Not calculating work (too much force)")
calc_work = False
continue
# Note: this condition only applies for ABFE, not RBFE
if abs(full_du_dls[0]) > 0.001 or abs(full_du_dls[-1]) > 0.001:
print("Not calculating work (du_dl endpoints are not ~0)")
calc_work = False
if calc_work:
work = np.trapz(full_du_dls,
insertion_lambda_schedule[::subsample_interval])
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",
)
if (not fewer_outfiles) or (step == eq_steps - 1):
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))),
"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 test_absolute_free_energy():
suppl = Chem.SDMolSupplier("tests/data/ligands_40.sdf", removeHs=False)
all_mols = [x for x in suppl]
mol = all_mols[1]
complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(
"tests/data/hif2a_nowater_min.pdb")
# build the water system.
solvent_system, solvent_coords, solvent_box, _ = builders.build_water_system(
4.0)
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
ff_params = ff.get_ordered_params()
seed = 2021
lambda_schedule = np.linspace(0, 1.0, 4)
equil_steps = 1000
prod_steps = 1000
afe = free_energy.AbsoluteFreeEnergy(mol, ff)
def absolute_model(ff_params):
dGs = []
for host_system, host_coords, host_box in [
(complex_system, complex_coords, complex_box),
(solvent_system, solvent_coords, solvent_box),
]:
# minimize the host to avoid clashes
host_coords = minimizer.minimize_host_4d([mol], host_system,
host_coords, ff, host_box)
unbound_potentials, sys_params, masses, coords = afe.prepare_host_edge(
ff_params, host_system, host_coords)
harmonic_bond_potential = unbound_potentials[0]
group_idxs = get_group_indices(
get_bond_list(harmonic_bond_potential))
x0 = coords
v0 = np.zeros_like(coords)
client = CUDAPoolClient(1)
temperature = 300.0
pressure = 1.0
integrator = LangevinIntegrator(temperature, 1.5e-3, 1.0, masses,
seed)
barostat = MonteCarloBarostat(x0.shape[0], pressure, temperature,
group_idxs, 25, seed)
model = estimator.FreeEnergyModel(
unbound_potentials,
client,
host_box,
x0,
v0,
integrator,
lambda_schedule,
equil_steps,
prod_steps,
barostat,
)
dG, _ = estimator.deltaG(model, sys_params)
dGs.append(dG)
return dGs[0] - dGs[1]
dG = absolute_model(ff_params)
assert np.abs(dG) < 1000.0
elif args.config == "production":
configuration = production_configuration # a lot
assert configuration is not None, "No configuration provided"
if not args.hosts:
num_gpus = args.num_gpus
if num_gpus is None:
num_gpus = NUM_GPUS
# set up multi-GPU client
client = CUDAPoolClient(max_workers=num_gpus)
else:
# Setup GRPC client
client = GRPCClient(hosts=args.hosts)
client.verify()
forcefield = Forcefield.load_from_file(args.path_to_ff)
relative_transformations: List[RelativeFreeEnergy] = []
# load pre-defined collection of relative transformations
for edge_path in args.path_to_edges:
with open(edge_path, "rb") as f:
relative_transformations.extend(load(f))
# if older transformation lack a complex_path, rely on --protein_path to set
protein_paths = set(x.complex_path for x in relative_transformations
if hasattr(x, "complex_path"))
if protein_path is not None:
protein_paths.add(protein_path)
if len(protein_paths) == 0:
print("No proteins provided by edges or with --protein_path")
sys.exit(1)
def do_relative_docking(host_pdbfile, mol_a, mol_b, core, num_switches,
transition_steps):
"""Runs non-equilibrium switching jobs:
1. Solvates a protein, minimizes w.r.t guest_A, equilibrates & spins off switching jobs
(deleting guest_A while inserting guest_B) every 1000th step, calculates work.
2. Does the same thing in solvent instead of protein
Does num_switches switching jobs per leg.
Parameters
----------
host_pdbfile (str): path to host pdb file
mol_a (rdkit mol): the starting ligand to swap from
mol_b (rdkit mol): the ending ligand to swap to
core (np.array[[int, int], [int, int], ...]): the common core atoms between mol_a and mol_b
num_switches (int): number of switching trajectories to run per compound pair per leg
transition_stpes (int): length of each switching trajectory
Returns
-------
{str: float}: map of leg label to work values of switching mol_a to mol_b in that leg,
{'protein': [work values], 'solvent': [work_values]}
Output
------
stdout noting the step number, lambda value, and energy at various steps
stdout noting the work of transition, if applicable
stdout noting how long it took to run
Note
----
The work will not be calculated if any norm of force per atom exceeds 20000 kJ/(mol*nm)
[MAX_NORM_FORCE defined in docking/report.py]
The simulations won't run if the atom maps are not factorizable
"""
# 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)
# 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,
water_coords,
water_box,
water_topology,
) = builders.build_water_system(water_box_width)
# it's okay if the water box here and the solvated protein box don't align -- they have PBCs
# Run the procedure
start_time = time.time()
guest_name_a = mol_a.GetProp("_Name")
guest_name_b = mol_b.GetProp("_Name")
combined_name = guest_name_a + "-->" + guest_name_b
guest_conformer_a = mol_a.GetConformer(0)
orig_guest_coords_a = np.array(guest_conformer_a.GetPositions(),
dtype=np.float64)
orig_guest_coords_a = orig_guest_coords_a / 10 # convert to md_units
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
all_works = {}
for system, coords, box, label in zip(
[solvated_host_system, water_system],
[solvated_host_coords, water_coords],
[host_box, water_box],
["protein", "solvent"],
):
# minimize w.r.t. both mol_a and mol_b?
min_coords = minimizer.minimize_host_4d([mol_a], system, coords, ff,
box)
try:
single_topology = topology.SingleTopology(mol_a, mol_b, core, ff)
rfe = free_energy.RelativeFreeEnergy(single_topology)
ups, sys_params, combined_masses, combined_coords = rfe.prepare_host_edge(
ff.get_ordered_params(), system, min_coords)
except topology.AtomMappingError as e:
print(f"NON-FACTORIZABLE PAIR: {combined_name}")
print(e)
return {}
combined_bps = []
for up, sp in zip(ups, sys_params):
combined_bps.append(up.bind(sp))
all_works[label] = run_leg(
combined_coords,
combined_bps,
combined_masses,
box,
combined_name,
label,
num_switches,
transition_steps,
)
end_time = time.time()
print(
f"{combined_name} {label} leg time:",
"%.2f" % (end_time - start_time),
"seconds",
)
return all_works
def test_serialization_of_ffs():
for path in glob("timemachine/ff/params/smirnoff_*.py"):
ff = Forcefield(deserialize_handlers(open(path).read()))
for handle in ff.get_ordered_handles():
assert handle is not None, f"{path} failed to deserialize correctly"
def get_110_ccc_ff():
forcefield = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
return forcefield
def test_relative_free_energy():
# test that we can properly build a single topology host guest system and
# that we can run a few steps in a stable way. This tests runs both the complex
# and the solvent stages.
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]
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],
])
complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(
"tests/data/hif2a_nowater_min.pdb")
# build the water system.
solvent_system, solvent_coords, solvent_box, _ = builders.build_water_system(
4.0)
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
ff_params = ff.get_ordered_params()
seed = 2021
lambda_schedule = np.linspace(0, 1.0, 4)
equil_steps = 1000
prod_steps = 1000
single_topology = topology.SingleTopology(mol_a, mol_b, core, ff)
rfe = free_energy.RelativeFreeEnergy(single_topology)
def vacuum_model(ff_params):
unbound_potentials, sys_params, masses, coords = rfe.prepare_vacuum_edge(
ff_params)
x0 = coords
v0 = np.zeros_like(coords)
client = CUDAPoolClient(1)
box = np.eye(3, dtype=np.float64) * 100
harmonic_bond_potential = unbound_potentials[0]
group_idxs = get_group_indices(get_bond_list(harmonic_bond_potential))
x0 = coords
v0 = np.zeros_like(coords)
client = CUDAPoolClient(1)
temperature = 300.0
pressure = 1.0
integrator = LangevinIntegrator(temperature, 1.5e-3, 1.0, masses, seed)
barostat = MonteCarloBarostat(x0.shape[0], pressure, temperature,
group_idxs, 25, seed)
model = estimator.FreeEnergyModel(unbound_potentials, client, box, x0,
v0, integrator, lambda_schedule,
equil_steps, prod_steps, barostat)
return estimator.deltaG(model, sys_params)[0]
dG = vacuum_model(ff_params)
assert np.abs(dG) < 1000.0
def binding_model(ff_params):
dGs = []
for host_system, host_coords, host_box in [
(complex_system, complex_coords, complex_box),
(solvent_system, solvent_coords, solvent_box),
]:
# minimize the host to avoid clashes
host_coords = minimizer.minimize_host_4d([mol_a], host_system,
host_coords, ff, host_box)
unbound_potentials, sys_params, masses, coords = rfe.prepare_host_edge(
ff_params, host_system, host_coords)
x0 = coords
v0 = np.zeros_like(coords)
client = CUDAPoolClient(1)
harmonic_bond_potential = unbound_potentials[0]
group_idxs = get_group_indices(
get_bond_list(harmonic_bond_potential))
temperature = 300.0
pressure = 1.0
integrator = LangevinIntegrator(temperature, 1.5e-3, 1.0, masses,
seed)
barostat = MonteCarloBarostat(x0.shape[0], pressure, temperature,
group_idxs, 25, seed)
model = estimator.FreeEnergyModel(
unbound_potentials,
client,
host_box,
x0,
v0,
integrator,
lambda_schedule,
equil_steps,
prod_steps,
barostat,
)
dG, _ = estimator.deltaG(model, sys_params)
dGs.append(dG)
return dGs[0] - dGs[1]
dG = binding_model(ff_params)
assert np.abs(dG) < 1000.0
def benchmark_hif2a(verbose=False, num_batches=100, steps_per_batch=1000):
from timemachine.testsystems.relative import hif2a_ligand_pair as testsystem
mol_a, mol_b, core = testsystem.mol_a, testsystem.mol_b, testsystem.core
ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py")
single_topology = SingleTopology(mol_a, mol_b, core, ff)
rfe = free_energy.RelativeFreeEnergy(single_topology)
ff_params = ff.get_ordered_params()
# build the protein system.
complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(
"tests/data/hif2a_nowater_min.pdb"
)
solvent_system, solvent_coords, solvent_box, _ = builders.build_water_system(4.0)
for stage, host_system, host_coords, host_box in [
("hif2a", complex_system, complex_coords, complex_box),
("solvent", solvent_system, solvent_coords, solvent_box),
]:
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.0)
# resolve host clashes
min_host_coords = minimizer.minimize_host_4d([mol_a, mol_b], host_system, host_coords, ff, host_box)
x0 = min_host_coords
v0 = np.zeros_like(x0)
# lamb = 0.0
benchmark(
stage + "-apo",
host_masses,
0.0,
x0,
v0,
host_box,
host_fns,
verbose=verbose,
num_batches=num_batches,
steps_per_batch=steps_per_batch,
)
benchmark(
stage + "-apo-barostat-interval-25",
host_masses,
0.0,
x0,
v0,
host_box,
host_fns,
verbose=verbose,
num_batches=num_batches,
steps_per_batch=steps_per_batch,
barostat_interval=25,
)
# RBFE
unbound_potentials, sys_params, masses, coords = rfe.prepare_host_edge(ff_params, host_system, x0)
bound_potentials = [x.bind(y) for (x, y) in zip(unbound_potentials, sys_params)]
x0 = coords
v0 = np.zeros_like(x0)
# lamb = 0.5
benchmark(
stage + "-rbfe-with-du-dp",
masses,
0.5,
x0,
v0,
host_box,
bound_potentials,
verbose=verbose,
num_batches=num_batches,
steps_per_batch=steps_per_batch,
)
for du_dl_interval in [0, 1, 5]:
benchmark(
stage + "-rbfe-du-dl-interval-" + str(du_dl_interval),
masses,
0.5,
x0,
v0,
host_box,
bound_potentials,
verbose=verbose,
num_batches=num_batches,
steps_per_batch=steps_per_batch,
compute_du_dl_interval=du_dl_interval,
)