def test_jax_nonbonded_block():
"""Assert that nonbonded_block and nonbonded_on_specific_pairs agree"""
system, positions, box, _ = builders.build_water_system(3.0)
bps, masses = openmm_deserializer.deserialize_system(system, cutoff=1.2)
nb = bps[-1]
params = nb.params
conf = positions.value_in_unit(unit.nanometer)
N = conf.shape[0]
beta = nb.get_beta()
cutoff = nb.get_cutoff()
split = 70
def u_a(x, box, params):
xi = x[:split]
xj = x[split:]
pi = params[:split]
pj = params[split:]
return nonbonded_block(xi, xj, box, pi, pj, beta, cutoff)
i_s, j_s = np.indices((split, N - split))
indices_left = i_s.flatten()
indices_right = j_s.flatten() + split
def u_b(x, box, params):
vdw, es = nonbonded_v3_on_specific_pairs(x, params, box, indices_left,
indices_right, beta, cutoff)
return np.sum(vdw + es)
onp.testing.assert_almost_equal(u_a(conf, box, params),
u_b(conf, box, params))
def setUp(self):
# This test checks hilbert curve re-ordering gives identical results
pdb_path = "tests/data/5dfr_solv_equil.pdb"
host_pdb = app.PDBFile(pdb_path)
protein_ff = app.ForceField("amber99sbildn.xml", "tip3p.xml")
host_system = protein_ff.createSystem(
host_pdb.topology, nonbondedMethod=app.NoCutoff, constraints=None, rigidWater=False
)
host_coords = host_pdb.positions
box = host_pdb.topology.getPeriodicBoxVectors()
self.box = np.asarray(box / box.unit)
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.0)
for f in host_fns:
if isinstance(f, potentials.Nonbonded):
self.nonbonded_fn = f
host_conf = []
for x, y, z in host_coords:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
self.host_conf = np.array(host_conf)
self.beta = 2.0
self.cutoff = 1.1
self.lamb = 0.1
def generate_waterbox_nb_args() -> NonbondedArgs:
system, positions, box, _ = builders.build_water_system(3.0)
bps, masses = openmm_deserializer.deserialize_system(system, cutoff=1.2)
nb = bps[-1]
params = nb.params
conf = positions.value_in_unit(unit.nanometer)
N = conf.shape[0]
beta = nb.get_beta()
cutoff = nb.get_cutoff()
lamb = 0.0
charge_rescale_mask = onp.ones((N, N))
lj_rescale_mask = onp.ones((N, N))
lambda_plane_idxs = np.zeros(N, dtype=int)
lambda_offset_idxs = np.zeros(N, dtype=int)
args = (
conf,
params,
box,
lamb,
charge_rescale_mask,
lj_rescale_mask,
beta,
cutoff,
lambda_plane_idxs,
lambda_offset_idxs,
)
return args
def prepare_host_edge(self, ff_params, host_system):
"""
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
Returns
-------
3-tuple
unbound_potentials, system_params, combined_masses
"""
ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()]
ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()]
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_a, ligand_masses_b])
return final_potentials, final_params, combined_masses
def _example_nonbonded_params(_example_system):
host_system, _, _ = _example_system
host_fns, _ = openmm_deserializer.deserialize_system(host_system,
cutoff=1.0)
nonbonded_fn = None
for f in host_fns:
if isinstance(f, potentials.Nonbonded):
nonbonded_fn = f
assert nonbonded_fn is not None
return nonbonded_fn.params
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 test_nblist_box_resize(self):
# test that running the coordinates under two different boxes produces correct results
# since we should be rebuilding the nblist when the box sizes change.
host_system, host_coords, box, _ = builders.build_water_system(3.0)
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.0)
for f in host_fns:
if isinstance(f, potentials.Nonbonded):
test_nonbonded_fn = f
host_conf = []
for x, y, z in host_coords:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
lamb = 0.1
ref_nonbonded_fn = prepare_reference_nonbonded(
test_nonbonded_fn.params,
test_nonbonded_fn.get_exclusion_idxs(),
test_nonbonded_fn.get_scale_factors(),
test_nonbonded_fn.get_lambda_plane_idxs(),
test_nonbonded_fn.get_lambda_offset_idxs(),
test_nonbonded_fn.get_beta(),
test_nonbonded_fn.get_cutoff(),
)
big_box = box + np.eye(3) * 1000
# print(big_box, small_box)
# (ytz): note the ordering should be from large box to small box. though in the current code
# the rebuild is triggered as long as the box *changes*.
for test_box in [big_box, box]:
for precision, rtol, atol in [(np.float64, 1e-8, 1e-10), (np.float32, 1e-4, 3e-5)]:
self.compare_forces(
host_conf,
test_nonbonded_fn.params,
test_box,
lamb,
ref_nonbonded_fn,
test_nonbonded_fn,
rtol=rtol,
atol=atol,
precision=precision,
)
def combine_potentials(ff_handlers, guest_mol, host_system, precision):
"""
This function is responsible for figuring out how to take two separate hamiltonians
and combining them into one sensible alchemical system.
Parameters
----------
ff_handlers: list of forcefield handlers
Small molecule forcefield handlers
guest_mol: Chem.ROMol
RDKit molecule
host_system: openmm.System
Host system to be deserialized
precision: np.float32 or np.float64
Numerical precision of the functional form
Returns
-------
tuple
Returns a list of lib.potentials objects, combined masses, and a list of
their corresponding vjp_fns back into the forcefield
"""
host_potentials, host_masses = openmm_deserializer.deserialize_system(
host_system, precision, cutoff=1.0)
host_nb_bp = None
combined_potentials = []
combined_vjp_fns = []
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:
combined_potentials.append(bp)
combined_vjp_fns.append([])
guest_masses = np.array([a.GetMass() for a in guest_mol.GetAtoms()],
dtype=np.float64)
num_guest_atoms = len(guest_masses)
num_host_atoms = len(host_masses)
combined_masses = np.concatenate([host_masses, guest_masses])
for handle in ff_handlers:
results = handle.parameterize(guest_mol)
if isinstance(handle, bonded.HarmonicBondHandler):
bond_idxs, (bond_params, vjp_fn) = results
bond_idxs += num_host_atoms
combined_potentials.append(
potentials.HarmonicBond(bond_idxs,
precision=precision).bind(bond_params))
combined_vjp_fns.append([(handle, vjp_fn)])
elif isinstance(handle, bonded.HarmonicAngleHandler):
angle_idxs, (angle_params, vjp_fn) = results
angle_idxs += num_host_atoms
combined_potentials.append(
potentials.HarmonicAngle(
angle_idxs, precision=precision).bind(angle_params))
combined_vjp_fns.append([(handle, vjp_fn)])
elif isinstance(handle, bonded.ProperTorsionHandler):
torsion_idxs, (torsion_params, vjp_fn) = results
torsion_idxs += num_host_atoms
combined_potentials.append(
potentials.PeriodicTorsion(
torsion_idxs, precision=precision).bind(torsion_params))
combined_vjp_fns.append([(handle, vjp_fn)])
elif isinstance(handle, bonded.ImproperTorsionHandler):
torsion_idxs, (torsion_params, vjp_fn) = results
torsion_idxs += num_host_atoms
combined_potentials.append(
potentials.PeriodicTorsion(
torsion_idxs, precision=precision).bind(torsion_params))
combined_vjp_fns.append([(handle, vjp_fn)])
elif isinstance(handle, nonbonded.AM1CCCHandler):
charge_handle = handle
guest_charge_params, guest_charge_vjp_fn = results
elif isinstance(handle, nonbonded.LennardJonesHandler):
guest_lj_params, guest_lj_vjp_fn = results
lj_handle = handle
else:
print("Warning: skipping handler", handle)
pass
# process nonbonded terms
combined_nb_params, (charge_vjp_fn, lj_vjp_fn) = nonbonded_vjps(
guest_charge_params, guest_charge_vjp_fn, guest_lj_params,
guest_lj_vjp_fn, host_nb_bp.params)
# these vjp_fns take in adjoints of combined_params and returns derivatives
# appropriate to the underlying handler
combined_vjp_fns.append([(charge_handle, charge_vjp_fn),
(lj_handle, lj_vjp_fn)])
# tbd change scale 14 for electrostatics
guest_exclusion_idxs, guest_scale_factors = nonbonded.generate_exclusion_idxs(
guest_mol, scale12=1.0, scale13=1.0, scale14=0.5)
# allow the ligand to be alchemically decoupled
# a value of one indicates that we allow the atom to be adjusted by the lambda value
guest_lambda_offset_idxs = np.ones(len(guest_masses), dtype=np.int32)
# use same scale factors until we modify 1-4s for electrostatics
guest_scale_factors = np.stack([guest_scale_factors, guest_scale_factors],
axis=1)
combined_lambda_offset_idxs = np.concatenate(
[host_nb_bp.get_lambda_offset_idxs(), guest_lambda_offset_idxs])
combined_exclusion_idxs = np.concatenate([
host_nb_bp.get_exclusion_idxs(), guest_exclusion_idxs + num_host_atoms
])
combined_scales = np.concatenate(
[host_nb_bp.get_scale_factors(), guest_scale_factors])
combined_beta = 2.0
combined_cutoff = 1.0 # nonbonded cutoff
combined_potentials.append(
potentials.Nonbonded(
combined_exclusion_idxs,
combined_scales,
combined_lambda_offset_idxs,
combined_beta,
combined_cutoff,
precision=precision,
).bind(combined_nb_params))
return combined_potentials, combined_masses, combined_vjp_fns
# 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
combined_coords = np.concatenate([host_coords, ligand_coords])
num_host_atoms = host_coords.shape[0]
final_potentials = []
final_vjp_and_handles = []
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()
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,
)
def benchmark_dhfr(verbose=False, num_batches=100, steps_per_batch=1000):
pdb_path = "tests/data/5dfr_solv_equil.pdb"
host_pdb = app.PDBFile(pdb_path)
protein_ff = app.ForceField("amber99sbildn.xml", "tip3p.xml")
host_system = protein_ff.createSystem(
host_pdb.topology, nonbondedMethod=app.NoCutoff, constraints=None, rigidWater=False
)
host_coords = host_pdb.positions
box = host_pdb.topology.getPeriodicBoxVectors()
box = np.asarray(box / box.unit)
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.0)
host_conf = []
for x, y, z in host_coords:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
x0 = host_conf
v0 = np.zeros_like(host_conf)
benchmark(
"dhfr-apo",
host_masses,
0.0,
x0,
v0,
box,
host_fns,
verbose=verbose,
num_batches=num_batches,
steps_per_batch=steps_per_batch,
)
benchmark(
"dhfr-apo-barostat-interval-25",
host_masses,
0.0,
x0,
v0,
box,
host_fns,
verbose=verbose,
num_batches=num_batches,
steps_per_batch=steps_per_batch,
barostat_interval=25,
)
benchmark(
"dhfr-apo-hmr-barostat-interval-25",
host_masses,
0.0,
x0,
v0,
box,
host_fns,
verbose=verbose,
hmr=True,
num_batches=num_batches,
steps_per_batch=steps_per_batch,
barostat_interval=25,
)