def test_avg_potential_param_sizes_is_zero(self): np.random.seed(814) N = 8 D = 3 x0 = np.random.rand(N, D).astype(dtype=np.float64) * 2 masses = np.random.rand(N) v0 = np.random.rand(x0.shape[0], x0.shape[1]) num_steps = 3 ca = np.random.rand() cbs = -np.random.rand(len(masses)) / 1 ccs = np.zeros_like(cbs) dt = 2e-3 lamb = np.random.rand() box = np.eye(3) * 1.5 intg = custom_ops.LangevinIntegrator(dt, ca, cbs, ccs, 814) # Construct a 'bad' centroid restraint potential = potentials.CentroidRestraint( np.random.randint(N, size=5, dtype=np.int32), np.random.randint(N, size=5, dtype=np.int32), 10.0, 0.0) # Bind to empty params bp = potential.bind(np.zeros(0)).bound_impl(precision=np.float64) ctxt = custom_ops.Context(x0, v0, box, intg, [bp]) for _ in range(num_steps): ctxt.step(lamb)
def equilibrate(integrator, barostat, potentials, coords, box, lamb, equil_steps) -> Tuple: all_impls = [] v0 = np.zeros_like(coords) for bp in potentials: impl = bp.bound_impl(np.float32) all_impls.append(impl) if integrator.seed == 0: integrator = copy.deepcopy(integrator) integrator.seed = np.random.randint(np.iinfo(np.int32).max) if barostat.seed == 0: barostat = copy.deepcopy(barostat) barostat.seed = np.random.randint(np.iinfo(np.int32).max) intg_impl = integrator.impl() baro_impl = barostat.impl(all_impls) # context components: positions, velocities, box, integrator, energy fxns ctxt = custom_ops.Context( coords, v0, box, intg_impl, all_impls, barostat=baro_impl, ) # equilibration equil_schedule = np.ones(equil_steps) * lamb ctxt.multiple_steps(equil_schedule) return CoordsVelBox(coords=ctxt.get_x_t(), velocities=ctxt.get_v_t(), box=ctxt.get_box())
def test_set_and_get(self): """ This test the setters and getters in the context. """ np.random.seed(4321) N = 8 D = 3 x0 = np.random.rand(N, D).astype(dtype=np.float64) * 2 E = 2 lambda_plane_idxs = np.random.randint(low=0, high=2, size=N, dtype=np.int32) lambda_offset_idxs = np.random.randint(low=0, high=2, size=N, dtype=np.int32) params, _, test_nrg = prepare_nb_system( x0, E, lambda_plane_idxs, lambda_offset_idxs, p_scale=3.0, cutoff=1.0, ) masses = np.random.rand(N) v0 = np.random.rand(x0.shape[0], x0.shape[1]) temperature = 300 dt = 2e-3 friction = 0.0 ca, cbs, ccs = langevin_coefficients(temperature, dt, friction, masses) box = np.eye(3) * 3.0 intg = custom_ops.LangevinIntegrator(dt, ca, cbs, ccs, 1234) bp = test_nrg.bind(params).bound_impl(precision=np.float64) bps = [bp] ctxt = custom_ops.Context(x0, v0, box, intg, bps) np.testing.assert_equal(ctxt.get_x_t(), x0) np.testing.assert_equal(ctxt.get_v_t(), v0) np.testing.assert_equal(ctxt.get_box(), box) new_x = np.random.rand(N, 3) ctxt.set_x_t(new_x) np.testing.assert_equal(ctxt.get_x_t(), new_x)
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_deletion(x0, v0, combined_bps, combined_masses, box, guest_name, leg_type): 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, box, intg, u_impls) subsample_freq = 2 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, TRANSITION_STEPS ) calc_work = True for step, lamb in enumerate(deletion_lambda_schedule): ctxt.step(lamb) if step % 100 == 0: report.report_step( ctxt, step, lamb, box, combined_bps, u_impls, guest_name, TRANSITION_STEPS, f"{leg_type.upper()}_DELETION", ) if step in (0, int(TRANSITION_STEPS/2), TRANSITION_STEPS-1): if report.too_much_force(ctxt, lamb, box, combined_bps, u_impls): calc_work = False return # 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(), deletion_lambda_schedule[::subsample_freq] ) print(f"guest_name: {guest_name}\t{leg_type}_work: {work:.2f}")
def equilibrate_solvent_phase( ubps, params, masses, coords, # minimized_coords box, temperature, pressure, num_steps, seed=None, ): """ Generate samples in the solvent phase. """ dt = 1e-4 friction = 1.0 bps = [] for p, bp in zip(params, ubps): bps.append(bp.bind(p)) all_impls = [bp.bound_impl(np.float32) for bp in bps] intg_equil = lib.LangevinIntegrator(temperature, dt, friction, masses, seed) intg_equil_impl = intg_equil.impl() bond_list = get_bond_list(ubps[0]) group_idxs = get_group_indices(bond_list) barostat_interval = 5 barostat = lib.MonteCarloBarostat(len(masses), pressure, temperature, group_idxs, barostat_interval, seed + 1) barostat_impl = barostat.impl(all_impls) # equilibration/minimization doesn't need a barostat equil_ctxt = custom_ops.Context(coords, np.zeros_like(coords), box, intg_equil_impl, all_impls, barostat_impl) lamb = 0.0 # TODO: revert to 50k equil_schedule = np.ones(num_steps) * lamb equil_ctxt.multiple_steps(equil_schedule) x0 = equil_ctxt.get_x_t() # (ytz): This has to be zeros_like for now since if we freeze ligand # coordinates it would start to move during rejected moves. v0 = np.zeros_like(x0) return CoordsVelBox(x0, v0, equil_ctxt.get_box())
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 move(self, x: CoordsVelBox): # note: context creation overhead here is actually very small! ctxt = custom_ops.Context(x.coords, x.velocities, x.box, self.integrator_impl, self.bound_impls, self.barostat_impl) # arguments: lambda_schedule, du_dl_interval, x_interval _ = ctxt.multiple_steps(self.lamb * np.ones(self.n_steps), 0, 0) x_t = ctxt.get_x_t() v_t = ctxt.get_v_t() box = ctxt.get_box() after_npt = CoordsVelBox(x_t, v_t, box) self.n_proposed += 1 self.n_accepted += 1 return after_npt
def run(args): lamb, intg, bound_potentials, masses, x0, box, gpu_idx, stage = args # print("running on", gpu_idx) os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_idx) u_impls = [] for bp in bound_potentials: u_impls.append(bp.bound_impl(precision=np.float32)) # important that we reseed here. intg.seed = np.random.randint(np.iinfo(np.int32).max) intg_impl = intg.impl() v0 = np.zeros_like(x0) ctxt = custom_ops.Context(x0, v0, box, intg_impl, u_impls) # secondary minimization needed only for stage 1 if stage == 1: for l in np.linspace(0.35, lamb, 500): ctxt.step(l) # equilibration for step in range(20000): # for step in range(1000): ctxt.step(lamb) # print(ctxt.get_x_t()) du_dl_obs = custom_ops.AvgPartialUPartialLambda(u_impls, 10) ctxt.add_observable(du_dl_obs) # add observable for <du/dl> for step in range(50000): # for step in range(5000): ctxt.step(lamb) print(lamb, du_dl_obs.avg_du_dl()) assert np.any(np.abs(ctxt.get_x_t()) > 100) == False assert np.any(np.isnan(ctxt.get_x_t())) == False assert np.any(np.isinf(ctxt.get_x_t())) == False return du_dl_obs.avg_du_dl()
def minimize(args): bound_potentials, masses, x0, box = args u_impls = [] for bp in bound_potentials: u_impls.append(bp.bound_impl(precision=np.float32)) seed = np.random.randint(np.iinfo(np.int32).max) intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, masses, seed).impl() v0 = np.zeros_like(x0) ctxt = custom_ops.Context(x0, v0, box, intg, u_impls) lambda_schedule = np.linspace(0.35, 0.0, 500) for lamb in lambda_schedule: ctxt.step(lamb) return ctxt.get_x_t()
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 benchmark( label, masses, lamb, x0, v0, box, bound_potentials, hmr=False, verbose=True, num_batches=100, steps_per_batch=1000, compute_du_dl_interval=0, barostat_interval=0, ): """ TODO: configuration blob containing num_batches, steps_per_batch, and any other options """ seed = 1234 dt = 1.5e-3 temperature = 300 pressure = 1.0 seconds_per_day = 86400 harmonic_bond_potential = bound_potentials[0] bond_list = get_bond_list(harmonic_bond_potential) if hmr: dt = 2.5e-3 masses = apply_hmr(masses, bond_list) intg = LangevinIntegrator(temperature, dt, 1.0, np.array(masses), seed).impl() bps = [] for potential in bound_potentials: bps.append(potential.bound_impl(precision=np.float32)) # get the bound implementation baro_impl = None if barostat_interval > 0: group_idxs = get_group_indices(bond_list) baro = MonteCarloBarostat( x0.shape[0], pressure, temperature, group_idxs, barostat_interval, seed, ) baro_impl = baro.impl(bps) ctxt = custom_ops.Context( x0, v0, box, intg, bps, barostat=baro_impl, ) batch_times = [] lambda_schedule = np.ones(steps_per_batch) * lamb # run once before timer starts ctxt.multiple_steps(lambda_schedule, compute_du_dl_interval) start = time.time() for batch in range(num_batches): # time the current batch batch_start = time.time() du_dls, _, _ = ctxt.multiple_steps(lambda_schedule, compute_du_dl_interval) batch_end = time.time() delta = batch_end - batch_start batch_times.append(delta) steps_per_second = steps_per_batch / np.mean(batch_times) steps_per_day = steps_per_second * seconds_per_day ps_per_day = dt * steps_per_day ns_per_day = ps_per_day * 1e-3 if verbose: print(f"steps per second: {steps_per_second:.3f}") print(f"ns per day: {ns_per_day:.3f}") assert np.all(np.abs(ctxt.get_x_t()) < 1000) print( f"{label}: N={x0.shape[0]} speed: {ns_per_day:.2f}ns/day dt: {dt*1e3}fs (ran {steps_per_batch * num_batches} steps in {(time.time() - start):.2f}s)" )
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")
# wrapper function -- since contexts are not pickle-able -- which will # be useful later in timemachine's multi-device parallelization strategy) # note: OpenMM unit system used throughout # (temperature: kelvin, timestep: picosecond, collision_rate: picosecond^-1) intg = LangevinIntegrator(300.0, 1.5e-3, 1.0, combined_masses, seed).impl() x0 = combined_coords v0 = np.zeros_like(x0) u_impls = [] for bp in final_potentials: u_impls.append(bp.bound_impl(np.float32)) # context components: positions, velocities, box, integrator, energy fxns ctxt = custom_ops.Context(x0, v0, box, intg, u_impls) for step, lamb in enumerate(np.linspace(1.0, final_lamb, 1000)): if step % 500 == 0: writer.write_frame(ctxt.get_x_t() * 10) ctxt.step(lamb) # print("insertion energy", ctxt._get_u_t_minus_1()) # note: these 5000 steps are "equilibration", before we attach a reporter / # "observable" to the context and start running "production" for step in range(5000): if step % 500 == 0: writer.write_frame(ctxt.get_x_t() * 10) ctxt.step(final_lamb)
def simulate( lamb, box, x0, v0, final_potentials, integrator, barostat, equil_steps, prod_steps, x_interval, u_interval, lambda_windows, ): """ Run a simulation and collect relevant statistics for this simulation. Parameters ---------- lamb: float lambda value used for the equilibrium simulation box: np.array 3x3 numpy array of the box, dtype should be np.float64 x0: np.array Nx3 numpy array of the coordinates v0: np.array Nx3 numpy array of the velocities final_potentials: list list of unbound potentials integrator: timemachine.Integrator integrator to be used for dynamics barostat: timemachine.Barostat barostat to be used for equilibration equil_steps: int number of equilibration steps prod_steps: int number of production steps x_interval: int how often we store coordinates. If x_interval == 0 then no frames are returned. u_interval: int how often we store energies. If u_interval == 0 then no energies are returned lambda_windows: list of float lambda windows we evaluate energies at. Returns ------- SimulationResult Results of the simulation. """ all_impls = [] for bp in final_potentials: impl = bp.bound_impl(np.float32) all_impls.append(impl) # fire minimize once again, needed for parameter interpolation x0 = minimizer.fire_minimize(x0, all_impls, box, np.ones(100, dtype=np.float64) * lamb) # sanity check that forces are well behaved for bp in all_impls: du_dx, du_dl, u = bp.execute(x0, box, lamb) norm_forces = np.linalg.norm(du_dx, axis=1) assert np.all(norm_forces < 25000 ), "Forces much greater than expected after minimization" if integrator.seed == 0: # this deepcopy is needed if we're running if client == None integrator = copy.deepcopy(integrator) integrator.seed = np.random.randint(np.iinfo(np.int32).max) if barostat.seed == 0: barostat = copy.deepcopy(barostat) barostat.seed = np.random.randint(np.iinfo(np.int32).max) intg_impl = integrator.impl() # technically we need to only pass in the nonbonded impl barostat_impl = barostat.impl(all_impls) # context components: positions, velocities, box, integrator, energy fxns ctxt = custom_ops.Context(x0, v0, box, intg_impl, all_impls, barostat_impl) # equilibration equil_schedule = np.ones(equil_steps) * lamb ctxt.multiple_steps(equil_schedule) # (ytz): intentionally hard-coded, I'd rather the end-user *not* # muck with this unless they have a good reason to. barostat_impl.set_interval(25) full_us, xs, boxes = ctxt.multiple_steps_U(lamb, prod_steps, np.array(lambda_windows), u_interval, x_interval) result = SimulationResult( xs=xs.astype("float32"), boxes=boxes.astype("float32"), lambda_us=full_us, ) return result
def test_barostat_partial_group_idxs(): """Verify that the barostat can handle a subset of the molecules rather than all of them. This test only verify that it runs, not the behavior""" temperature = 300.0 * unit.kelvin initial_waterbox_width = 3.0 * unit.nanometer timestep = 1.5 * unit.femtosecond barostat_interval = 3 collision_rate = 1.0 / unit.picosecond seed = 2021 np.random.seed(seed) pressure = 1.0 * unit.atmosphere mol_a = hif2a_ligand_pair.mol_a ff = hif2a_ligand_pair.ff complex_system, complex_coords, complex_box, complex_top = build_water_system( initial_waterbox_width.value_in_unit(unit.nanometer)) min_complex_coords = minimize_host_4d([mol_a], complex_system, complex_coords, ff, complex_box) afe = AbsoluteFreeEnergy(mol_a, ff) unbound_potentials, sys_params, masses, coords = afe.prepare_host_edge( ff.get_ordered_params(), complex_system, min_complex_coords) # get list of molecules for barostat by looking at bond table harmonic_bond_potential = unbound_potentials[0] bond_list = get_bond_list(harmonic_bond_potential) group_indices = get_group_indices(bond_list) # Cut the number of groups in half group_indices = group_indices[len(group_indices) // 2:] lam = 1.0 bound_potentials = [] for params, unbound_pot in zip(sys_params, unbound_potentials): bp = unbound_pot.bind(np.asarray(params)) bound_potentials.append(bp) u_impls = [] for bp in bound_potentials: bp_impl = bp.bound_impl(precision=np.float32) u_impls.append(bp_impl) integrator = LangevinIntegrator( temperature.value_in_unit(unit.kelvin), timestep.value_in_unit(unit.picosecond), collision_rate.value_in_unit(unit.picosecond**-1), masses, seed, ) integrator_impl = integrator.impl() v_0 = sample_velocities(masses * unit.amu, temperature) baro = custom_ops.MonteCarloBarostat( coords.shape[0], pressure.value_in_unit(unit.bar), temperature.value_in_unit(unit.kelvin), group_indices, barostat_interval, u_impls, seed, ) ctxt = custom_ops.Context(coords, v_0, complex_box, integrator_impl, u_impls, barostat=baro) ctxt.multiple_steps(np.ones(1000) * lam)
def _simulate(lamb, box, x0, v0, final_potentials, integrator, equil_steps, prod_steps): all_impls = [] bonded_impls = [] nonbonded_impls = [] # set up observables for du_dps here as well. du_dp_obs = [] for bps in final_potentials: obs_list = [] for bp in bps: impl = bp.bound_impl(np.float32) if isinstance(bp, potentials.InterpolatedPotential) or isinstance(bp, potentials.LambdaPotential): bp = bp.get_u_fn() if isinstance(bp, potentials.Nonbonded): nonbonded_impls.append(impl) else: bonded_impls.append(impl) all_impls.append(impl) obs_list.append(custom_ops.AvgPartialUPartialParam(impl, 5)) du_dp_obs.append(obs_list) intg_impl = integrator.impl() # context components: positions, velocities, box, integrator, energy fxns ctxt = custom_ops.Context( x0, v0, box, intg_impl, all_impls ) # equilibration for step in range(equil_steps): ctxt.step(lamb) bonded_du_dl_obs = custom_ops.FullPartialUPartialLambda(bonded_impls, 5) nonbonded_du_dl_obs = custom_ops.FullPartialUPartialLambda(nonbonded_impls, 5) # add observable ctxt.add_observable(bonded_du_dl_obs) ctxt.add_observable(nonbonded_du_dl_obs) for obs_list in du_dp_obs: for obs in obs_list: ctxt.add_observable(obs) for _ in range(prod_steps): ctxt.step(lamb) bonded_full_du_dls = bonded_du_dl_obs.full_du_dl() nonbonded_full_du_dls = nonbonded_du_dl_obs.full_du_dl() bonded_mean, bonded_std = np.mean(bonded_full_du_dls), np.std(bonded_full_du_dls) nonbonded_mean, nonbonded_std = np.mean(nonbonded_full_du_dls), np.std(nonbonded_full_du_dls) # keep the structure of grads the same as that of final_potentials so we can properly # form their vjps. grads = [] for obs_list in du_dp_obs: grad_list = [] for obs in obs_list: grad_list.append(obs.avg_du_dp()) grads.append(grad_list) return (bonded_mean, bonded_std), (nonbonded_mean, nonbonded_std), grads
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 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
def simulate( lamb, box, x0, v0, final_potentials, integrator, equil_steps, prod_steps, barostat, x_interval=1000, du_dl_interval=5, ) -> SimulationResult: """ Run a simulation and collect relevant statistics for this simulation. Parameters ---------- lamb: float lambda parameter box: np.array 3x3 numpy array of the box, dtype should be np.float64 x0: np.array Nx3 numpy array of the coordinates v0: np.array Nx3 numpy array of the velocities final_potentials: list list of unbound potentials integrator: timemachine.Integrator integrator to be used for dynamics equil_steps: int number of equilibration steps prod_steps: int number of production steps x_interval: int how often we store coordinates. if x_interval == 0 then no frames are returned. du_dl_interval: int how often we store du_dls. if du_dl_interval == 0 then no du_dls are returned barostat: timemachine.lib.MonteCarloBarostat Monte carlo barostat to use when simulating. Returns ------- SimulationResult Results of the simulation. """ all_impls = [] bonded_impls = [] nonbonded_impls = [] for bp in final_potentials: impl = bp.bound_impl(np.float32) if isinstance(bp, potentials.Nonbonded): nonbonded_impls.append(impl) else: bonded_impls.append(impl) all_impls.append(impl) if integrator.seed == 0: integrator = copy.deepcopy(integrator) integrator.seed = np.random.randint(np.iinfo(np.int32).max) if barostat.seed == 0: barostat = copy.deepcopy(barostat) barostat.seed = np.random.randint(np.iinfo(np.int32).max) intg_impl = integrator.impl() baro_impl = barostat.impl(all_impls) # context components: positions, velocities, box, integrator, energy fxns ctxt = custom_ops.Context( x0, v0, box, intg_impl, all_impls, barostat=baro_impl, ) base_interval = baro_impl.get_interval() # Use an interval of 5 for equilibration baro_impl.set_interval(5) # equilibration equil_schedule = np.ones(equil_steps) * lamb ctxt.multiple_steps(equil_schedule) baro_impl.set_interval(base_interval) prod_schedule = np.ones(prod_steps) * lamb full_du_dls, xs, _ = ctxt.multiple_steps(prod_schedule, du_dl_interval, x_interval) result = SimulationResult(xs=xs.astype("float32"), du_dls=full_du_dls) return result
def test_fwd_mode(self): """ This test ensures that we can reverse-mode differentiate observables that are dU_dlambdas of each state. We provide adjoints with respect to each computed dU/dLambda. """ np.random.seed(4321) N = 8 B = 5 A = 0 T = 0 D = 3 x0 = np.random.rand(N, D).astype(dtype=np.float64) * 2 E = 2 lambda_plane_idxs = np.random.randint(low=0, high=2, size=N, dtype=np.int32) lambda_offset_idxs = np.random.randint(low=0, high=2, size=N, dtype=np.int32) params, ref_nrg_fn, test_nrg = prepare_nb_system( x0, E, lambda_plane_idxs, lambda_offset_idxs, p_scale=3.0, # cutoff=0.5, cutoff=1.5) masses = np.random.rand(N) v0 = np.random.rand(x0.shape[0], x0.shape[1]) N = len(masses) num_steps = 5 lambda_schedule = np.random.rand(num_steps) ca = np.random.rand() cbs = -np.random.rand(len(masses)) / 1 ccs = np.zeros_like(cbs) dt = 2e-3 lamb = np.random.rand() def loss_fn(du_dls): return jnp.sum(du_dls * du_dls) / du_dls.shape[0] def sum_loss_fn(du_dls): du_dls = np.sum(du_dls, axis=0) return jnp.sum(du_dls * du_dls) / du_dls.shape[0] def integrate_once_through(x_t, v_t, box, params): dU_dx_fn = jax.grad(ref_nrg_fn, argnums=(0, )) dU_dp_fn = jax.grad(ref_nrg_fn, argnums=(1, )) dU_dl_fn = jax.grad(ref_nrg_fn, argnums=(3, )) all_du_dls = [] all_du_dps = [] all_xs = [] all_du_dxs = [] all_us = [] for step in range(num_steps): u = ref_nrg_fn(x_t, params, box, lamb) all_us.append(u) du_dl = dU_dl_fn(x_t, params, box, lamb)[0] all_du_dls.append(du_dl) du_dp = dU_dp_fn(x_t, params, box, lamb)[0] all_du_dps.append(du_dp) du_dx = dU_dx_fn(x_t, params, box, lamb)[0] all_du_dxs.append(du_dx) v_t = ca * v_t + np.expand_dims(cbs, axis=-1) * du_dx x_t = x_t + v_t * dt all_xs.append(x_t) # note that we do not calculate the du_dl of the last frame. return all_xs, all_du_dxs, all_du_dps, all_du_dls, all_us box = np.eye(3) * 1.5 # when we have multiple parameters, we need to set this up correctly ref_all_xs, ref_all_du_dxs, ref_all_du_dps, ref_all_du_dls, ref_all_us = integrate_once_through( x0, v0, box, params) intg = custom_ops.LangevinIntegrator(dt, ca, cbs, ccs, 1234) bp = test_nrg.bind(params).bound_impl(precision=np.float64) bps = [bp] ctxt = custom_ops.Context(x0, v0, box, intg, bps) test_obs = custom_ops.AvgPartialUPartialParam(bp, 1) test_obs_f2 = custom_ops.AvgPartialUPartialParam(bp, 2) test_obs_du_dl = custom_ops.AvgPartialUPartialLambda(bps, 1) test_obs_f2_du_dl = custom_ops.AvgPartialUPartialLambda(bps, 2) test_obs_f3_du_dl = custom_ops.FullPartialUPartialLambda(bps, 2) obs = [ test_obs, test_obs_f2, test_obs_du_dl, test_obs_f2_du_dl, test_obs_f3_du_dl ] for o in obs: ctxt.add_observable(o) for step in range(num_steps): print("comparing step", step) ctxt.step(lamb) test_x_t = ctxt.get_x_t() test_v_t = ctxt.get_v_t() test_du_dx_t = ctxt._get_du_dx_t_minus_1() # test_u_t = ctxt._get_u_t_minus_1() # np.testing.assert_allclose(test_u_t, ref_all_us[step]) np.testing.assert_allclose(test_du_dx_t, ref_all_du_dxs[step]) np.testing.assert_allclose(test_x_t, ref_all_xs[step]) ref_avg_du_dls = np.mean(ref_all_du_dls, axis=0) ref_avg_du_dls_f2 = np.mean(ref_all_du_dls[::2], axis=0) np.testing.assert_allclose(test_obs_du_dl.avg_du_dl(), ref_avg_du_dls) np.testing.assert_allclose(test_obs_f2_du_dl.avg_du_dl(), ref_avg_du_dls_f2) full_du_dls = test_obs_f3_du_dl.full_du_dl() assert len(full_du_dls) == np.ceil(num_steps / 2) np.testing.assert_allclose(np.mean(full_du_dls), ref_avg_du_dls_f2) ref_avg_du_dps = np.mean(ref_all_du_dps, axis=0) ref_avg_du_dps_f2 = np.mean(ref_all_du_dps[::2], axis=0) # the fixed point accumulator makes it hard to converge some of these # if the derivative is super small - in which case they probably don't matter # anyways np.testing.assert_allclose(test_obs.avg_du_dp()[:, 0], ref_avg_du_dps[:, 0], 1.5e-6) np.testing.assert_allclose(test_obs.avg_du_dp()[:, 1], ref_avg_du_dps[:, 1], 1.5e-6) np.testing.assert_allclose(test_obs.avg_du_dp()[:, 2], ref_avg_du_dps[:, 2], 5e-5)
def Simulate(self, request, context): if request.precision == 'single': precision = np.float32 elif request.precision == 'double': precision = np.float64 else: raise Exception("Unknown precision") simulation = pickle.loads(request.simulation) bps = [] pots = [] for potential in simulation.potentials: bps.append(potential.bound_impl()) # get the bound implementation intg = simulation.integrator.impl() ctxt = custom_ops.Context(simulation.x, simulation.v, simulation.box, intg, bps) lamb = request.lamb for step, minimize_lamb in enumerate( np.linspace(1.0, lamb, request.prep_steps)): ctxt.step(minimize_lamb) energies = [] frames = [] if request.observe_du_dl_freq > 0: du_dl_obs = custom_ops.AvgPartialUPartialLambda( bps, request.observe_du_dl_freq) ctxt.add_observable(du_dl_obs) if request.observe_du_dp_freq > 0: du_dps = [] # for name, bp in zip(names, bps): # if name == 'LennardJones' or name == 'Electrostatics': for bp in bps: du_dp_obs = custom_ops.AvgPartialUPartialParam( bp, request.observe_du_dp_freq) ctxt.add_observable(du_dp_obs) du_dps.append(du_dp_obs) # dynamics for step in range(request.prod_steps): if step % 100 == 0: u = ctxt._get_u_t_minus_1() energies.append(u) if request.n_frames > 0: interval = max(1, request.prod_steps // request.n_frames) if step % interval == 0: frames.append(ctxt.get_x_t()) ctxt.step(lamb) frames = np.array(frames) if request.observe_du_dl_freq > 0: avg_du_dls = du_dl_obs.avg_du_dl() else: avg_du_dls = None if request.observe_du_dp_freq > 0: avg_du_dps = [] for obs in du_dps: avg_du_dps.append(obs.avg_du_dp()) else: avg_du_dps = None return service_pb2.SimulateReply( avg_du_dls=pickle.dumps(avg_du_dls), avg_du_dps=pickle.dumps(avg_du_dps), energies=pickle.dumps(energies), frames=pickle.dumps(frames), )
def test_fwd_mode(self): """ This test ensures that we can reverse-mode differentiate observables that are dU_dlambdas of each state. We provide adjoints with respect to each computed dU/dLambda. """ np.random.seed(4321) N = 8 D = 3 x0 = np.random.rand(N, D).astype(dtype=np.float64) * 2 E = 2 lambda_plane_idxs = np.random.randint(low=0, high=2, size=N, dtype=np.int32) lambda_offset_idxs = np.random.randint(low=0, high=2, size=N, dtype=np.int32) params, ref_nrg_fn, test_nrg = prepare_nb_system( x0, E, lambda_plane_idxs, lambda_offset_idxs, p_scale=3.0, # cutoff=0.5, cutoff=1.0, ) masses = np.random.rand(N) v0 = np.random.rand(x0.shape[0], x0.shape[1]) num_steps = 5 temperature = 300 dt = 2e-3 friction = 0.0 ca, cbs, ccs = langevin_coefficients(temperature, dt, friction, masses) # not convenient to simulate identical trajectories otherwise assert (ccs == 0).all() lamb = np.random.rand() lambda_windows = np.array([lamb + 0.05, lamb, lamb - 0.05]) def integrate_once_through(x_t, v_t, box, params): dU_dx_fn = jax.grad(ref_nrg_fn, argnums=(0, )) dU_dp_fn = jax.grad(ref_nrg_fn, argnums=(1, )) dU_dl_fn = jax.grad(ref_nrg_fn, argnums=(3, )) all_du_dls = [] all_du_dps = [] all_xs = [] all_du_dxs = [] all_us = [] all_lambda_us = [] for step in range(num_steps): u = ref_nrg_fn(x_t, params, box, lamb) all_us.append(u) du_dl = dU_dl_fn(x_t, params, box, lamb)[0] all_du_dls.append(du_dl) du_dp = dU_dp_fn(x_t, params, box, lamb)[0] all_du_dps.append(du_dp) du_dx = dU_dx_fn(x_t, params, box, lamb)[0] all_du_dxs.append(du_dx) all_xs.append(x_t) lus = [] for lamb_u in lambda_windows: lus.append(ref_nrg_fn(x_t, params, box, lamb_u)) all_lambda_us.append(lus) noise = np.random.randn(*v_t.shape) v_mid = v_t + np.expand_dims(cbs, axis=-1) * du_dx v_t = ca * v_mid + np.expand_dims(ccs, axis=-1) * noise x_t += 0.5 * dt * (v_mid + v_t) # note that we do not calculate the du_dl of the last frame. return all_xs, all_du_dxs, all_du_dps, all_du_dls, all_us, all_lambda_us box = np.eye(3) * 3.0 # when we have multiple parameters, we need to set this up correctly ( ref_all_xs, ref_all_du_dxs, ref_all_du_dps, ref_all_du_dls, ref_all_us, ref_all_lambda_us, ) = integrate_once_through(x0, v0, box, params) intg = custom_ops.LangevinIntegrator(dt, ca, cbs, ccs, 1234) bp = test_nrg.bind(params).bound_impl(precision=np.float64) bps = [bp] ctxt = custom_ops.Context(x0, v0, box, intg, bps) for step in range(num_steps): print("comparing step", step) test_x_t = ctxt.get_x_t() np.testing.assert_allclose(test_x_t, ref_all_xs[step]) ctxt.step(lamb) test_du_dx_t = ctxt._get_du_dx_t_minus_1() # test_u_t = ctxt._get_u_t_minus_1() # np.testing.assert_allclose(test_u_t, ref_all_us[step]) np.testing.assert_allclose(test_du_dx_t, ref_all_du_dxs[step]) # test the multiple_steps method ctxt_2 = custom_ops.Context(x0, v0, box, intg, bps) lambda_schedule = np.ones(num_steps) * lamb du_dl_interval = 3 x_interval = 2 start_box = ctxt_2.get_box() test_du_dls, test_xs, test_boxes = ctxt_2.multiple_steps( lambda_schedule, du_dl_interval, x_interval) end_box = ctxt_2.get_box() np.testing.assert_allclose(test_du_dls, ref_all_du_dls[::du_dl_interval]) np.testing.assert_allclose(test_xs, ref_all_xs[::x_interval]) np.testing.assert_array_equal(start_box, end_box) for i in range(test_boxes.shape[0]): np.testing.assert_array_equal(start_box, test_boxes[i]) self.assertEqual(test_boxes.shape[0], test_xs.shape[0]) self.assertEqual(test_boxes.shape[1], D) self.assertEqual(test_boxes.shape[2], test_xs.shape[2]) # test the multiple_steps_U method ctxt_3 = custom_ops.Context(x0, v0, box, intg, bps) u_interval = 3 test_us, test_xs, test_boxes = ctxt_3.multiple_steps_U( lamb, num_steps, lambda_windows, u_interval, x_interval) np.testing.assert_array_almost_equal(ref_all_lambda_us[::u_interval], test_us) np.testing.assert_array_almost_equal(ref_all_xs[::x_interval], test_xs) test_us, test_xs, test_boxes = ctxt_3.multiple_steps_U( lamb, num_steps, np.array([], dtype=np.float64), u_interval, x_interval) assert test_us.shape == (2, 0)
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()
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 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 minimize_host_4d(romol, host_system, host_coords, ff, box): """ Insert romol into a host system via 4D decoupling under a Langevin thermostat. The ligand coordinates are fixed during this, and only host_coordinates are minimized. Parameters ---------- romol: ROMol Ligand to be inserted. It must be embedded. 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. Returns ------- np.ndarray This returns minimized host_coords. """ host_bps, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.2) # keep the ligand rigid ligand_masses = [a.GetMass()*100000 for a in romol.GetAtoms()] combined_masses = np.concatenate([host_masses, ligand_masses]) ligand_coords = get_romol_conf(romol) combined_coords = np.concatenate([host_coords, ligand_coords]) num_host_atoms = host_coords.shape[0] final_potentials = [] for bp in host_bps: if isinstance(bp, potentials.Nonbonded): host_p = bp else: final_potentials.append(bp) 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]], ] for fn, handles in tuples: params, potential = fn(*[h.params for h in handles]) final_potentials.append(potential.bind(params)) seed = 2020 intg = LangevinIntegrator( 300.0, 1.5e-3, 1.0, combined_masses, seed ).impl() x0 = combined_coords v0 = np.zeros_like(x0) u_impls = [] for bp in final_potentials: fn = bp.bound_impl(precision=np.float32) u_impls.append(fn) # context components: positions, velocities, box, integrator, energy fxns ctxt = custom_ops.Context( x0, v0, box, intg, u_impls ) for lamb in np.linspace(1.0, 0, 1000): ctxt.step(lamb) return ctxt.get_x_t()[:num_host_atoms]
def run_leg( orig_host_coords, orig_guest_coords, combined_bps, combined_masses, host_box, guest_name, leg_type, host_mol, guest_mol, outdir, 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, TRANSITION_STEPS ) for step, lamb in enumerate(insertion_lambda_schedule): ctxt.step(lamb) if step % 100 == 0: report.report_step( ctxt, step, lamb, host_box, combined_bps, u_impls, guest_name, TRANSITION_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(TRANSITION_STEPS))), f"{leg_type}-ins", ) if step in (0, int(TRANSITION_STEPS/2), TRANSITION_STEPS-1): if report.too_much_force(ctxt, lamb, host_box, combined_bps, u_impls): return # equilibrate for step in range(EQ1_STEPS): ctxt.step(MIN_LAMBDA) if step % 1000 == 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 step in (0, int(EQ1_STEPS/2), EQ1_STEPS-1): if report.too_much_force(ctxt, MIN_LAMBDA, host_box, combined_bps, u_impls): return # equilibrate more & shoot off deletion jobs for step in range(EQ2_STEPS): ctxt.step(MIN_LAMBDA) if step % 1000 == 0: report.report_step( ctxt, step, MIN_LAMBDA, host_box, combined_bps, u_impls, guest_name, EQ2_STEPS, 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(step).zfill(len(str(EQ2_STEPS))), f"{leg_type}-eq2", ) if report.too_much_force(ctxt, MIN_LAMBDA, host_box, combined_bps, u_impls): return do_deletion( ctxt.get_x_t(), ctxt.get_v_t(), combined_bps, combined_masses, host_box, guest_name, leg_type, )
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 every 100 steps of insertion (outdir/<guest_name>/<guest_name>_<step>.[pdb/sdf]) A pdb & sdf file every 1000 steps of equilibration (outdir/<guest_name>/<guest_name>_<step>.[pdb/sdf]) stdout every 100(0) 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. """ 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...") # TODO: return topology from builders.build_protein_system ( solvated_host_system, solvated_host_coords, _, _, host_box, solvated_topology, ) = builders.build_protein_system(host_pdbfile) # sometimes water boxes are sad. Should be minimized first; this is a workaround host_box += np.eye(3) * 0.1 print("host box", host_box) solvated_host_pdb = os.path.join(outdir, "solvated_host.pdb") 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) final_host_potentials = [] host_potentials, host_masses = openmm_deserializer.deserialize_system( solvated_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_host_potentials.append(bp) # 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 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_top = topology.BaseTopology(guest_mol, ff) # combine host & guest hgt = topology.HostGuestTopology(host_nb_bp, guest_base_top) # 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]] combined_bps = list(final_host_potentials) # instantiate the vjps while parameterizing (forward pass) for fn, handle in bonded_tuples: params, potential = fn(handle.params) combined_bps.append(potential.bind(params)) nb_params, nb_potential = hgt.parameterize_nonbonded( ff.q_handle.params, ff.lj_handle.params) combined_bps.append(nb_potential.bind(nb_params)) guest_masses = [a.GetMass() for a in guest_mol.GetAtoms()] combined_masses = np.concatenate([host_masses, guest_masses]) x0 = np.concatenate([solvated_host_coords, orig_guest_coords]) v0 = np.zeros_like(x0) print( f"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) # collect a du_dl calculation once every other step subsample_freq = 2 du_dl_obs = custom_ops.FullPartialUPartialLambda( u_impls, subsample_freq) ctxt.add_observable(du_dl_obs) # insert guest insertion_lambda_schedule = np.linspace(max_lambda, 0.0, insertion_steps) calc_work = True for step, lamb in enumerate(insertion_lambda_schedule): ctxt.step(lamb) if step % 100 == 0: 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))), f"ins", ) if step in (0, int(insertion_steps / 2), insertion_steps - 1): if report.too_much_force(ctxt, lamb, host_box, combined_bps, u_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(), insertion_lambda_schedule[::subsample_freq]) 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') 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))), f"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 test_molecular_ideal_gas(): """ References ---------- OpenMM testIdealGas https://github.com/openmm/openmm/blob/d8ef57fed6554ec95684e53768188e1f666405c9/tests/TestMonteCarloBarostat.h#L86-L140 """ # simulation parameters initial_waterbox_width = 3.0 * unit.nanometer timestep = 1.5 * unit.femtosecond collision_rate = 1.0 / unit.picosecond n_moves = 10000 barostat_interval = 5 seed = 2021 # thermodynamic parameters temperatures = np.array([300, 600, 1000]) * unit.kelvin pressure = 100.0 * unit.bar # very high pressure, to keep the expected volume small # generate an alchemical system of a waterbox + alchemical ligand: # effectively discard ligands by running in AbsoluteFreeEnergy mode at lambda = 1.0 mol_a = hif2a_ligand_pair.mol_a ff = hif2a_ligand_pair.ff complex_system, complex_coords, complex_box, complex_top = build_water_system( initial_waterbox_width.value_in_unit(unit.nanometer)) min_complex_coords = minimize_host_4d([mol_a], complex_system, complex_coords, ff, complex_box) afe = AbsoluteFreeEnergy(mol_a, ff) _unbound_potentials, _sys_params, masses, coords = afe.prepare_host_edge( ff.get_ordered_params(), complex_system, min_complex_coords) # drop the nonbonded potential unbound_potentials = _unbound_potentials[:-1] sys_params = _sys_params[:-1] # get list of molecules for barostat by looking at bond table harmonic_bond_potential = unbound_potentials[0] bond_list = get_bond_list(harmonic_bond_potential) group_indices = get_group_indices(bond_list) volume_trajs = [] relative_tolerance = 1e-2 initial_relative_box_perturbation = 2 * relative_tolerance n_molecules = complex_top.getNumResidues() bound_potentials = [] for params, unbound_pot in zip(sys_params, unbound_potentials): bp = unbound_pot.bind(np.asarray(params)) bound_potentials.append(bp) u_impls = [] for bp in bound_potentials: bp_impl = bp.bound_impl(precision=np.float32) u_impls.append(bp_impl) # expected volume md_pressure_unit = ENERGY_UNIT / DISTANCE_UNIT**3 pressure_in_md = ( pressure * unit.AVOGADRO_CONSTANT_NA).value_in_unit(md_pressure_unit) expected_volume_in_md = (n_molecules + 1) * BOLTZ * temperatures.value_in_unit( unit.kelvin) / pressure_in_md for i, temperature in enumerate(temperatures): # define a thermostat integrator = LangevinIntegrator( temperature.value_in_unit(unit.kelvin), timestep.value_in_unit(unit.picosecond), collision_rate.value_in_unit(unit.picosecond**-1), masses, seed, ) integrator_impl = integrator.impl() v_0 = sample_velocities(masses * unit.amu, temperature) # rescale the box to be approximately the desired box volume already rescaler = CentroidRescaler(group_indices) initial_volume = compute_box_volume(complex_box) initial_center = compute_box_center(complex_box) length_scale = ((1 + initial_relative_box_perturbation) * expected_volume_in_md[i] / initial_volume)**(1.0 / 3) new_coords = rescaler.scale_centroids(coords, initial_center, length_scale) new_box = complex_box * length_scale baro = custom_ops.MonteCarloBarostat( new_coords.shape[0], pressure.value_in_unit(unit.bar), temperature.value_in_unit(unit.kelvin), group_indices, barostat_interval, u_impls, seed, ) ctxt = custom_ops.Context(new_coords, v_0, new_box, integrator_impl, u_impls, barostat=baro) vols = [] for move in range(n_moves // barostat_interval): ctxt.multiple_steps(np.ones(barostat_interval)) new_box = ctxt.get_box() volume = np.linalg.det(new_box) vols.append(volume) volume_trajs.append(vols) equil_time = len(volume_trajs[0]) // 2 # TODO: don't hard-code this? actual_volume_in_md = np.array( [np.mean(volume_traj[equil_time:]) for volume_traj in volume_trajs]) np.testing.assert_allclose(actual=actual_volume_in_md, desired=expected_volume_in_md, rtol=relative_tolerance)