def do_deletion(
x0,
v0,
combined_bps,
combined_masses,
box,
guest_name,
leg_type,
u_impls,
deletion_steps,
):
seed = 2021
intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, combined_masses, seed).impl()
ctxt = custom_ops.Context(x0, v0, box, intg, u_impls)
# du_dl_obs = custom_ops.FullPartialUPartialLambda(u_impls, subsample_freq)
# ctxt.add_observable(du_dl_obs)
deletion_lambda_schedule = np.linspace(MIN_LAMBDA, DELETION_MAX_LAMBDA,
deletion_steps)
subsample_freq = 1
full_du_dls, _, _ = ctxt.multiple_steps(deletion_lambda_schedule,
subsample_freq)
step = len(deletion_lambda_schedule) - 1
lamb = deletion_lambda_schedule[-1]
ctxt.step(lamb)
report.report_step(
ctxt,
step,
lamb,
box,
combined_bps,
u_impls,
guest_name,
deletion_steps,
f"{leg_type.upper()}_DELETION",
)
if report.too_much_force(ctxt, lamb, box, combined_bps, u_impls):
print("Not calculating work (too much force)")
return None
# Note: this condition only applies for ABFE, not RBFE
if abs(full_du_dls[0]) > 0.001 or abs(full_du_dls[-1]) > 0.001:
print("Not calculating work (du_dl endpoints are not ~0)")
return None
work = np.trapz(full_du_dls, deletion_lambda_schedule[::subsample_freq])
print(f"guest_name: {guest_name}\t{leg_type}_work: {work:.2f}")
return work
def do_switch(
x0,
v0,
combined_bps,
combined_masses,
box,
guest_name,
leg_type,
u_impls,
transition_steps,
):
seed = 2021
intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, combined_masses, seed).impl()
ctxt = custom_ops.Context(x0, v0, box, intg, u_impls)
switching_lambda_schedule = np.linspace(MIN_LAMBDA, MAX_LAMBDA,
transition_steps)
subsample_interval = 1
full_du_dls, _, _ = ctxt.multiple_steps(switching_lambda_schedule,
subsample_interval)
step = len(switching_lambda_schedule) - 1
lamb = switching_lambda_schedule[-1]
ctxt.step(lamb)
report.report_step(
ctxt,
step,
lamb,
box,
combined_bps,
u_impls,
guest_name,
transition_steps,
f"{leg_type.upper()}_SWITCH",
)
if report.too_much_force(ctxt, lamb, box, combined_bps, u_impls):
return
work = np.trapz(full_du_dls,
switching_lambda_schedule[::subsample_interval])
print(f"guest_name: {guest_name}\t{leg_type}_work: {work:.2f}")
return work
def run_leg(
combined_coords,
combined_bps,
combined_masses,
host_box,
guest_name,
leg_type,
num_switches,
transition_steps,
):
x0 = combined_coords
v0 = np.zeros_like(x0)
print(
f"{leg_type.upper()}_SYSTEM",
f"guest_name: {guest_name}",
f"num_atoms: {len(x0)}",
)
seed = 2021
intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, combined_masses, seed).impl()
u_impls = []
for bp in combined_bps:
bp_impl = bp.bound_impl(precision=np.float32)
u_impls.append(bp_impl)
ctxt = custom_ops.Context(x0, v0, host_box, intg, u_impls)
# TODO: pre-equilibrate?
# equilibrate & shoot off switching jobs
steps_per_batch = 1001
works = []
for b in range(num_switches):
equil2_lambda_schedule = np.ones(steps_per_batch) * MIN_LAMBDA
ctxt.multiple_steps(equil2_lambda_schedule, 0)
lamb = equil2_lambda_schedule[-1]
step = len(equil2_lambda_schedule) - 1
report.report_step(
ctxt,
(b + 1) * step,
lamb,
host_box,
combined_bps,
u_impls,
guest_name,
num_switches * steps_per_batch,
f"{leg_type.upper()}_EQUILIBRATION_2",
)
if report.too_much_force(ctxt, MIN_LAMBDA, host_box, combined_bps,
u_impls):
return
work = do_switch(
ctxt.get_x_t(),
ctxt.get_v_t(),
combined_bps,
combined_masses,
host_box,
guest_name,
leg_type,
u_impls,
transition_steps,
)
works.append(work)
return works
def dock_and_equilibrate(
host_pdbfile,
guests_sdfile,
max_lambda,
insertion_steps,
eq_steps,
outdir,
fewer_outfiles=False,
constant_atoms=[],
):
"""Solvates a host, inserts guest(s) into solvated host, equilibrates
Parameters
----------
host_pdbfile: path to host pdb file to dock into
guests_sdfile: path to input sdf with guests to pose/dock
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)
insertion_steps: how many steps to insert the guest over (recommended: 501)
eq_steps: how many steps of equilibration to do after insertion (recommended: 15001)
outdir: where to write output (will be created if it does not already exist)
fewer_outfiles: if True, will only write frames for the equilibration, not insertion
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 for the last step of insertion
(outdir/<guest_name>/<guest_name>_ins_<step>_[host.pdb/guest.sdf])
A pdb & sdf file every 1000 steps of equilibration
(outdir/<guest_name>/<guest_name>_eq_<step>_[host.pdb/guest.sdf])
stdout corresponding to the files written noting the lambda value and energy
stdout for each guest noting the work of transition, if applicable
stdout for each guest noting how long it took to run
Note
----
The work will not be calculated if the du_dl endpoints are not close to 0 or if any norm of
force per atom exceeds 20000 kJ/(mol*nm) [MAX_NORM_FORCE defined in docking/report.py]
"""
if not os.path.exists(outdir):
os.makedirs(outdir)
print(f"""
HOST_PDBFILE = {host_pdbfile}
GUESTS_SDFILE = {guests_sdfile}
OUTDIR = {outdir}
MAX_LAMBDA = {max_lambda}
INSERTION_STEPS = {insertion_steps}
EQ_STEPS = {eq_steps}
""")
# Prepare host
# TODO: handle extra (non-transitioning) guests?
print("Solvating host...")
(
solvated_host_system,
solvated_host_coords,
_,
_,
host_box,
solvated_topology,
) = builders.build_protein_system(host_pdbfile)
_, solvated_host_pdb = tempfile.mkstemp(suffix=".pdb", text=True)
writer = pdb_writer.PDBWriter([solvated_topology], solvated_host_pdb)
writer.write_frame(solvated_host_coords)
writer.close()
solvated_host_mol = Chem.MolFromPDBFile(solvated_host_pdb, removeHs=False)
os.remove(solvated_host_pdb)
ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
# Run the procedure
print("Getting guests...")
suppl = Chem.SDMolSupplier(guests_sdfile, removeHs=False)
for guest_mol in suppl:
start_time = time.time()
guest_name = guest_mol.GetProp("_Name")
guest_conformer = guest_mol.GetConformer(0)
orig_guest_coords = np.array(guest_conformer.GetPositions(),
dtype=np.float64)
orig_guest_coords = orig_guest_coords / 10 # convert to md_units
minimized_coords = minimizer.minimize_host_4d([guest_mol],
solvated_host_system,
solvated_host_coords, ff,
host_box)
afe = free_energy.AbsoluteFreeEnergy(guest_mol, ff)
ups, sys_params, combined_masses, _ = afe.prepare_host_edge(
ff.get_ordered_params(), solvated_host_system, minimized_coords)
combined_bps = []
for up, sp in zip(ups, sys_params):
combined_bps.append(up.bind(sp))
x0 = np.concatenate([minimized_coords, orig_guest_coords])
v0 = np.zeros_like(x0)
print("SYSTEM", f"guest_name: {guest_name}", f"num_atoms: {len(x0)}")
for atom_num in constant_atoms:
combined_masses[atom_num - 1] += 50000
seed = 2021
intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, combined_masses,
seed).impl()
u_impls = []
for bp in combined_bps:
bp_impl = bp.bound_impl(precision=np.float32)
u_impls.append(bp_impl)
ctxt = custom_ops.Context(x0, v0, host_box, intg, u_impls)
# insert guest
insertion_lambda_schedule = np.linspace(max_lambda, 0.0,
insertion_steps)
calc_work = True
# collect a du_dl calculation once every other step
subsample_interval = 1
full_du_dls, _, _ = ctxt.multiple_steps(insertion_lambda_schedule,
subsample_interval)
step = len(insertion_lambda_schedule) - 1
lamb = insertion_lambda_schedule[-1]
ctxt.step(lamb)
report.report_step(
ctxt,
step,
lamb,
host_box,
combined_bps,
u_impls,
guest_name,
insertion_steps,
"INSERTION",
)
if not fewer_outfiles:
host_coords = ctxt.get_x_t()[:len(solvated_host_coords)] * 10
guest_coords = ctxt.get_x_t()[len(solvated_host_coords):] * 10
report.write_frame(
host_coords,
solvated_host_mol,
guest_coords,
guest_mol,
guest_name,
outdir,
str(step).zfill(len(str(insertion_steps))),
"ins",
)
if report.too_much_force(ctxt, lamb, host_box, combined_bps, u_impls):
print("Not calculating work (too much force)")
calc_work = False
continue
# Note: this condition only applies for ABFE, not RBFE
if abs(full_du_dls[0]) > 0.001 or abs(full_du_dls[-1]) > 0.001:
print("Not calculating work (du_dl endpoints are not ~0)")
calc_work = False
if calc_work:
work = np.trapz(full_du_dls,
insertion_lambda_schedule[::subsample_interval])
print(f"guest_name: {guest_name}\tinsertion_work: {work:.2f}")
# equilibrate
for step in range(eq_steps):
ctxt.step(0.00)
if step % 1000 == 0:
report.report_step(
ctxt,
step,
0.00,
host_box,
combined_bps,
u_impls,
guest_name,
eq_steps,
"EQUILIBRATION",
)
if (not fewer_outfiles) or (step == eq_steps - 1):
host_coords = ctxt.get_x_t()[:len(solvated_host_coords
)] * 10
guest_coords = ctxt.get_x_t()[len(solvated_host_coords
):] * 10
report.write_frame(
host_coords,
solvated_host_mol,
guest_coords,
guest_mol,
guest_name,
outdir,
str(step).zfill(len(str(eq_steps))),
"eq",
)
if step in (0, int(eq_steps / 2), eq_steps - 1):
if report.too_much_force(ctxt, 0.00, host_box, combined_bps,
u_impls):
break
end_time = time.time()
print(f"{guest_name} took {(end_time - start_time):.2f} seconds")
def run_leg(
orig_host_coords,
orig_guest_coords,
combined_bps,
combined_masses,
host_box,
guest_name,
leg_type,
host_mol,
guest_mol,
outdir,
num_deletions,
deletion_steps,
insertion_max_lambda,
insertion_steps,
eq1_steps,
fewer_outfiles=False,
no_outfiles=False,
):
x0 = np.concatenate([orig_host_coords, orig_guest_coords])
v0 = np.zeros_like(x0)
print(
f"{leg_type.upper()}_SYSTEM",
f"guest_name: {guest_name}",
f"num_atoms: {len(x0)}",
)
seed = 2021
intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, combined_masses, seed).impl()
u_impls = []
for bp in combined_bps:
bp_impl = bp.bound_impl(precision=np.float32)
u_impls.append(bp_impl)
ctxt = custom_ops.Context(x0, v0, host_box, intg, u_impls)
# insert guest
insertion_lambda_schedule = np.linspace(insertion_max_lambda, MIN_LAMBDA,
insertion_steps)
ctxt.multiple_steps(insertion_lambda_schedule, 0) # do not collect du_dls
lamb = insertion_lambda_schedule[-1]
step = len(insertion_lambda_schedule) - 1
report.report_step(
ctxt,
step,
lamb,
host_box,
combined_bps,
u_impls,
guest_name,
insertion_steps,
f"{leg_type.upper()}_INSERTION",
)
if not fewer_outfiles and not no_outfiles:
host_coords = ctxt.get_x_t()[:len(orig_host_coords)] * 10
guest_coords = ctxt.get_x_t()[len(orig_host_coords):] * 10
report.write_frame(
host_coords,
host_mol,
guest_coords,
guest_mol,
guest_name,
outdir,
str(step).zfill(len(str(insertion_steps))),
f"{leg_type}-ins",
)
if report.too_much_force(ctxt, lamb, host_box, combined_bps, u_impls):
return []
# equilibrate
equil_lambda_schedule = np.ones(eq1_steps) * MIN_LAMBDA
lamb = equil_lambda_schedule[-1]
step = len(equil_lambda_schedule) - 1
ctxt.multiple_steps(equil_lambda_schedule, 0)
report.report_step(
ctxt,
step,
MIN_LAMBDA,
host_box,
combined_bps,
u_impls,
guest_name,
eq1_steps,
f"{leg_type.upper()}_EQUILIBRATION_1",
)
if not fewer_outfiles and not no_outfiles:
host_coords = ctxt.get_x_t()[:len(orig_host_coords)] * 10
guest_coords = ctxt.get_x_t()[len(orig_host_coords):] * 10
report.write_frame(
host_coords,
host_mol,
guest_coords,
guest_mol,
guest_name,
outdir,
str(step).zfill(len(str(eq1_steps))),
f"{leg_type}-eq1",
)
if report.too_much_force(ctxt, MIN_LAMBDA, host_box, combined_bps,
u_impls):
print("Too much force")
return []
# equilibrate more & shoot off deletion jobs
steps_per_batch = 1001
works = []
for b in range(num_deletions):
deletion_lambda_schedule = np.ones(steps_per_batch) * MIN_LAMBDA
ctxt.multiple_steps(deletion_lambda_schedule, 0)
lamb = deletion_lambda_schedule[-1]
step = len(deletion_lambda_schedule) - 1
report.report_step(
ctxt,
(b + 1) * step,
MIN_LAMBDA,
host_box,
combined_bps,
u_impls,
guest_name,
num_deletions * steps_per_batch,
f"{leg_type.upper()}_EQUILIBRATION_2",
)
# TODO: if guest has undocked, stop simulation
if not no_outfiles:
host_coords = ctxt.get_x_t()[:len(orig_host_coords)] * 10
guest_coords = ctxt.get_x_t()[len(orig_host_coords):] * 10
report.write_frame(
host_coords,
host_mol,
guest_coords,
guest_mol,
guest_name,
outdir,
str((b + 1) * step).zfill(
len(str(num_deletions * steps_per_batch))),
f"{leg_type}-eq2",
)
if report.too_much_force(ctxt, MIN_LAMBDA, host_box, combined_bps,
u_impls):
print("Too much force")
return works
work = do_deletion(
ctxt.get_x_t(),
ctxt.get_v_t(),
combined_bps,
combined_masses,
host_box,
guest_name,
leg_type,
u_impls,
deletion_steps,
)
works.append(work)
return works