def benchmark_dhfr():
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)
seed = 1234
dt = 1.5e-3
intg = LangevinIntegrator(
300,
dt,
1.0,
np.array(host_masses),
seed
).impl()
bps = []
for potential in host_fns:
bps.append(potential.bound_impl(precision=np.float32)) # get the bound implementation
x0 = host_conf
v0 = np.zeros_like(host_conf)
ctxt = custom_ops.Context(
x0,
v0,
box,
intg,
bps
)
# initialize observables
obs = []
for bp in bps:
du_dp_obs = custom_ops.AvgPartialUPartialParam(bp, 100)
ctxt.add_observable(du_dp_obs)
obs.append(du_dp_obs)
lamb = 0.0
start = time.time()
# num_steps = 50000
num_steps = 50000
# num_steps = 10
writer = PDBWriter([host_pdb.topology], "dhfr.pdb")
for step in range(num_steps):
ctxt.step(lamb)
if step % 1000 == 0:
delta = time.time()-start
steps_per_second = step/delta
seconds_per_day = 86400
steps_per_day = steps_per_second*seconds_per_day
ps_per_day = dt*steps_per_day
ns_per_day = ps_per_day*1e-3
print(step, "ns/day", ns_per_day)
# coords = recenter(ctxt.get_x_t(), box)
# writer.write_frame(coords*10)
print("total time", time.time() - start)
writer.close()
# bond angle torsions nonbonded
for potential, du_dp_obs in zip(host_fns, obs):
dp = du_dp_obs.avg_du_dp()
print(potential, dp.shape)
print(dp)
def test_benchmark(self):
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)
for f in host_fns:
if isinstance(f, potentials.Nonbonded):
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)
beta = 2.0
cutoff = 1.1
lamb = 0.0
N = host_conf.shape[0]
test_conf = host_conf[:N]
# process exclusions
test_exclusions = []
test_scales = []
for (i, j), (sa, sb) in zip(nonbonded_fn.get_exclusion_idxs(),
nonbonded_fn.get_scale_factors()):
if i < N and j < N:
test_exclusions.append((i, j))
test_scales.append((sa, sb))
test_exclusions = np.array(test_exclusions, dtype=np.int32)
test_scales = np.array(test_scales, dtype=np.float64)
test_params = nonbonded_fn.params[:N, :]
test_lambda_plane_idxs = np.zeros(N, dtype=np.int32)
test_lambda_offset_idxs = np.zeros(N, dtype=np.int32)
test_nonbonded_fn = potentials.Nonbonded(test_exclusions, test_scales,
test_lambda_plane_idxs,
test_lambda_offset_idxs, beta,
cutoff)
precision = np.float32
impl = test_nonbonded_fn.unbound_impl(precision)
for _ in range(100):
impl.execute_du_dx(test_conf, test_params, box, lamb)
def run_simulation(mol_name, mol, smirnoff_params, pdb, inference):
omm_forcefield = app.ForceField('amber99sb.xml', 'tip3p.xml')
system = omm_forcefield.createSystem(
pdb.topology,
nonbondedMethod=app.NoCutoff,
constraints=None,
rigidWater=False)
coords = []
for x,y,z in pdb.positions:
coords.append([to_md_units(x),to_md_units(y),to_md_units(z)])
coords = np.array(coords)
host_conf = coords
host_potentials, (host_params, host_param_groups), host_masses = serialize.deserialize_system(system)
smirnoff = ForceField("test_forcefields/smirnoff99Frosst.offxml")
# CAREFUL THIS SHOULD BE THE SINGLE SET OF CONSISTENT FORCEFIELDS
am1 = True
if am1:
guest_potentials, smirnoff_params, smirnoff_param_groups, guest_conf, guest_masses = forcefield.parameterize(mol, smirnoff, am1)
else:
guest_potentials, _, smirnoff_param_groups, guest_conf, guest_masses = forcefield.parameterize(mol, smirnoff, am1)
# guest_conf = rescale_and_center(guest_conf, scale_factor=3.0)
combined_potentials, combined_params, combined_param_groups, combined_conf, combined_masses = forcefield.combiner(
host_potentials, guest_potentials,
host_params, smirnoff_params,
host_param_groups, smirnoff_param_groups,
host_conf, guest_conf,
host_masses, guest_masses)
num_host_atoms = host_conf.shape[0]
def filter_groups(param_groups, groups):
roll = np.zeros_like(param_groups)
for g in groups:
roll = np.logical_or(roll, param_groups == g)
return roll
# host_dp_idxs = np.argwhere(filter_groups(host_param_groups, [7])).reshape(-1)
# guest_dp_idxs = np.argwhere(filter_groups(smirnoff_param_groups, [7])).reshape(-1)
# 1. host bond lengths
# 7. host charges
# 8. host vdw sigma
# 9. host vdw epsilon
# 17. charge
combined_dp_idxs = np.argwhere(filter_groups(combined_param_groups, [17])).reshape(-1)
if inference:
combined_dp_idxs = []
outfile = open("frames/"+mol_name+".pdb", "w")
combined_pdb = Chem.CombineMols(Chem.MolFromPDBFile(pdb.filepath, removeHs=False), mol)
combined_pdb_str = StringIO(Chem.MolToPDBBlock(combined_pdb))
cpdb = app.PDBFile(combined_pdb_str)
PDBFile.writeHeader(cpdb.topology, outfile)
def write_fn(x, frame_idx):
PDBFile.writeModel(cpdb.topology, x, outfile, frame_idx)
dG, dG_grads, all_xis = minimize(
mol_name,
num_host_atoms,
combined_potentials,
combined_params,
combined_conf,
combined_masses,
combined_dp_idxs,
writer_fn=write_fn
)
PDBFile.writeFooter(cpdb.topology, outfile)
outfile.flush()
# print("CDPI", combined_dp_idxs)
# (ytz): we need to offset this by number of host params to compute the indices
# of the original ligand parameters.
if not inference:
ligand_dp_idxs = combined_dp_idxs - len(host_params)
else:
ligand_dp_idxs = None
return dG, dG_grads, ligand_dp_idxs
def test_dhfr(self):
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)
for f in host_fns:
if isinstance(f, potentials.Nonbonded):
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)
beta = 2.0
cutoff = 1.1
lamb = 0.1
max_N = host_conf.shape[0]
for N in [33, 65, 231, 1050, 4080]:
print("N", N)
test_conf = host_conf[:N]
# process exclusions
test_exclusions = []
test_scales = []
for (i, j), (sa, sb) in zip(nonbonded_fn.get_exclusion_idxs(),
nonbonded_fn.get_scale_factors()):
if i < N and j < N:
test_exclusions.append((i, j))
test_scales.append((sa, sb))
test_exclusions = np.array(test_exclusions, dtype=np.int32)
test_scales = np.array(test_scales, dtype=np.float64)
test_params = nonbonded_fn.params[:N, :]
test_lambda_plane_idxs = np.random.randint(low=-2,
high=2,
size=N,
dtype=np.int32)
test_lambda_offset_idxs = np.random.randint(low=-2,
high=2,
size=N,
dtype=np.int32)
test_nonbonded_fn = potentials.Nonbonded(test_exclusions,
test_scales,
test_lambda_plane_idxs,
test_lambda_offset_idxs,
beta, cutoff)
ref_nonbonded_fn = prepare_reference_nonbonded(
test_params, test_exclusions, test_scales,
test_lambda_plane_idxs, test_lambda_offset_idxs, beta, cutoff)
for precision, rtol in [(np.float64, 1e-8), (np.float32, 1e-4)]:
self.compare_forces(test_conf,
test_params,
box,
lamb,
ref_nonbonded_fn,
test_nonbonded_fn,
rtol,
precision=precision)
def pose_dock(
guests_sdfile,
host_pdbfile,
transition_type,
n_steps,
transition_steps,
max_lambda,
outdir,
random_rotation=False,
constant_atoms=[],
):
"""Runs short simulations in which the guests phase in or out over time
Parameters
----------
guests_sdfile: path to input sdf with guests to pose/dock
host_pdbfile: path to host pdb file to dock into
transition_type: "insertion" or "deletion"
n_steps: how many total steps of simulation to do (recommended: <= 1000)
transition_steps: how many steps to insert/delete the guest over (recommended: <= 500)
(must be <= n_steps)
max_lambda: lambda value the guest should insert from or delete to
(recommended: 1.0 for work calulation, 0.25 to stay close to original pose)
(must be =1 for work calculation to be applicable)
outdir: where to write output (will be created if it does not already exist)
random_rotation: whether to apply a random rotation to each guest before inserting
constant_atoms: atom numbers from the host_pdbfile to hold mostly fixed across the simulation
(1-indexed, like PDB files)
Output
------
A pdb & sdf file every 100 steps (outdir/<guest_name>_<step>.pdb)
stdout every 100 steps noting the step number, lambda value, and energy
stdout for each guest noting the work of transition
stdout for each guest noting how long it took to run
Note
----
If any norm of force per atom exceeds 20000 kJ/(mol*nm) [MAX_NORM_FORCE defined in docking/report.py],
the simulation for that guest will stop and the work will not be calculated.
"""
assert transition_steps <= n_steps
assert transition_type in ("insertion", "deletion")
if random_rotation:
assert transition_type == "insertion"
if not os.path.exists(outdir):
os.makedirs(outdir)
host_mol = Chem.MolFromPDBFile(host_pdbfile, removeHs=False)
amber_ff = app.ForceField("amber99sbildn.xml", "tip3p.xml")
host_file = PDBFile(host_pdbfile)
host_system = amber_ff.createSystem(
host_file.topology,
nonbondedMethod=app.NoCutoff,
constraints=None,
rigidWater=False,
)
host_conf = []
for x, y, z in host_file.positions:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
final_potentials = []
host_potentials, host_masses = openmm_deserializer.deserialize_system(
host_system, cutoff=1.2)
host_nb_bp = None
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:
final_potentials.append(bp)
# TODO (ytz): we should really fix this later on. This padding was done to
# address the particles that are too close to the boundary.
padding = 0.1
box_lengths = np.amax(host_conf, axis=0) - np.amin(host_conf, axis=0)
box_lengths = box_lengths + padding
box = np.eye(3, dtype=np.float64) * box_lengths
suppl = Chem.SDMolSupplier(guests_sdfile, removeHs=False)
for guest_mol in suppl:
start_time = time.time()
guest_name = guest_mol.GetProp("_Name")
guest_ff_handlers = deserialize_handlers(
open(
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"..",
"ff/params/smirnoff_1_1_0_ccc.py",
)).read())
ff = Forcefield(guest_ff_handlers)
guest_base_topology = topology.BaseTopology(guest_mol, ff)
# combine
hgt = topology.HostGuestTopology(host_nb_bp, guest_base_topology)
# setup the parameter handlers for the ligand
bonded_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]]
these_potentials = list(final_potentials)
# instantiate the vjps while parameterizing (forward pass)
for fn, handle in bonded_tuples:
params, potential = fn(handle.params)
these_potentials.append(potential.bind(params))
nb_params, nb_potential = hgt.parameterize_nonbonded(
ff.q_handle.params, ff.lj_handle.params)
these_potentials.append(nb_potential.bind(nb_params))
bps = these_potentials
guest_masses = [a.GetMass() for a in guest_mol.GetAtoms()]
masses = np.concatenate([host_masses, guest_masses])
for atom_num in constant_atoms:
masses[atom_num - 1] += 50000
conformer = guest_mol.GetConformer(0)
mol_conf = np.array(conformer.GetPositions(), dtype=np.float64)
mol_conf = mol_conf / 10 # convert to md_units
if random_rotation:
center = np.mean(mol_conf, axis=0)
mol_conf -= center
from scipy.stats import special_ortho_group
mol_conf = np.matmul(mol_conf, special_ortho_group.rvs(3))
mol_conf += center
x0 = np.concatenate([host_conf, mol_conf]) # combined geometry
v0 = np.zeros_like(x0)
seed = 2021
intg = LangevinIntegrator(300, 1.5e-3, 1.0, masses, seed).impl()
impls = []
precision = np.float32
for b in bps:
p_impl = b.bound_impl(precision)
impls.append(p_impl)
ctxt = custom_ops.Context(x0, v0, box, intg, impls)
# collect a du_dl calculation once every other step
subsample_freq = 2
du_dl_obs = custom_ops.FullPartialUPartialLambda(impls, subsample_freq)
ctxt.add_observable(du_dl_obs)
if transition_type == "insertion":
new_lambda_schedule = np.concatenate([
np.linspace(max_lambda, 0.0, transition_steps),
np.zeros(n_steps - transition_steps),
])
elif transition_type == "deletion":
new_lambda_schedule = np.concatenate([
np.linspace(0.0, max_lambda, transition_steps),
np.ones(n_steps - transition_steps) * max_lambda,
])
else:
raise (RuntimeError(
'invalid `transition_type` (must be one of ["insertion", "deletion"])'
))
calc_work = True
for step, lamb in enumerate(new_lambda_schedule):
ctxt.step(lamb)
if step % 100 == 0:
report.report_step(ctxt, step, lamb, box, bps, impls,
guest_name, n_steps, 'pose_dock')
host_coords = ctxt.get_x_t()[:len(host_conf)] * 10
guest_coords = ctxt.get_x_t()[len(host_conf):] * 10
report.write_frame(host_coords, host_mol, guest_coords,
guest_mol, guest_name, outdir, step, 'pd')
if step in (0, int(n_steps / 2), n_steps - 1):
if report.too_much_force(ctxt, lamb, box, bps, impls):
calc_work = False
break
# Note: this condition only applies for ABFE, not RBFE
if (abs(du_dl_obs.full_du_dl()[0]) > 0.001
or abs(du_dl_obs.full_du_dl()[-1]) > 0.001):
print("Error: du_dl endpoints are not ~0")
calc_work = False
if calc_work:
work = np.trapz(du_dl_obs.full_du_dl(),
new_lambda_schedule[::subsample_freq])
print(f"guest_name: {guest_name}\twork: {work:.2f}")
end_time = time.time()
print(f"{guest_name} took {(end_time - start_time):.2f} seconds")
def create_system(guest_mol, host_pdb, handlers, stage, core_atoms,
restr_force, restr_alpha, restr_count):
"""
Initialize a self-encompassing System object that we can serialize and simulate.
Parameters
----------
guest_mol: rdkit.ROMol
protein: openmm.System
"""
# host_system = protein_system
guest_masses = np.array([a.GetMass() for a in guest_mol.GetAtoms()],
dtype=np.float64)
amber_ff = app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')
host_system = amber_ff.createSystem(host_pdb.topology,
nonbondedMethod=app.NoCutoff,
constraints=None,
rigidWater=False)
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system)
num_host_atoms = len(host_masses)
num_guest_atoms = guest_mol.GetNumAtoms()
# Name, Args, vjp_fn
final_gradients = []
final_vjp_fns = []
for item in host_fns:
if item[0] == 'LennardJones':
host_lj_params = item[1]
elif item[0] == 'Charges':
host_charge_params = item[1]
elif item[0] == 'GBSA':
host_gb_params = item[1][0]
host_gb_props = item[1][1:]
elif item[0] == 'Exclusions':
host_exclusions = item[1]
else:
final_gradients.append((item[0], item[1]))
final_vjp_fns.append(None)
# print("Ligand A Name:", a_name)
guest_exclusion_idxs, guest_scales = nonbonded.generate_exclusion_idxs(
guest_mol, scale12=1.0, scale13=1.0, scale14=0.5)
guest_exclusion_idxs += num_host_atoms
guest_lj_exclusion_scales = guest_scales
guest_charge_exclusion_scales = guest_scales
host_exclusion_idxs = host_exclusions[0]
host_lj_exclusion_scales = host_exclusions[1]
host_charge_exclusion_scales = host_exclusions[2]
combined_exclusion_idxs = np.concatenate(
[host_exclusion_idxs, guest_exclusion_idxs])
combined_lj_exclusion_scales = np.concatenate(
[host_lj_exclusion_scales, guest_lj_exclusion_scales])
combined_charge_exclusion_scales = np.concatenate(
[host_charge_exclusion_scales, guest_charge_exclusion_scales])
# handler_vjps = []
for handle in handlers:
results = handle.parameterize(guest_mol)
if isinstance(handle, bonded.HarmonicBondHandler):
bond_idxs, (bond_params, bond_vjp_fn) = results
bond_idxs += num_host_atoms
final_gradients.append(("HarmonicBond", (bond_idxs, bond_params)))
final_vjp_fns.append((bond_vjp_fn))
# handler_vjps.append(bond_vjp_fn)
elif isinstance(handle, bonded.HarmonicAngleHandler):
angle_idxs, (angle_params, angle_vjp_fn) = results
angle_idxs += num_host_atoms
final_gradients.append(
("HarmonicAngle", (angle_idxs, angle_params)))
final_vjp_fns.append(angle_vjp_fn)
# handler_vjps.append(angle_vjp_fn)
elif isinstance(handle, bonded.ProperTorsionHandler):
torsion_idxs, (torsion_params, torsion_vjp_fn) = results
torsion_idxs += num_host_atoms
final_gradients.append(
("PeriodicTorsion", (torsion_idxs, torsion_params)))
final_vjp_fns.append(torsion_vjp_fn)
# handler_vjps.append(torsion_vjp_fn)
# guest_vjp_fns.append(torsion_vjp_fn)
elif isinstance(handle, bonded.ImproperTorsionHandler):
torsion_idxs, (torsion_params, torsion_vjp_fn) = results
torsion_idxs += num_host_atoms
final_gradients.append(
("PeriodicTorsion", (torsion_idxs, torsion_params)))
final_vjp_fns.append(torsion_vjp_fn)
# handler_vjps.append(torsion_vjp_fn)
elif isinstance(handle, nonbonded.LennardJonesHandler):
guest_lj_params, guest_lj_vjp_fn = results
combined_lj_params, combined_lj_vjp_fn = concat_with_vjps(
host_lj_params, guest_lj_params, None, guest_lj_vjp_fn)
# handler_vjps.append(lj_adjoint_fn)
elif isinstance(handle, nonbonded.SimpleChargeHandler):
guest_charge_params, guest_charge_vjp_fn = results
combined_charge_params, combined_charge_vjp_fn = concat_with_vjps(
host_charge_params, guest_charge_params, None,
guest_charge_vjp_fn)
# handler_vjps.append(charge_adjoint_fn)
elif isinstance(handle, nonbonded.GBSAHandler):
guest_gb_params, guest_gb_vjp_fn = results
combined_gb_params, combined_gb_vjp_fn = concat_with_vjps(
host_gb_params, guest_gb_params, None, guest_gb_vjp_fn)
# handler_vjps.append(gb_adjoint_fn)
elif isinstance(handle, nonbonded.AM1BCCHandler):
# ill defined behavior if both SimpleChargeHandler and AM1Handler is present
guest_charge_params, guest_charge_vjp_fn = results
combined_charge_params, combined_charge_vjp_fn = concat_with_vjps(
host_charge_params, guest_charge_params, None,
guest_charge_vjp_fn)
# handler_vjps.append(gb_adjoint_fn)
elif isinstance(handle, nonbonded.AM1CCCHandler):
guest_charge_params, guest_charge_vjp_fn = results
combined_charge_params, combined_charge_vjp_fn = concat_with_vjps(
host_charge_params, guest_charge_params, None,
guest_charge_vjp_fn)
# handler_vjps.append(gb_adjoint_fn)
else:
raise Exception("Unknown Handler", handle)
# (use the below vjps for correctness)
# combined_charge_params, charge_adjoint_fn = concat_with_vjps(host_charge_params, guest_charge_params, None, guest_charge_vjp_fn)
# combined_lj_params, lj_adjoint_fn = concat_with_vjps(host_lj_params, guest_lj_params, None, guest_lj_vjp_fn)
# combined_gb_params, gb_adjoint_fn = concat_with_vjps(host_gb_params, guest_gb_params, None, guest_gb_vjp_fn)
# WIP
N_C = num_host_atoms + num_guest_atoms
N_A = num_host_atoms
if stage == 0:
combined_lambda_plane_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs = np.zeros(N_C, dtype=np.int32)
elif stage == 1:
combined_lambda_plane_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs[N_A:] = 1
elif stage == 2:
combined_lambda_plane_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_plane_idxs[N_A:] = 1
combined_lambda_offset_idxs = np.zeros(N_C, dtype=np.int32)
# assert 0
cutoff = 100000.0
final_gradients.append(
('Nonbonded',
(np.asarray(combined_charge_params), np.asarray(combined_lj_params),
combined_exclusion_idxs, combined_charge_exclusion_scales,
combined_lj_exclusion_scales, combined_lambda_plane_idxs,
combined_lambda_offset_idxs, cutoff)))
final_vjp_fns.append((combined_charge_vjp_fn, combined_lj_vjp_fn))
final_gradients.append(
('GBSA', (np.asarray(combined_charge_params),
np.asarray(combined_gb_params), combined_lambda_plane_idxs,
combined_lambda_offset_idxs, *host_gb_props, cutoff,
cutoff)))
final_vjp_fns.append((combined_charge_vjp_fn, combined_gb_vjp_fn))
host_conf = []
for x, y, z in host_pdb.positions:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
conformer = guest_mol.GetConformer(0)
mol_a_conf = np.array(conformer.GetPositions(), dtype=np.float64)
mol_a_conf = mol_a_conf / 10 # convert to md_units
x0 = np.concatenate([host_conf, mol_a_conf]) # combined geometry
v0 = np.zeros_like(x0)
# build restraints using the coordinates
backbone_atoms = []
for r_idx, residue in enumerate(host_pdb.getTopology().residues()):
for a in residue.atoms():
if a.name == 'CA':
backbone_atoms.append(a.index)
final_gradients.append(
setup_core_restraints(restr_force,
restr_alpha,
restr_count,
x0,
num_host_atoms,
core_atoms,
backbone_atoms,
stage=stage))
final_vjp_fns.append(None)
combined_masses = np.concatenate([host_masses, guest_masses])
return x0, combined_masses, final_gradients, final_vjp_fns
def create_system(guest_mol, host_pdb, handlers):
"""
Initialize a self-encompassing System object that we can serialize and simulate.
Parameters
----------
guest_mol: rdkit.ROMol
guest molecule
host_pdb: openmm.PDBFile
host system from OpenMM
handlers: list of timemachine.ops.Gradients
forcefield handlers used to parameterize the small molecule
Returns
-------
3-tuple
x0, combined_masses, final_gradients
"""
guest_masses = np.array([a.GetMass() for a in guest_mol.GetAtoms()],
dtype=np.float64)
amber_ff = app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')
host_system = amber_ff.createSystem(host_pdb.topology,
nonbondedMethod=app.NoCutoff,
constraints=None,
rigidWater=False)
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system)
num_host_atoms = len(host_masses)
num_guest_atoms = guest_mol.GetNumAtoms()
# Name, Args, vjp_fn
final_gradients = []
for item in host_fns:
if item[0] == 'LennardJones':
host_lj_params = item[1]
elif item[0] == 'Charges':
host_charge_params = item[1]
elif item[0] == 'GBSA':
host_gb_params = item[1][0]
host_gb_props = item[1][1:]
elif item[0] == 'Exclusions':
host_exclusions = item[1]
else:
final_gradients.append((item[0], item[1]))
guest_exclusion_idxs, guest_scales = nonbonded.generate_exclusion_idxs(
guest_mol, scale12=1.0, scale13=1.0, scale14=0.5)
guest_exclusion_idxs += num_host_atoms
guest_lj_exclusion_scales = guest_scales
guest_charge_exclusion_scales = guest_scales
host_exclusion_idxs = host_exclusions[0]
host_lj_exclusion_scales = host_exclusions[1]
host_charge_exclusion_scales = host_exclusions[2]
combined_exclusion_idxs = np.concatenate(
[host_exclusion_idxs, guest_exclusion_idxs])
combined_lj_exclusion_scales = np.concatenate(
[host_lj_exclusion_scales, guest_lj_exclusion_scales])
combined_charge_exclusion_scales = np.concatenate(
[host_charge_exclusion_scales, guest_charge_exclusion_scales])
for handle in handlers:
results = handle.parameterize(guest_mol)
if isinstance(handle, bonded.HarmonicBondHandler):
bond_idxs, (bond_params, _) = results
bond_idxs += num_host_atoms
final_gradients.append(("HarmonicBond", (bond_idxs, bond_params)))
elif isinstance(handle, bonded.HarmonicAngleHandler):
angle_idxs, (angle_params, _) = results
angle_idxs += num_host_atoms
final_gradients.append(
("HarmonicAngle", (angle_idxs, angle_params)))
elif isinstance(handle, bonded.ProperTorsionHandler):
torsion_idxs, (torsion_params, _) = results
torsion_idxs += num_host_atoms
final_gradients.append(
("PeriodicTorsion", (torsion_idxs, torsion_params)))
elif isinstance(handle, bonded.ImproperTorsionHandler):
torsion_idxs, (torsion_params, _) = results
torsion_idxs += num_host_atoms
final_gradients.append(
("PeriodicTorsion", (torsion_idxs, torsion_params)))
elif isinstance(handle, nonbonded.LennardJonesHandler):
guest_lj_params, _ = results
combined_lj_params = np.concatenate(
[host_lj_params, guest_lj_params])
elif isinstance(handle, nonbonded.SimpleChargeHandler):
guest_charge_params, _ = results
combined_charge_params = np.concatenate(
[host_charge_params, guest_charge_params])
elif isinstance(handle, nonbonded.GBSAHandler):
guest_gb_params, _ = results
combined_gb_params = np.concatenate(
[host_gb_params, guest_gb_params])
elif isinstance(handle, nonbonded.AM1BCCHandler):
guest_charge_params, _ = results
combined_charge_params = np.concatenate(
[host_charge_params, guest_charge_params])
elif isinstance(handle, nonbonded.AM1CCCHandler):
guest_charge_params, _ = results
combined_charge_params = np.concatenate(
[host_charge_params, guest_charge_params])
else:
raise Exception("Unknown Handler", handle)
host_conf = []
for x, y, z in host_pdb.positions:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
conformer = guest_mol.GetConformer(0)
mol_a_conf = np.array(conformer.GetPositions(), dtype=np.float64)
mol_a_conf = mol_a_conf / 10 # convert to md_units
center = np.mean(mol_a_conf, axis=0)
mol_a_conf -= center
from scipy.stats import special_ortho_group
mol_a_conf = np.matmul(mol_a_conf, special_ortho_group.rvs(3))
mol_a_conf += center
# assert 0
x0 = np.concatenate([host_conf, mol_a_conf]) # combined geometry
v0 = np.zeros_like(x0)
N_C = num_host_atoms + num_guest_atoms
N_A = num_host_atoms
cutoff = 100000.0
combined_lambda_plane_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs[num_host_atoms:] = 1
final_gradients.append(
('Nonbonded',
(np.asarray(combined_charge_params), np.asarray(combined_lj_params),
combined_exclusion_idxs, combined_charge_exclusion_scales,
combined_lj_exclusion_scales, combined_lambda_plane_idxs,
combined_lambda_offset_idxs, cutoff)))
final_gradients.append(
('GBSA', (np.asarray(combined_charge_params),
np.asarray(combined_gb_params), combined_lambda_plane_idxs,
combined_lambda_offset_idxs, *host_gb_props, cutoff,
cutoff)))
combined_masses = np.concatenate([host_masses, guest_masses])
return x0, combined_masses, final_gradients
def create_system(guest_mol, host_pdb, handlers, restr_search_radius,
restr_force_constant, intg_temperature, stage):
"""
Initialize a self-encompassing System object that we can serialize and simulate.
Parameters
----------
guest_mol: rdkit.ROMol
guest molecule
host_pdb: openmm.PDBFile
host system from OpenMM
handlers: list of timemachine.ops.Gradients
forcefield handlers used to parameterize the system
restr_search_radius: float
how far away we search from the ligand to define the binding pocket atoms.
restr_force_constant: float
strength of the harmonic oscillator for the restraint
intg_temperature: float
temperature of the integrator in Kelvin
stage: int (0 or 1)
a free energy specific variable that determines how we decouple.
"""
guest_masses = np.array([a.GetMass() for a in guest_mol.GetAtoms()],
dtype=np.float64)
amber_ff = app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')
host_system = amber_ff.createSystem(host_pdb.topology,
nonbondedMethod=app.NoCutoff,
constraints=None,
rigidWater=False)
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system)
num_host_atoms = len(host_masses)
num_guest_atoms = guest_mol.GetNumAtoms()
# Name, Args, vjp_fn
final_gradients = []
for item in host_fns:
if item[0] == 'LennardJones':
host_lj_params = item[1]
elif item[0] == 'Charges':
host_charge_params = item[1]
elif item[0] == 'GBSA':
host_gb_params = item[1][0]
host_gb_props = item[1][1:]
elif item[0] == 'Exclusions':
host_exclusions = item[1]
else:
final_gradients.append((item[0], item[1]))
guest_exclusion_idxs, guest_scales = nonbonded.generate_exclusion_idxs(
guest_mol, scale12=1.0, scale13=1.0, scale14=0.5)
guest_exclusion_idxs += num_host_atoms
guest_lj_exclusion_scales = guest_scales
guest_charge_exclusion_scales = guest_scales
host_exclusion_idxs = host_exclusions[0]
host_lj_exclusion_scales = host_exclusions[1]
host_charge_exclusion_scales = host_exclusions[2]
combined_exclusion_idxs = np.concatenate(
[host_exclusion_idxs, guest_exclusion_idxs])
combined_lj_exclusion_scales = np.concatenate(
[host_lj_exclusion_scales, guest_lj_exclusion_scales])
combined_charge_exclusion_scales = np.concatenate(
[host_charge_exclusion_scales, guest_charge_exclusion_scales])
# We build up a map of handles to a corresponding vjp_fn that takes in adjoints of output parameters
# for nonbonded terms, the vjp_fn has been modified to take in combined parameters
handler_vjp_fns = {}
for handle in handlers:
results = handle.parameterize(guest_mol)
if isinstance(handle, bonded.HarmonicBondHandler):
bond_idxs, (bond_params, handler_vjp_fn) = results
bond_idxs += num_host_atoms
final_gradients.append(("HarmonicBond", (bond_idxs, bond_params)))
elif isinstance(handle, bonded.HarmonicAngleHandler):
angle_idxs, (angle_params, handler_vjp_fn) = results
angle_idxs += num_host_atoms
final_gradients.append(
("HarmonicAngle", (angle_idxs, angle_params)))
elif isinstance(handle, bonded.ProperTorsionHandler):
torsion_idxs, (torsion_params, handler_vjp_fn) = results
torsion_idxs += num_host_atoms
final_gradients.append(
("PeriodicTorsion", (torsion_idxs, torsion_params)))
elif isinstance(handle, bonded.ImproperTorsionHandler):
torsion_idxs, (torsion_params, handler_vjp_fn) = results
torsion_idxs += num_host_atoms
final_gradients.append(
("PeriodicTorsion", (torsion_idxs, torsion_params)))
elif isinstance(handle, nonbonded.LennardJonesHandler):
guest_lj_params, guest_lj_vjp_fn = results
combined_lj_params, handler_vjp_fn = concat_with_vjps(
host_lj_params, guest_lj_params, None, guest_lj_vjp_fn)
elif isinstance(handle, nonbonded.SimpleChargeHandler):
guest_charge_params, guest_charge_vjp_fn = results
combined_charge_params, handler_vjp_fn = concat_with_vjps(
host_charge_params, guest_charge_params, None,
guest_charge_vjp_fn)
elif isinstance(handle, nonbonded.GBSAHandler):
guest_gb_params, guest_gb_vjp_fn = results
combined_gb_params, handler_vjp_fn = concat_with_vjps(
host_gb_params, guest_gb_params, None, guest_gb_vjp_fn)
elif isinstance(handle, nonbonded.AM1BCCHandler):
guest_charge_params, guest_charge_vjp_fn = results
combined_charge_params, handler_vjp_fn = concat_with_vjps(
host_charge_params, guest_charge_params, None,
guest_charge_vjp_fn)
elif isinstance(handle, nonbonded.AM1CCCHandler):
guest_charge_params, guest_charge_vjp_fn = results
combined_charge_params, handler_vjp_fn = concat_with_vjps(
host_charge_params, guest_charge_params, None,
guest_charge_vjp_fn)
else:
raise Exception("Unknown Handler", handle)
handler_vjp_fns[handle] = handler_vjp_fn
host_conf = []
for x, y, z in host_pdb.positions:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
conformer = guest_mol.GetConformer(0)
mol_a_conf = np.array(conformer.GetPositions(), dtype=np.float64)
mol_a_conf = mol_a_conf / 10 # convert to md_units
x0 = np.concatenate([host_conf, mol_a_conf]) # combined geometry
v0 = np.zeros_like(x0)
pocket_atoms = find_protein_pocket_atoms(x0, num_host_atoms,
restr_search_radius)
N_C = num_host_atoms + num_guest_atoms
N_A = num_host_atoms
cutoff = 100000.0
if stage == 0:
combined_lambda_plane_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs = np.zeros(N_C, dtype=np.int32)
elif stage == 1:
combined_lambda_plane_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs = np.zeros(N_C, dtype=np.int32)
combined_lambda_offset_idxs[num_host_atoms:] = 1
else:
assert 0
final_gradients.append(
('Nonbonded',
(np.asarray(combined_charge_params), np.asarray(combined_lj_params),
combined_exclusion_idxs, combined_charge_exclusion_scales,
combined_lj_exclusion_scales, combined_lambda_plane_idxs,
combined_lambda_offset_idxs, cutoff)))
final_gradients.append(
('GBSA', (np.asarray(combined_charge_params),
np.asarray(combined_gb_params), combined_lambda_plane_idxs,
combined_lambda_offset_idxs, *host_gb_props, cutoff,
cutoff)))
ligand_idxs = np.arange(N_A, N_C, dtype=np.int32)
# restraints
if stage == 0:
lamb_flag = 1
lamb_offset = 0
if stage == 1:
lamb_flag = 0
lamb_offset = 1
# unweighted center of mass restraints
avg_xi = np.mean(x0[ligand_idxs], axis=0)
avg_xj = np.mean(x0[pocket_atoms], axis=0)
ctr_dij = np.sqrt(np.sum((avg_xi - avg_xj)**2))
combined_masses = np.concatenate([host_masses, guest_masses])
# restraints
final_gradients.append(
('CentroidRestraint',
(ligand_idxs, pocket_atoms, combined_masses, restr_force_constant,
ctr_dij, lamb_flag, lamb_offset)))
ssc = standard_state.harmonic_com_ssc(restr_force_constant, ctr_dij,
intg_temperature)
return x0, combined_masses, ssc, final_gradients, handler_vjp_fns