def test_harmonic_bond_singularity(self):
"""Test that two particles sitting directly on top of each other should generate a proper force."""
x = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
params = np.array([[2.0, 0.0]], dtype=np.float64)
bond_idxs = np.array([[0, 1]], dtype=np.int32)
lamb = 0.0
box = np.eye(3) * 100
# specific to harmonic bond force
relative_tolerance_at_precision = {np.float32: 2e-5, np.float64: 1e-9}
for precision, rtol in relative_tolerance_at_precision.items():
test_potential = potentials.HarmonicBond(bond_idxs)
ref_potential = functools.partial(bonded.harmonic_bond, bond_idxs=bond_idxs)
# we assert finite-ness of the forces.
self.compare_forces(x, params, box, lamb, ref_potential, test_potential, rtol, precision=precision)
# test with both zero and non zero terms
x = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
params = np.array([[2.0, 0.0], [2.0, 1.0]], dtype=np.float64)
bond_idxs = np.array([[0, 1], [0, 2]], dtype=np.int32)
# specific to harmonic bond force
relative_tolerance_at_precision = {np.float32: 2e-5, np.float64: 1e-9}
for precision, rtol in relative_tolerance_at_precision.items():
test_potential = potentials.HarmonicBond(bond_idxs)
ref_potential = functools.partial(bonded.harmonic_bond, bond_idxs=bond_idxs)
# we assert finite-ness of the forces.
self.compare_forces(x, params, box, lamb, ref_potential, test_potential, rtol, precision=precision)
def test_harmonic_bond(self, n_particles=64, n_bonds=35, dim=3):
"""Randomly connect pairs of particles, then validate the resulting HarmonicBond force"""
np.random.seed(125) # TODO: where should this seed be set?
x = self.get_random_coords(n_particles, dim)
atom_idxs = np.arange(n_particles)
params = np.random.rand(n_bonds, 2).astype(np.float64)
bond_idxs = []
for _ in range(n_bonds):
bond_idxs.append(np.random.choice(atom_idxs, size=2,
replace=False))
bond_idxs = np.array(bond_idxs, dtype=np.int32)
lamb = 0.0
box = np.eye(3) * 100
# specific to harmonic bond force
relative_tolerance_at_precision = {np.float32: 2e-5, np.float64: 1e-9}
for precision, rtol in relative_tolerance_at_precision.items():
test_potential = potentials.HarmonicBond(bond_idxs)
ref_potential = functools.partial(bonded.harmonic_bond,
bond_idxs=bond_idxs)
self.compare_forces(x,
params,
box,
lamb,
ref_potential,
test_potential,
rtol,
precision=precision)
def get_harmonic_bond(n_atoms, n_bonds):
atom_idxs = np.arange(n_atoms)
params = np.random.rand(n_bonds, 2).astype(np.float64)
bond_idxs = []
for _ in range(n_bonds):
bond_idxs.append(np.random.choice(atom_idxs, size=2, replace=False))
bond_idxs = np.array(bond_idxs, dtype=np.int32)
lamb_mult = np.random.randint(-5, 5, size=n_bonds, dtype=np.int32)
lamb_offset = np.random.randint(-5, 5, size=n_bonds, dtype=np.int32)
return potentials.HarmonicBond(bond_idxs, lamb_mult, lamb_offset), params
def prepare_bonded_system(x, B, A, T, precision):
assert x.ndim == 2
N = x.shape[0]
# D = x.shape[1]
atom_idxs = np.arange(N)
bond_params = np.random.rand(B, 2).astype(np.float64)
bond_idxs = []
for _ in range(B):
bond_idxs.append(np.random.choice(atom_idxs, size=2, replace=False))
bond_idxs = np.array(bond_idxs, dtype=np.int32)
# params = np.concatenate([params, bond_params])
# angle_params = np.random.rand(P_angles).astype(np.float64)
# angle_param_idxs = np.random.randint(low=0, high=P_angles, size=(A,2), dtype=np.int32) + len(params)
# angle_idxs = []
# for _ in range(A):
# angle_idxs.append(np.random.choice(atom_idxs, size=3, replace=False))
# angle_idxs = np.array(angle_idxs, dtype=np.int32)
# params = np.concatenate([params, angle_params])
# torsion_params = np.random.rand(P_torsions).astype(np.float64)
# torsion_param_idxs = np.random.randint(low=0, high=P_torsions, size=(T,3), dtype=np.int32) + len(params)
# torsion_idxs = []
# for _ in range(T):
# torsion_idxs.append(np.random.choice(atom_idxs, size=4, replace=False))
# torsion_idxs = np.array(torsion_idxs, dtype=np.int32)
# params = np.concatenate([params, torsion_params])
print("precision", precision)
custom_bonded = potentials.HarmonicBond(bond_idxs,
bond_params,
precision=precision)
harmonic_bond_fn = functools.partial(bonded.harmonic_bond,
box=None,
bond_idxs=bond_idxs)
# custom_angles = potentials.HarmonicAngle(angle_idxs, angle_param_idxs, precision=precision)
# harmonic_angle_fn = functools.partial(bonded.harmonic_angle, box=None, angle_idxs=angle_idxs, param_idxs=angle_param_idxs)
# custom_torsions = potentials.PeriodicTorsion(torsion_idxs, torsion_param_idxs, precision=precision)
# periodic_torsion_fn = functools.partial(bonded.periodic_torsion, box=None, torsion_idxs=torsion_idxs, param_idxs=torsion_param_idxs)
return (bond_params, harmonic_bond_fn), custom_bonded
def get_harmonic_restraints(n_atoms, n_restraints):
assert n_restraints * 2 <= n_atoms
params = np.random.rand(n_restraints, 2).astype(np.float64)
bond_idxs_src = []
bond_idxs_dst = []
atom_idxs_src = np.arange(n_atoms // 2)
atom_idxs_dst = np.arange(n_atoms // 2) + n_atoms // 2
bond_idxs_src = np.random.choice(atom_idxs_src,
size=n_restraints,
replace=False)
bond_idxs_dst = np.random.choice(atom_idxs_dst,
size=n_restraints,
replace=False)
bond_idxs = np.array([bond_idxs_src, bond_idxs_dst], dtype=np.int32).T
lamb_mult = np.random.randint(-5, 5, size=n_restraints, dtype=np.int32)
lamb_offset = np.random.randint(-5, 5, size=n_restraints, dtype=np.int32)
return potentials.HarmonicBond(bond_idxs, lamb_mult, lamb_offset), params
def parameterize_harmonic_bond(self, ff_params):
params, idxs = self.ff.hb_handle.partial_parameterize(ff_params, self.mol)
return params, potentials.HarmonicBond(idxs)
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
def harmonic_bond():
bond_idxs = np.array([[0, 1], [0, 2]], dtype=np.int32)
params = np.ones(shape=(2, 2), dtype=np.float32)
return potentials.HarmonicBond(bond_idxs).bind(params)
def _futures_a_to_b(self, ff_params, mol_a, mol_b, combined_core_idxs, x0,
box0, prefix, seed):
num_host_atoms = x0.shape[0] - mol_a.GetNumAtoms() - mol_b.GetNumAtoms(
)
# (ytz): super ugly, undo combined_core_idxs to get back original idxs
core_idxs = combined_core_idxs - num_host_atoms
core_idxs[:, 1] -= mol_a.GetNumAtoms()
dual_topology = self.setup_topology(mol_a, mol_b)
rfe = free_energy_rabfe.RelativeFreeEnergy(dual_topology)
unbound_potentials, sys_params, masses = rfe.prepare_host_edge(
ff_params, self.host_system)
k_core = 30.0
core_params = np.zeros_like(combined_core_idxs).astype(np.float64)
core_params[:, 0] = k_core
restraint_potential = potentials.HarmonicBond(combined_core_idxs, )
unbound_potentials.append(restraint_potential)
sys_params.append(core_params)
# tbd sample from boltzmann distribution later
v0 = np.zeros_like(x0)
beta = 1 / (constants.BOLTZ * self.temperature)
bond_list = np.concatenate(
[unbound_potentials[0].get_idxs(), core_idxs])
masses = model_utils.apply_hmr(masses, bond_list)
friction = 1.0
integrator = LangevinIntegrator(self.temperature, self.dt, friction,
masses, seed)
bond_list = list(map(tuple, bond_list))
group_indices = get_group_indices(bond_list)
barostat_interval = 5
barostat = MonteCarloBarostat(x0.shape[0], self.pressure,
self.temperature, group_indices,
barostat_interval, seed)
endpoint_correct = True
model = estimator_abfe.FreeEnergyModel(
unbound_potentials,
endpoint_correct,
self.client,
box0, # important, use equilibrated box.
x0,
v0,
integrator,
barostat,
self.host_schedule,
self.equil_steps,
self.prod_steps,
beta,
prefix,
)
bound_potentials = []
for params, unbound_pot in zip(sys_params, model.unbound_potentials):
bp = unbound_pot.bind(np.asarray(params))
bound_potentials.append(bp)
all_args = []
for lamb_idx, lamb in enumerate(model.lambda_schedule):
subsample_interval = 1000
all_args.append((
lamb,
model.box,
model.x0,
model.v0,
bound_potentials,
model.integrator,
model.barostat,
model.equil_steps,
model.prod_steps,
subsample_interval,
subsample_interval,
model.lambda_schedule,
))
if endpoint_correct:
assert isinstance(bound_potentials[-1], potentials.HarmonicBond)
all_args.append((
1.0,
model.box,
model.x0,
model.v0,
bound_potentials[:-1], # strip out the restraints
model.integrator,
model.barostat,
model.equil_steps,
model.prod_steps,
subsample_interval,
subsample_interval,
[], # no need to evaluate Us for the endpoint correction
))
futures = []
if self.client is None:
for args in all_args:
futures.append(_MockFuture(estimator_abfe.simulate(*args)))
else:
for args in all_args:
futures.append(
self.client.submit(estimator_abfe.simulate, *args))
return sys_params, model, futures
def from_rdkit(cls, mol, ff_handlers):
"""
Initialize a system from an RDKit ROMol.
Parameters
----------
mol: Chem.ROMol
RDKit ROMol. Should have graphical hydrogens in the topology.
ff_handlers: list of forcefield handlers.
openforcefield small molecule handlers.
"""
masses = np.array([a.GetMass() for a in mol.GetAtoms()], dtype=np.float64)
bound_potentials = []
for handle in ff_handlers:
results = handle.parameterize(mol)
if isinstance(handle, bonded.HarmonicBondHandler):
bond_params, bond_idxs = results
bound_potentials.append(potentials.HarmonicBond(bond_idxs).bind(bond_params))
elif isinstance(handle, bonded.HarmonicAngleHandler):
angle_params, angle_idxs = results
bound_potentials.append(potentials.HarmonicAngle(angle_idxs).bind(angle_params))
elif isinstance(handle, bonded.ProperTorsionHandler):
torsion_params, torsion_idxs = results
bound_potentials.append(potentials.PeriodicTorsion(torsion_idxs).bind(torsion_params))
elif isinstance(handle, bonded.ImproperTorsionHandler):
torsion_params, torsion_idxs = results
bound_potentials.append(potentials.PeriodicTorsion(torsion_idxs).bind(torsion_params))
elif isinstance(handle, nonbonded.AM1CCCHandler):
charge_handle = handle
charge_params = results
elif isinstance(handle, nonbonded.LennardJonesHandler):
lj_params = results
lj_handle = handle
else:
print("WARNING: skipping handler", handle)
pass
lambda_plane_idxs = np.zeros(len(masses), dtype=np.int32)
lambda_offset_idxs = np.zeros(len(masses), dtype=np.int32)
exclusion_idxs, scale_factors = nonbonded.generate_exclusion_idxs(
mol,
scale12=1.0,
scale13=1.0,
scale14=0.5
)
# use same scale factors until we modify 1-4s for electrostatics
scale_factors = np.stack([scale_factors, scale_factors], axis=1)
# (ytz) fix this later to not be so hard coded
alpha = 2.0 # same as ewald alpha
cutoff = 1.0 # nonbonded cutoff
qlj_params = jnp.concatenate([
jnp.reshape(charge_params, (-1, 1)),
jnp.reshape(lj_params, (-1, 2))
], axis=1)
bound_potentials.append(potentials.Nonbonded(
exclusion_idxs,
scale_factors,
lambda_plane_idxs,
lambda_offset_idxs,
alpha,
cutoff).bind(qlj_params))
return cls(masses, bound_potentials)
def deserialize_system(system, cutoff):
"""
Deserialize an OpenMM XML file
Parameters
----------
system: openmm.System
A system object to be deserialized
Returns
-------
list of lib.Potential, masses
Note: We add a small epsilon (1e-3) to all zero eps values to prevent
a singularity from occuring in the lennard jones derivatives
"""
masses = []
for p in range(system.getNumParticles()):
masses.append(value(system.getParticleMass(p)))
N = len(masses)
# this should not be a dict since we may have more than one instance of a given
# force.
bps = []
for force in system.getForces():
if isinstance(force, mm.HarmonicBondForce):
bond_idxs = []
bond_params = []
for b_idx in range(force.getNumBonds()):
src_idx, dst_idx, length, k = force.getBondParameters(b_idx)
length = value(length)
k = value(k)
bond_idxs.append([src_idx, dst_idx])
bond_params.append((k, length))
bond_idxs = np.array(bond_idxs, dtype=np.int32)
bond_params = np.array(bond_params, dtype=np.float64)
bps.append(potentials.HarmonicBond(bond_idxs).bind(bond_params))
if isinstance(force, mm.HarmonicAngleForce):
angle_idxs = []
angle_params = []
for a_idx in range(force.getNumAngles()):
src_idx, mid_idx, dst_idx, angle, k = force.getAngleParameters(
a_idx)
angle = value(angle)
k = value(k)
angle_idxs.append([src_idx, mid_idx, dst_idx])
angle_params.append((k, angle))
angle_idxs = np.array(angle_idxs, dtype=np.int32)
angle_params = np.array(angle_params, dtype=np.float64)
bps.append(potentials.HarmonicAngle(angle_idxs).bind(angle_params))
if isinstance(force, mm.PeriodicTorsionForce):
torsion_idxs = []
torsion_params = []
for t_idx in range(force.getNumTorsions()):
a_idx, b_idx, c_idx, d_idx, period, phase, k = force.getTorsionParameters(
t_idx)
phase = value(phase)
k = value(k)
torsion_params.append((k, phase, period))
torsion_idxs.append([a_idx, b_idx, c_idx, d_idx])
torsion_idxs = np.array(torsion_idxs, dtype=np.int32)
torsion_params = np.array(torsion_params, dtype=np.float64)
bps.append(
potentials.PeriodicTorsion(torsion_idxs).bind(torsion_params))
if isinstance(force, mm.NonbondedForce):
num_atoms = force.getNumParticles()
charge_params = []
lj_params = []
for a_idx in range(num_atoms):
charge, sig, eps = force.getParticleParameters(a_idx)
charge = value(charge) * np.sqrt(constants.ONE_4PI_EPS0)
sig = value(sig)
eps = value(eps)
# increment eps by 1e-3 if we have eps==0 to avoid a singularity in parameter derivatives
# override default amber types
# this doesn't work for water!
# if eps == 0:
# print("Warning: overriding eps by 1e-3 to avoid a singularity")
# eps += 1e-3
# charge_params.append(charge_idx)
charge_params.append(charge)
lj_params.append((sig, eps))
charge_params = np.array(charge_params, dtype=np.float64)
# print("Protein net charge:", np.sum(np.array(global_params)[charge_param_idxs]))
lj_params = np.array(lj_params, dtype=np.float64)
# 1 here means we fully remove the interaction
# 1-2, 1-3
# scale_full = insert_parameters(1.0, 20)
# 1-4, remove half of the interaction
# scale_half = insert_parameters(0.5, 21)
exclusion_idxs = []
scale_factors = []
all_sig = lj_params[:, 0]
all_eps = lj_params[:, 1]
# validate exclusions/exceptions to make sure they make sense
for a_idx in range(force.getNumExceptions()):
# tbd charge scale factors
src, dst, new_cp, new_sig, new_eps = force.getExceptionParameters(
a_idx)
new_sig = value(new_sig)
new_eps = value(new_eps)
src_sig = all_sig[src]
dst_sig = all_sig[dst]
src_eps = all_eps[src]
dst_eps = all_eps[dst]
expected_sig = (src_sig + dst_sig) / 2
expected_eps = np.sqrt(src_eps * dst_eps)
exclusion_idxs.append([src, dst])
# sanity check this (expected_eps can be zero), redo this thing
# the lj_scale factor measures how much we *remove*
if expected_eps == 0:
if new_eps == 0:
lj_scale_factor = 1
else:
raise RuntimeError(
"Divide by zero in epsilon calculation")
else:
lj_scale_factor = 1 - new_eps / expected_eps
scale_factors.append(lj_scale_factor)
# tbd fix charge_scale_factors using new_cp
if new_eps != 0:
np.testing.assert_almost_equal(expected_sig, new_sig)
exclusion_idxs = np.array(exclusion_idxs, dtype=np.int32)
lambda_plane_idxs = np.zeros(N, dtype=np.int32)
lambda_offset_idxs = np.zeros(N, dtype=np.int32)
# cutoff = 1000.0
nb_params = np.concatenate(
[np.expand_dims(charge_params, axis=1), lj_params], axis=1)
# optimizations
nb_params[:, 1] = nb_params[:, 1] / 2
nb_params[:, 2] = np.sqrt(nb_params[:, 2])
beta = 2.0 # erfc correction
# use the same scale factors for electrostatics and lj
scale_factors = np.stack([scale_factors, scale_factors], axis=1)
bps.append(
potentials.Nonbonded(exclusion_idxs, scale_factors,
lambda_plane_idxs, lambda_offset_idxs,
beta, cutoff).bind(nb_params))
# nrg_fns.append(('Exclusions', (exclusion_idxs, scale_factors, es_scale_factors)))
return bps, masses