def prepare_vacuum_edge(self, ff_params):
"""
Prepares the vacuum system.
Parameters
----------
ff_params: tuple of np.array
Exactly equal to bond_params, angle_params, proper_params, improper_params, charge_params, lj_params
Returns
-------
4 tuple
unbound_potentials, system_parameters, combined_masses, combined_coords
"""
ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()]
ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()]
ligand_coords_a = get_romol_conf(self.mol_a)
ligand_coords_b = get_romol_conf(self.mol_b)
final_params, final_potentials = self._get_system_params_and_potentials(
ff_params, self.top)
combined_masses = np.mean(self.top.interpolate_params(
ligand_masses_a, ligand_masses_b),
axis=0)
combined_coords = np.mean(self.top.interpolate_params(
ligand_coords_a, ligand_coords_b),
axis=0)
return final_potentials, final_params, combined_masses, combined_coords
def test_setting_up_restraints_using_distance():
seed = 814
smi_a = "CCCONNN"
smi_b = "CCCNNN"
mol_a = Chem.MolFromSmiles(smi_a)
mol_a = Chem.AddHs(mol_a)
mol_b = Chem.MolFromSmiles(smi_b)
mol_b = Chem.AddHs(mol_b)
for mol in [mol_a, mol_b]:
AllChem.EmbedMolecule(mol, randomSeed=seed)
mol_a_coords = get_romol_conf(mol_a)
mol_b_coords = get_romol_conf(mol_b)
core = setup_relative_restraints_by_distance(mol_a, mol_b)
assert core.shape == (5, 2)
# If we have a 0 cutoff, expect nothing to overlap
core = setup_relative_restraints_by_distance(mol_a, mol_b, cutoff=0.0)
assert core.shape == (0, )
for cutoff in [0.08, 0.1, 0.2, 1.0]:
core = setup_relative_restraints_by_distance(mol_a,
mol_b,
cutoff=cutoff)
assert core.size > 0
for a, b in core.tolist():
assert np.linalg.norm(mol_a_coords[a] - mol_b_coords[b]) < cutoff
# Adds seven hydrogen (terminal) atoms if allow terminal matches
core = setup_relative_restraints_by_distance(mol_a, mol_b, terminal=True)
assert core.shape == (12, 2)
def test_nonbonded_optimal_map(self):
"""Similar test as test_nonbonbed, ie. assert that coordinates and nonbonded parameters
can be averaged in benzene -> phenol transformation. However, use the maximal mapping possible."""
# map benzene H to phenol O, leaving a dangling phenol H
core = np.array(
[[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6]],
dtype=np.int32)
st = topology.SingleTopology(self.mol_a, self.mol_b, core, self.ff)
x_a = get_romol_conf(self.mol_a)
x_b = get_romol_conf(self.mol_b)
# test interpolation of coordinates.
x_src, x_dst = st.interpolate_params(x_a, x_b)
x_avg = np.mean([x_src, x_dst], axis=0)
assert x_avg.shape == (st.get_num_atoms(), 3)
np.testing.assert_array_equal((x_a[:7] + x_b[:7]) / 2,
x_avg[:7]) # core parts
np.testing.assert_array_equal(x_b[-1], x_avg[7]) # dangling H
params, vjp_fn, pot_c = jax.vjp(st.parameterize_nonbonded,
self.ff.q_handle.params,
self.ff.lj_handle.params,
has_aux=True)
vjp_fn(np.random.rand(*params.shape))
assert params.shape == (2 * st.get_num_atoms(), 3) # qlj
# test interpolation of parameters
bt_a = topology.BaseTopology(self.mol_a, self.ff)
qlj_a, pot_a = bt_a.parameterize_nonbonded(self.ff.q_handle.params,
self.ff.lj_handle.params)
bt_b = topology.BaseTopology(self.mol_b, self.ff)
qlj_b, pot_b = bt_b.parameterize_nonbonded(self.ff.q_handle.params,
self.ff.lj_handle.params)
n_base_params = len(
params
) // 2 # params is actually interpolated, so its 2x number of base params
# qlj_c = np.mean([params[:n_base_params], params[n_base_params:]], axis=0)
params_src = params[:n_base_params]
params_dst = params[n_base_params:]
# core testing
np.testing.assert_array_equal(qlj_a[:7], params_src[:7])
np.testing.assert_array_equal(qlj_b[:7], params_dst[:7])
# r-group atoms in A are all part of the core. so no testing is needed.
# test r-group in B
np.testing.assert_array_equal(qlj_b[7], params_dst[8])
np.testing.assert_array_equal(np.array([0, qlj_b[7][1], 0]),
params_src[8])
def setup_relative_restraints_by_distance(mol_a: Chem.Mol,
mol_b: Chem.Mol,
cutoff: float = 0.1,
terminal: bool = False):
"""
Setup restraints between atoms in two molecules using
a cutoff distance between atoms
Parameters
----------
mol_a: Chem.Mol
First molecule
mol_b: Chem.Mol
Second molecule
cutoff: float=0.1
Distance between atoms to consider as a match
terminal: bool=false
Map terminal atoms
Returns
-------
np.array (N, 2)
Atom mapping between atoms in mol_a to atoms in mol_b.
"""
ligand_coords_a = get_romol_conf(mol_a)
ligand_coords_b = get_romol_conf(mol_b)
core_idxs_a = []
core_idxs_b = []
for idx, a in enumerate(mol_a.GetAtoms()):
if not terminal and a.GetDegree() == 1:
continue
for b_idx, b in enumerate(mol_b.GetAtoms()):
if not terminal and b.GetDegree() == 1:
continue
if np.linalg.norm(ligand_coords_a[idx] -
ligand_coords_b[b_idx]) < cutoff:
core_idxs_a.append(idx)
core_idxs_b.append(b_idx)
assert len(core_idxs_a) == len(core_idxs_b), "Core sizes were inconsistent"
rij = cdist(ligand_coords_a[core_idxs_a], ligand_coords_b[core_idxs_b])
row_idxs, col_idxs = linear_sum_assignment(rij)
core_idxs = []
for core_a, core_b in zip(row_idxs, col_idxs):
core_idxs.append((core_idxs_a[core_a], core_idxs_b[core_b]))
core_idxs = np.array(core_idxs, dtype=np.int32)
return core_idxs
def prepare_host_edge(self, ff_params, host_system, host_coords):
"""
Prepares the host-edge system
Parameters
----------
ff_params: tuple of np.array
Exactly equal to bond_params, angle_params, proper_params, improper_params, charge_params, lj_params
host_system: openmm.System
openmm System object to be deserialized
host_coords: np.array
Nx3 array of atomic coordinates
Returns
-------
4 tuple
unbound_potentials, system_params, combined_masses, combined_coords
"""
ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()]
ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()]
# extract the 0th conformer
ligand_coords_a = get_romol_conf(self.mol_a)
ligand_coords_b = get_romol_conf(self.mol_b)
host_bps, host_masses = openmm_deserializer.deserialize_system(
host_system, cutoff=1.2)
hgt = topology.HostGuestTopology(host_bps, self.top)
final_params, final_potentials = self._get_system_params_and_potentials(
ff_params, hgt)
if isinstance(self.top, topology.SingleTopology):
combined_masses = np.concatenate([
host_masses,
np.mean(self.top.interpolate_params(ligand_masses_a,
ligand_masses_b),
axis=0)
])
combined_coords = np.concatenate([
host_coords,
np.mean(self.top.interpolate_params(ligand_coords_a,
ligand_coords_b),
axis=0)
])
else:
combined_masses = np.concatenate(
[host_masses, ligand_masses_a, ligand_masses_b])
combined_coords = np.concatenate(
[host_coords, ligand_coords_a, ligand_coords_b])
return final_potentials, final_params, combined_masses, combined_coords
def _wrap_simulate(args):
print("IN WRAP SIMULATE")
(
mol,
U_proposal,
U_target,
temperature,
masses,
steps_per_batch,
batches_per_worker,
burn_in_batches,
num_workers,
seed,
) = args
# wraps a callable fn so it runs in a subprocess with the device_count set explicitly
assert multiprocessing.get_start_method() == "spawn"
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=" + str(num_workers)
kT = temperature * BOLTZ
x0 = get_romol_conf(mol)
xs_proposal = simulate(
x0,
U_proposal,
temperature,
masses,
steps_per_batch,
batches_per_worker + burn_in_batches,
num_workers,
seed,
)
num_atoms = mol.GetNumAtoms()
# discard burn-in batches and reshape into a single flat array
xs_proposal = xs_proposal[:, burn_in_batches:, :, :]
batch_U_proposal_fn = jax.pmap(jax.vmap(U_proposal))
batch_U_target_fn = jax.pmap(jax.vmap(U_target))
Us_target = batch_U_target_fn(xs_proposal)
Us_proposal = batch_U_proposal_fn(xs_proposal)
log_numerator = -Us_target.reshape(-1) / kT
log_denominator = -Us_proposal.reshape(-1) / kT
log_weights = log_numerator - log_denominator
# reshape into flat array by removing num_workers dimension
xs_proposal = xs_proposal.reshape(-1, num_atoms, 3)
return xs_proposal, log_weights
def prepare_host_edge(self, ff_params, host_system, host_coords):
"""
Prepares the host-edge system
Parameters
----------
ff_params: tuple of np.array
Exactly equal to bond_params, angle_params, proper_params, improper_params, charge_params, lj_params
host_system: openmm.System
openmm System object to be deserialized
host_coords: np.array
Nx3 array of atomic coordinates
Returns
-------
4 tuple
unbound_potentials, system_params, combined_masses, combined_coords
"""
ligand_masses = [a.GetMass() for a in self.mol.GetAtoms()]
ligand_coords = get_romol_conf(self.mol)
host_bps, host_masses = openmm_deserializer.deserialize_system(
host_system, cutoff=1.2)
hgt = topology.HostGuestTopology(host_bps, self.top)
final_params, final_potentials = self._get_system_params_and_potentials(
ff_params, hgt)
combined_masses = np.concatenate([host_masses, ligand_masses])
combined_coords = np.concatenate([host_coords, ligand_coords])
return final_potentials, final_params, combined_masses, combined_coords
def test_neighborlist_ligand_host():
ligand = hif2a_ligand_pair.mol_a
ligand_coords = get_romol_conf(ligand)
system, host_coords, box, top = build_water_system(4.0)
num_host_atoms = host_coords.shape[0]
host_coords = np.array(host_coords)
coords = np.concatenate([host_coords, ligand_coords])
N = coords.shape[0]
D = 3
cutoff = 1.0
block_size = 32
padding = 0.1
np.random.seed(1234)
diag = np.amax(coords, axis=0) - np.amin(coords, axis=0) + padding
box = np.diag(diag)
# Can only sort the host coords, but not the row/ligand
sort = True
if sort:
perm = hilbert_sort(
coords[:num_host_atoms] + np.argmin(coords[:num_host_atoms]), D)
coords[:num_host_atoms] = coords[:num_host_atoms][perm]
col_coords = np.expand_dims(coords[:num_host_atoms], axis=0)
# Compute the reference interactions of the ligand
ref_ixn_list = []
num_ligand_atoms = coords[num_host_atoms:].shape[0]
num_blocks_of_32 = (num_ligand_atoms + block_size - 1) // block_size
box_diag = np.diag(box)
for rbidx in range(num_blocks_of_32):
row_start = num_host_atoms + (rbidx * block_size)
row_end = min(num_host_atoms + ((rbidx + 1) * block_size), N)
row_coords = coords[row_start:row_end]
row_coords = np.expand_dims(row_coords, axis=1)
deltas = row_coords - col_coords
deltas -= box_diag * np.floor(deltas / box_diag + 0.5)
dij = np.linalg.norm(deltas, axis=-1)
# Since the row and columns are unique, don't need to handle duplicates
idxs = np.argwhere(np.any(dij < cutoff, axis=0))
ref_ixn_list.append(idxs.reshape(-1).tolist())
for nblist in (
custom_ops.Neighborlist_f32(num_host_atoms, num_ligand_atoms),
custom_ops.Neighborlist_f64(num_host_atoms, num_ligand_atoms),
):
for _ in range(2):
test_ixn_list = nblist.get_nblist_host_ligand(
coords[:num_host_atoms], coords[num_host_atoms:], box, cutoff)
# compute the sparsity of the tile
assert len(ref_ixn_list) == len(
test_ixn_list
), "Number of blocks with interactions don't agree"
for bidx, (a, b) in enumerate(zip(ref_ixn_list, test_ixn_list)):
if sorted(a) != sorted(b):
print("TESTING bidx", bidx)
print(sorted(a))
print(sorted(b))
np.testing.assert_equal(sorted(a), sorted(b))
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 test_nonbonded(self):
# leaving benzene H unmapped, and phenol OH unmapped
core = np.array(
[
[0, 0],
[1, 1],
[2, 2],
[3, 3],
[4, 4],
[5, 5],
],
dtype=np.int32,
)
st = topology.SingleTopology(self.mol_a, self.mol_b, core, self.ff)
x_a = get_romol_conf(self.mol_a)
x_b = get_romol_conf(self.mol_b)
# test interpolation of coordinates.
x_src, x_dst = st.interpolate_params(x_a, x_b)
x_avg = np.mean([x_src, x_dst], axis=0)
assert x_avg.shape == (st.get_num_atoms(), 3)
np.testing.assert_array_equal((x_a[:6] + x_b[:6]) / 2, x_avg[:6]) # C
np.testing.assert_array_equal(x_a[6], x_avg[6]) # H
np.testing.assert_array_equal(x_b[6:], x_avg[7:]) # OH
# NOTE: unused result
st.parameterize_nonbonded(self.ff.q_handle.params,
self.ff.lj_handle.params)
params, vjp_fn, pot_c = jax.vjp(st.parameterize_nonbonded,
self.ff.q_handle.params,
self.ff.lj_handle.params,
has_aux=True)
vjp_fn(np.random.rand(*params.shape))
assert params.shape == (2 * st.get_num_atoms(), 3) # qlj
# test interpolation of parameters
bt_a = topology.BaseTopology(self.mol_a, self.ff)
qlj_a, pot_a = bt_a.parameterize_nonbonded(self.ff.q_handle.params,
self.ff.lj_handle.params)
bt_b = topology.BaseTopology(self.mol_b, self.ff)
qlj_b, pot_b = bt_b.parameterize_nonbonded(self.ff.q_handle.params,
self.ff.lj_handle.params)
n_base_params = len(
params
) // 2 # params is actually interpolated, so its 2x number of base params
# qlj_c = np.mean([params[:n_base_params], params[n_base_params:]], axis=0)
params_src = params[:n_base_params]
params_dst = params[n_base_params:]
# core testing
np.testing.assert_array_equal(qlj_a[:6], params_src[:6])
np.testing.assert_array_equal(qlj_b[:6], params_dst[:6])
# test r-group in A
np.testing.assert_array_equal(qlj_a[6], params_src[6])
np.testing.assert_array_equal(np.array([0, qlj_a[6][1], 0]),
params_dst[6])
# test r-group in B
np.testing.assert_array_equal(qlj_b[6:], params_dst[7:])
np.testing.assert_array_equal(
np.array([[0, qlj_b[6][1], 0], [0, qlj_b[7][1], 0]]),
params_src[7:])
def setup_relative_restraints_using_smarts(mol_a, mol_b, smarts):
"""
Setup restraints between atoms in two molecules using
a pre-defined SMARTS pattern.
Parameters
----------
mol_a: Chem.Mol
First molecule
mol_b: Chem.Mol
Second molecule
smarts: string
Smarts pattern defining the common core.
Returns
-------
np.array (N, 2)
Atom mapping between atoms in mol_a to atoms in mol_b.
"""
# check to ensure the core is connected
# technically allow for this but we need to do more validation before
# we can be fully comfortable
assert "." not in smarts
core = Chem.MolFromSmarts(smarts)
# we want *all* possible combinations.
limit = 1000
all_core_idxs_a = np.array(
mol_a.GetSubstructMatches(core, uniquify=False, maxMatches=limit))
all_core_idxs_b = np.array(
mol_b.GetSubstructMatches(core, uniquify=False, maxMatches=limit))
assert len(all_core_idxs_a) < limit
assert len(all_core_idxs_b) < limit
best_rmsd = np.inf
best_core_idxs_a = None
best_core_idxs_b = None
ligand_coords_a = get_romol_conf(mol_a)
ligand_coords_b = get_romol_conf(mol_b)
# setup relative orientational restraints
# rough sketch of algorithm:
# find core atoms in mol_a
# find core atoms in mol_b
# for all matches in mol_a
# for all matches in mol_b
# use the hungarian algorithm to assign matching
# if sum is smaller than best, then store.
for core_idxs_a in all_core_idxs_a:
for core_idxs_b in all_core_idxs_b:
ri = np.expand_dims(ligand_coords_a[core_idxs_a], 1)
rj = np.expand_dims(ligand_coords_b[core_idxs_b], 0)
rij = np.sqrt(np.sum(np.power(ri - rj, 2), axis=-1))
row_idxs, col_idxs = linear_sum_assignment(rij)
rmsd = np.linalg.norm(ligand_coords_a[core_idxs_a[row_idxs]] -
ligand_coords_b[core_idxs_b[col_idxs]])
if rmsd < best_rmsd:
best_rmsd = rmsd
best_core_idxs_a = core_idxs_a
best_core_idxs_b = core_idxs_b
core_idxs = np.stack([best_core_idxs_a, best_core_idxs_b],
axis=1).astype(np.int32)
print("core_idxs", core_idxs, "rmsd", best_rmsd)
return core_idxs
# 1. build water box
# 2. build ligand
# 3. convert to recipe
# construct an RDKit molecule of aspirin
# note: not using OpenFF Molecule because want to avoid the dependency (YTZ?)
romol = Chem.AddHs(Chem.MolFromSmiles("CC(=O)OC1=CC=CC=C1C(=O)O"))
ligand_masses = [a.GetMass() for a in romol.GetAtoms()]
# generate conformers
AllChem.EmbedMolecule(romol)
# extract the 0th conformer
ligand_coords = get_romol_conf(romol)
# construct a 4-nanometer water box (from openmmtools approach: selecting out
# of a large pre-equilibrated water box snapshot)
system, host_coords, box, omm_topology = builders.build_water_system(4.0)
host_bps, host_masses = openmm_deserializer.deserialize_system(system,
cutoff=1.2)
combined_masses = np.concatenate([host_masses, ligand_masses])
# write some conformations into this PDB file
writer = pdb_writer.PDBWriter([omm_topology, romol], "debug.pdb")
# note the order in which the coordinates are concatenated in this step --
# in a later step we will need to combine recipes in the same order
def minimize_host_4d(mols, host_system, host_coords, ff, box, mol_coords=None) -> np.ndarray:
"""
Insert mols into a host system via 4D decoupling using Fire minimizer at lambda=1.0,
0 Kelvin Langevin integration at a sequence of lambda from 1.0 to 0.0, and Fire minimizer again at lambda=0.0
The ligand coordinates are fixed during this, and only host_coords are minimized.
Parameters
----------
mols: list of Chem.Mol
Ligands to be inserted. This must be of length 1 or 2 for now.
host_system: openmm.System
OpenMM System representing the host
host_coords: np.ndarray
N x 3 coordinates of the host. units of nanometers.
ff: ff.Forcefield
Wrapper class around a list of handlers
box: np.ndarray [3,3]
Box matrix for periodic boundary conditions. units of nanometers.
mol_coords: list of np.ndarray
Pre-specify a list of mol coords. Else use the mol.GetConformer(0)
Returns
-------
np.ndarray
This returns minimized host_coords.
"""
assert box.shape == (3, 3)
host_bps, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.2)
num_host_atoms = host_coords.shape[0]
if len(mols) == 1:
top = topology.BaseTopology(mols[0], ff)
elif len(mols) == 2:
top = topology.DualTopologyMinimization(mols[0], mols[1], ff)
else:
raise ValueError("mols must be length 1 or 2")
mass_list = [np.array(host_masses)]
conf_list = [np.array(host_coords)]
for mol in mols:
# mass increase is to keep the ligand fixed
mass_list.append(np.array([a.GetMass() * 100000 for a in mol.GetAtoms()]))
if mol_coords is not None:
for mc in mol_coords:
conf_list.append(mc)
else:
for mol in mols:
conf_list.append(get_romol_conf(mol))
combined_masses = np.concatenate(mass_list)
combined_coords = np.concatenate(conf_list)
hgt = topology.HostGuestTopology(host_bps, top)
u_impls = bind_potentials(hgt, ff)
# this value doesn't matter since we will turn off the noise.
seed = 0
intg = LangevinIntegrator(0.0, 1.5e-3, 1.0, combined_masses, seed).impl()
x0 = combined_coords
v0 = np.zeros_like(x0)
x0 = fire_minimize(x0, u_impls, box, np.ones(50))
# context components: positions, velocities, box, integrator, energy fxns
ctxt = custom_ops.Context(x0, v0, box, intg, u_impls)
ctxt.multiple_steps(np.linspace(1.0, 0, 1000))
final_coords = fire_minimize(ctxt.get_x_t(), u_impls, box, np.zeros(50))
for impl in u_impls:
du_dx, _, _ = impl.execute(final_coords, box, 0.0)
norm = np.linalg.norm(du_dx, axis=-1)
assert np.all(norm < 25000)
return final_coords[:num_host_atoms]
def equilibrate_host(
mol: Chem.Mol,
host_system: openmm.System,
host_coords: NDArray,
temperature: float,
pressure: float,
ff: Forcefield,
box: NDArray,
n_steps: int,
seed: Optional[int] = None,
) -> Tuple[NDArray, NDArray]:
"""
Equilibrate a host system given a reference molecule using the MonteCarloBarostat.
Useful for preparing a host that will be used for multiple FEP calculations using the same reference, IE a starmap.
Performs the following:
- Minimize host with rigid mol
- Minimize host and mol
- Run n_steps with HMR enabled and MonteCarloBarostat every 5 steps
Parameters
----------
mol: Chem.Mol
Ligand for the host to equilibrate with.
host_system: openmm.System
OpenMM System representing the host.
host_coords: np.ndarray
N x 3 coordinates of the host. units of nanometers.
temperature: float
Temperature at which to run the simulation. Units of kelvins.
pressure: float
Pressure at which to run the simulation. Units of bars.
ff: ff.Forcefield
Wrapper class around a list of handlers.
box: np.ndarray [3,3]
Box matrix for periodic boundary conditions. units of nanometers.
n_steps: int
Number of steps to run the simulation for.
seed: int or None
Value to seed simulation with
Returns
-------
tuple (coords, box)
Returns equilibrated system coords as well as the box.
"""
# insert mol into the binding pocket.
host_bps, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.2)
min_host_coords = minimize_host_4d([mol], host_system, host_coords, ff, box)
ligand_masses = [a.GetMass() for a in mol.GetAtoms()]
ligand_coords = get_romol_conf(mol)
combined_masses = np.concatenate([host_masses, ligand_masses])
combined_coords = np.concatenate([min_host_coords, ligand_coords])
top = topology.BaseTopology(mol, ff)
hgt = topology.HostGuestTopology(host_bps, top)
# setup the parameter handlers for the ligand
tuples = [
[hgt.parameterize_harmonic_bond, [ff.hb_handle]],
[hgt.parameterize_harmonic_angle, [ff.ha_handle]],
[hgt.parameterize_periodic_torsion, [ff.pt_handle, ff.it_handle]],
[hgt.parameterize_nonbonded, [ff.q_handle, ff.lj_handle]],
]
u_impls = []
bound_potentials = []
for fn, handles in tuples:
params, potential = fn(*[h.params for h in handles])
bp = potential.bind(params)
bound_potentials.append(bp)
u_impls.append(bp.bound_impl(precision=np.float32))
bond_list = get_bond_list(bound_potentials[0])
combined_masses = model_utils.apply_hmr(combined_masses, bond_list)
dt = 2.5e-3
friction = 1.0
if seed is None:
seed = np.random.randint(np.iinfo(np.int32).max)
integrator = LangevinIntegrator(temperature, dt, friction, combined_masses, seed).impl()
x0 = combined_coords
v0 = np.zeros_like(x0)
group_indices = get_group_indices(bond_list)
barostat_interval = 5
barostat = MonteCarloBarostat(x0.shape[0], pressure, temperature, group_indices, barostat_interval, seed).impl(
u_impls
)
# Re-minimize with the mol being flexible
x0 = fire_minimize(x0, u_impls, box, np.ones(50))
# context components: positions, velocities, box, integrator, energy fxns
ctxt = custom_ops.Context(x0, v0, box, integrator, u_impls, barostat)
ctxt.multiple_steps(np.linspace(0.0, 0.0, n_steps))
return ctxt.get_x_t(), ctxt.get_box()
from timemachine.md import builders, minimizer
# construct an RDKit molecule of aspirin
# note: not using OpenFF Molecule because want to avoid the dependency (YTZ?)
romol_a = Chem.AddHs(Chem.MolFromSmiles("CC(=O)OC1=CC=CC=C1C(=O)O"))
romol_b = Chem.AddHs(Chem.MolFromSmiles("CC(=O)OC1=CC=CC=C1C(=O)OC"))
ligand_masses_a = [a.GetMass() for a in romol_a.GetAtoms()]
ligand_masses_b = [a.GetMass() for a in romol_b.GetAtoms()]
# generate conformers
AllChem.EmbedMolecule(romol_a)
AllChem.EmbedMolecule(romol_b)
# extract the 0th conformer
ligand_coords_a = get_romol_conf(romol_a)
ligand_coords_b = get_romol_conf(romol_b)
# construct a 4-nanometer water box (from openmmtools approach: selecting out
# of a large pre-equilibrated water box snapshot)
system, host_coords, box, omm_topology = builders.build_water_system(4.0)
# padding to avoid jank
box = box + np.eye(3) * 0.1
host_bps, host_masses = openmm_deserializer.deserialize_system(system, cutoff=1.2)
combined_masses = np.concatenate([host_masses, ligand_masses_a, ligand_masses_b])
# minimize coordinates
