def prepare_host_edge(self, ff_params, host_system): """ Prepares the host-edge system Parameters ---------- ff_params: tuple of np.array Exactly equal to bond_params, angle_params, proper_params, improper_params, charge_params, lj_params host_system: openmm.System openmm System object to be deserialized Returns ------- 3-tuple unbound_potentials, system_params, combined_masses """ ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()] ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()] host_bps, host_masses = openmm_deserializer.deserialize_system( host_system, cutoff=1.2) hgt = topology.HostGuestTopology(host_bps, self.top) final_params, final_potentials = self._get_system_params_and_potentials( ff_params, hgt) combined_masses = np.concatenate( [host_masses, ligand_masses_a, ligand_masses_b]) return final_potentials, final_params, combined_masses
def prepare_host_edge(self, ff_params, host_system, host_coords): """ Prepares the host-edge system Parameters ---------- ff_params: tuple of np.array Exactly equal to bond_params, angle_params, proper_params, improper_params, charge_params, lj_params host_system: openmm.System openmm System object to be deserialized host_coords: np.array Nx3 array of atomic coordinates Returns ------- 4 tuple unbound_potentials, system_params, combined_masses, combined_coords """ ligand_masses_a = [a.GetMass() for a in self.mol_a.GetAtoms()] ligand_masses_b = [b.GetMass() for b in self.mol_b.GetAtoms()] # extract the 0th conformer ligand_coords_a = get_romol_conf(self.mol_a) ligand_coords_b = get_romol_conf(self.mol_b) host_bps, host_masses = openmm_deserializer.deserialize_system( host_system, cutoff=1.2) hgt = topology.HostGuestTopology(host_bps, self.top) final_params, final_potentials = self._get_system_params_and_potentials( ff_params, hgt) if isinstance(self.top, topology.SingleTopology): combined_masses = np.concatenate([ host_masses, np.mean(self.top.interpolate_params(ligand_masses_a, ligand_masses_b), axis=0) ]) combined_coords = np.concatenate([ host_coords, np.mean(self.top.interpolate_params(ligand_coords_a, ligand_coords_b), axis=0) ]) else: combined_masses = np.concatenate( [host_masses, ligand_masses_a, ligand_masses_b]) combined_coords = np.concatenate( [host_coords, ligand_coords_a, ligand_coords_b]) return final_potentials, final_params, combined_masses, combined_coords
final_potentials = [] final_vjp_and_handles = [] # keep the bonded terms in the host the same. # but we keep the nonbonded term for a subsequent modification for bp in host_bps: if isinstance(bp, potentials.Nonbonded): host_p = bp else: final_potentials.append(bp) final_vjp_and_handles.append(None) ff = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py") 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]], ] # instantiate the vjps while parameterizing (forward pass) for fn, handles in tuples: params, vjp_fn, potential = jax.vjp(fn, *[h.params for h in handles], has_aux=True)
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 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()