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 test_bad_factor(self):
# test a bad mapping that results in a non-cancellable endpoint
suppl = Chem.SDMolSupplier('tests/data/ligands_40.sdf', removeHs=False)
all_mols = [x for x in suppl]
mol_a = all_mols[0]
mol_b = all_mols[1]
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_recharge.py').read())
ff = Forcefield(ff_handlers)
core = np.array([[4, 1], [5, 2], [6, 3], [7, 4], [8, 5], [9,
6], [10, 7],
[11, 8], [12, 9], [13, 10], [15, 11], [16, 12],
[18, 14], [34, 31], [17, 13], [23, 23], [33, 30],
[32, 28], [31, 27], [30, 26], [19, 15], [20, 16],
[21, 17]])
with self.assertRaises(topology.AtomMappingError):
st = topology.SingleTopology(mol_a, mol_b, core, ff)
def test_good_factor(self):
# test a good mapping
suppl = Chem.SDMolSupplier('tests/data/ligands_40.sdf', removeHs=False)
all_mols = [x for x in suppl]
mol_a = all_mols[1]
mol_b = all_mols[4]
ff_handlers = deserialize_handlers(
open('ff/params/smirnoff_1_1_0_recharge.py').read())
ff = Forcefield(ff_handlers)
core = np.array([[0, 0], [2, 2], [1, 1], [6, 6], [5, 5], [4, 4],
[3, 3], [15, 16], [16, 17], [17, 18], [18, 19],
[19, 20], [20, 21], [32, 30], [26, 25], [27, 26],
[7, 7], [8, 8], [9, 9], [10, 10], [29, 11], [11, 12],
[12, 13], [14, 15], [31, 29], [13, 14], [23, 24],
[30, 28], [28, 27], [21, 22]])
st = topology.SingleTopology(mol_a, mol_b, core, ff)
# test that the vjps work
_ = jax.vjp(st.parameterize_harmonic_bond,
ff.hb_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_harmonic_angle,
ff.ha_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_proper_torsion,
ff.pt_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_improper_torsion,
ff.it_handle.params,
has_aux=True)
_ = jax.vjp(st.parameterize_nonbonded,
ff.q_handle.params,
ff.lj_handle.params,
has_aux=True)
def test_bonded(self):
# other bonded terms use an identical protocol, so we assume they're correct if the harmonic bond tests pass.
# 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)
(params_src, params_dst,
params_uni), vjp_fn, (potential_src, potential_dst,
potential_uni) = jax.vjp(
st.parameterize_harmonic_bond,
self.ff.hb_handle.params,
has_aux=True)
# test that vjp_fn works
vjp_fn([
np.random.rand(*params_src.shape),
np.random.rand(*params_dst.shape),
np.random.rand(*params_uni.shape)
])
assert len(potential_src.get_idxs() == 6)
assert len(potential_dst.get_idxs() == 6)
assert len(potential_uni.get_idxs() == 3)
cc = self.ff.hb_handle.lookup_smirks("[#6X3:1]:[#6X3:2]")
cH = self.ff.hb_handle.lookup_smirks("[#6X3:1]-[#1:2]")
cO = self.ff.hb_handle.lookup_smirks("[#6X3:1]-[#8X2H1:2]")
OH = self.ff.hb_handle.lookup_smirks("[#8:1]-[#1:2]")
np.testing.assert_array_equal(params_src, [cc, cc, cc, cc, cc, cc])
np.testing.assert_array_equal(params_dst, [cc, cc, cc, cc, cc, cc])
np.testing.assert_array_equal(params_uni, [cH, cO, OH])
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)
(params_src, params_dst,
params_uni), vjp_fn, (potential_src, potential_dst,
potential_uni) = jax.vjp(
st.parameterize_harmonic_bond,
self.ff.hb_handle.params,
has_aux=True)
assert len(potential_src.get_idxs() == 7)
assert len(potential_dst.get_idxs() == 7)
assert len(potential_uni.get_idxs() == 1)
np.testing.assert_array_equal(params_src, [cc, cc, cc, cc, cc, cc, cH])
np.testing.assert_array_equal(params_dst, [cc, cc, cc, cc, cc, cc, cO])
np.testing.assert_array_equal(params_uni, [OH])
# test that vjp_fn works
vjp_fn([
np.random.rand(*params_src.shape),
np.random.rand(*params_dst.shape),
np.random.rand(*params_uni.shape)
])
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
res = 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 __init__(self, mol_a, mol_b, core, ff):
self.mol_a = mol_a
self.mol_b = mol_b
self.core = core
self.ff = ff
self.top = topology.SingleTopology(mol_a, mol_b, core, ff)