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 __init__(self, mol, ff):
"""
VacuumState allows us to enable/disable various parts of a forcefield so that
we can more easily sample across rotational barriers in the vacuum.
Parameters
----------
mol: Chem.ROMol
rdkit molecule
ff: Forcefield
forcefield
"""
self.mol = mol
bt = topology.BaseTopology(mol, ff)
self.bond_params, self.hb_potential = bt.parameterize_harmonic_bond(ff.hb_handle.params)
self.angle_params, self.ha_potential = bt.parameterize_harmonic_angle(ff.ha_handle.params)
self.proper_torsion_params, self.pt_potential = bt.parameterize_proper_torsion(ff.pt_handle.params)
(
self.improper_torsion_params,
self.it_potential,
) = bt.parameterize_improper_torsion(ff.it_handle.params)
self.nb_params, self.nb_potential = bt.parameterize_nonbonded(ff.q_handle.params, ff.lj_handle.params)
self.box = None
self.lamb = 0.0
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_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 __init__(self, mol, ff):
"""
Compute the absolute free energy of a molecule via 4D decoupling.
Parameters
----------
mol: rdkit mol
Ligand to be decoupled
ff: ff.Forcefield
Ligand forcefield
"""
self.mol = mol
self.ff = ff
self.top = topology.BaseTopology(mol, ff)
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_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:])
num_host_atoms = host_coords.shape[0]
final_potentials = []
final_vjp_and_handles = []
# keep the bonded terms in the host the same.
# but we keep the nonbonded term for a subsequent modification
for bp in host_bps:
if isinstance(bp, potentials.Nonbonded):
host_p = bp
else:
final_potentials.append(bp)
final_vjp_and_handles.append(None)
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
gbt = topology.BaseTopology(romol, ff)
hgt = topology.HostGuestTopology(host_p, gbt)
# setup the parameter handlers for the ligand
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]],
[hgt.parameterize_nonbonded, [ff.q_handle, ff.lj_handle]],
]
# instantiate the vjps while parameterizing (forward pass)
for fn, handles in tuples:
params, vjp_fn, potential = jax.vjp(fn,
*[h.params for h in 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()
