def test_nonbonded_optimal_map(self): """Similar test as test_nonbonbed, ie. assert that coordinates and nonbonded parameters can be averaged in benzene -> phenol transformation. However, use the maximal mapping possible.""" # map benzene H to phenol O, leaving a dangling phenol H core = np.array( [[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6]], dtype=np.int32) st = topology.SingleTopology(self.mol_a, self.mol_b, core, self.ff) x_a = get_romol_conf(self.mol_a) x_b = get_romol_conf(self.mol_b) # test interpolation of coordinates. x_src, x_dst = st.interpolate_params(x_a, x_b) x_avg = np.mean([x_src, x_dst], axis=0) assert x_avg.shape == (st.get_num_atoms(), 3) np.testing.assert_array_equal((x_a[:7] + x_b[:7]) / 2, x_avg[:7]) # core parts np.testing.assert_array_equal(x_b[-1], x_avg[7]) # dangling H params, vjp_fn, pot_c = jax.vjp(st.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True) vjp_fn(np.random.rand(*params.shape)) assert params.shape == (2 * st.get_num_atoms(), 3) # qlj # test interpolation of parameters bt_a = topology.BaseTopology(self.mol_a, self.ff) qlj_a, pot_a = bt_a.parameterize_nonbonded(self.ff.q_handle.params, self.ff.lj_handle.params) bt_b = topology.BaseTopology(self.mol_b, self.ff) qlj_b, pot_b = bt_b.parameterize_nonbonded(self.ff.q_handle.params, self.ff.lj_handle.params) n_base_params = len( params ) // 2 # params is actually interpolated, so its 2x number of base params # qlj_c = np.mean([params[:n_base_params], params[n_base_params:]], axis=0) params_src = params[:n_base_params] params_dst = params[n_base_params:] # core testing np.testing.assert_array_equal(qlj_a[:7], params_src[:7]) np.testing.assert_array_equal(qlj_b[:7], params_dst[:7]) # r-group atoms in A are all part of the core. so no testing is needed. # test r-group in B np.testing.assert_array_equal(qlj_b[7], params_dst[8]) np.testing.assert_array_equal(np.array([0, qlj_b[7][1], 0]), params_src[8])
def __init__(self, mol, ff): """ VacuumState allows us to enable/disable various parts of a forcefield so that we can more easily sample across rotational barriers in the vacuum. Parameters ---------- mol: Chem.ROMol rdkit molecule ff: Forcefield forcefield """ self.mol = mol bt = topology.BaseTopology(mol, ff) self.bond_params, self.hb_potential = bt.parameterize_harmonic_bond(ff.hb_handle.params) self.angle_params, self.ha_potential = bt.parameterize_harmonic_angle(ff.ha_handle.params) self.proper_torsion_params, self.pt_potential = bt.parameterize_proper_torsion(ff.pt_handle.params) ( self.improper_torsion_params, self.it_potential, ) = bt.parameterize_improper_torsion(ff.it_handle.params) self.nb_params, self.nb_potential = bt.parameterize_nonbonded(ff.q_handle.params, ff.lj_handle.params) self.box = None self.lamb = 0.0
def test_base_topology_conversion_ring_torsion(): # test that the conversion protocol behaves as intended on a # simple linked cycle. ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py") mol = Chem.MolFromSmiles("C1CC1C1CC1") vanilla_mol_top = topology.BaseTopology(mol, ff) vanilla_torsion_params, _ = vanilla_mol_top.parameterize_proper_torsion( ff.pt_handle.params) mol_top = topology.BaseTopologyConversion(mol, ff) conversion_torsion_params, torsion_potential = mol_top.parameterize_proper_torsion( ff.pt_handle.params) np.testing.assert_array_equal(vanilla_torsion_params, conversion_torsion_params) assert torsion_potential.get_lambda_mult() is None assert torsion_potential.get_lambda_offset() is None vanilla_qlj_params, _ = vanilla_mol_top.parameterize_nonbonded( ff.q_handle.params, ff.lj_handle.params) qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded( ff.q_handle.params, ff.lj_handle.params) assert isinstance(nonbonded_potential, potentials.NonbondedInterpolated) src_qlj_params = qlj_params[:len(qlj_params) // 2] dst_qlj_params = qlj_params[len(qlj_params) // 2:] np.testing.assert_array_equal(vanilla_qlj_params, src_qlj_params) np.testing.assert_array_equal(topology.standard_qlj_typer(mol), dst_qlj_params)
def test_base_topology_standard_decoupling(): # this class is typically used in the second step of the RABFE protocol for the solvent leg. # we expected the charges to be zero, and the lj parameters to be standardized. In addition, # the torsions should be turned off. ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py") mol = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1O")) vanilla_mol_top = topology.BaseTopology(mol, ff) vanilla_torsion_params, _ = vanilla_mol_top.parameterize_proper_torsion( ff.pt_handle.params) mol_top = topology.BaseTopologyStandardDecoupling(mol, ff) decouple_torsion_params, torsion_potential = mol_top.parameterize_proper_torsion( ff.pt_handle.params) np.testing.assert_array_equal(vanilla_torsion_params, decouple_torsion_params) # in the conversion phase, torsions that bridge the two rings should be set to # be alchemically turned off. # is_in_ring = [1, 1, 1, 1, 1, 1, 0, 0] combined_decouple_torsion_params, combined_torsion_potential = mol_top.parameterize_periodic_torsion( ff.pt_handle.params, ff.it_handle.params) assert len(combined_torsion_potential.get_lambda_mult()) == len( combined_torsion_potential.get_idxs()) assert len(combined_torsion_potential.get_lambda_mult()) == len( combined_torsion_potential.get_lambda_offset()) # impropers should always be turned on. # num_proper_torsions = len(torsion_potential.get_idxs()) assert np.all(combined_torsion_potential.get_lambda_mult() == 0) assert np.all(combined_torsion_potential.get_lambda_offset() == 1) qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded( ff.q_handle.params, ff.lj_handle.params) assert not isinstance(nonbonded_potential, potentials.NonbondedInterpolated) np.testing.assert_array_equal(topology.standard_qlj_typer(mol), qlj_params) np.testing.assert_array_equal(nonbonded_potential.get_lambda_plane_idxs(), np.zeros(mol.GetNumAtoms(), dtype=np.int32)) np.testing.assert_array_equal(nonbonded_potential.get_lambda_offset_idxs(), np.ones(mol.GetNumAtoms(), dtype=np.int32))
def __init__(self, mol, ff): """ Compute the absolute free energy of a molecule via 4D decoupling. Parameters ---------- mol: rdkit mol Ligand to be decoupled ff: ff.Forcefield Ligand forcefield """ self.mol = mol self.ff = ff self.top = topology.BaseTopology(mol, ff)
def test_dual_topology_rhfe(): # used in testing the relative hydration protocol. The nonbonded charges and epsilons are reduced # to half strength ff = Forcefield.load_from_file("smirnoff_1_1_0_sc.py") mol_a = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1O")) mol_b = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1F")) mol_c = Chem.CombineMols(mol_a, mol_b) mol_top = topology.DualTopologyRHFE(mol_a, mol_b, ff) C = mol_a.GetNumAtoms() + mol_b.GetNumAtoms() ref_qlj_params, _ = topology.BaseTopology(mol_c, ff).parameterize_nonbonded( ff.q_handle.params, ff.lj_handle.params) qlj_params, nonbonded_potential = mol_top.parameterize_nonbonded( ff.q_handle.params, ff.lj_handle.params) assert isinstance(nonbonded_potential, potentials.NonbondedInterpolated) src_qlj_params = qlj_params[:len(qlj_params) // 2] dst_qlj_params = qlj_params[len(qlj_params) // 2:] np.testing.assert_array_equal(src_qlj_params[:, 0], ref_qlj_params[:, 0] / 2) np.testing.assert_array_equal(src_qlj_params[:, 1], ref_qlj_params[:, 1]) np.testing.assert_array_equal(src_qlj_params[:, 2], ref_qlj_params[:, 2] / 2) np.testing.assert_array_equal(dst_qlj_params, ref_qlj_params) combined_lambda_plane_idxs = nonbonded_potential.get_lambda_plane_idxs() combined_lambda_offset_idxs = nonbonded_potential.get_lambda_offset_idxs() A = mol_a.GetNumAtoms() B = mol_b.GetNumAtoms() C = mol_c.GetNumAtoms() np.testing.assert_array_equal(combined_lambda_plane_idxs, np.zeros(C)) np.testing.assert_array_equal(combined_lambda_offset_idxs[:A], np.zeros(A)) np.testing.assert_array_equal(combined_lambda_offset_idxs[A:], np.ones(B))
def test_nonbonded(self): # leaving benzene H unmapped, and phenol OH unmapped core = np.array( [ [0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], ], dtype=np.int32, ) st = topology.SingleTopology(self.mol_a, self.mol_b, core, self.ff) x_a = get_romol_conf(self.mol_a) x_b = get_romol_conf(self.mol_b) # test interpolation of coordinates. x_src, x_dst = st.interpolate_params(x_a, x_b) x_avg = np.mean([x_src, x_dst], axis=0) assert x_avg.shape == (st.get_num_atoms(), 3) np.testing.assert_array_equal((x_a[:6] + x_b[:6]) / 2, x_avg[:6]) # C np.testing.assert_array_equal(x_a[6], x_avg[6]) # H np.testing.assert_array_equal(x_b[6:], x_avg[7:]) # OH # NOTE: unused result st.parameterize_nonbonded(self.ff.q_handle.params, self.ff.lj_handle.params) params, vjp_fn, pot_c = jax.vjp(st.parameterize_nonbonded, self.ff.q_handle.params, self.ff.lj_handle.params, has_aux=True) vjp_fn(np.random.rand(*params.shape)) assert params.shape == (2 * st.get_num_atoms(), 3) # qlj # test interpolation of parameters bt_a = topology.BaseTopology(self.mol_a, self.ff) qlj_a, pot_a = bt_a.parameterize_nonbonded(self.ff.q_handle.params, self.ff.lj_handle.params) bt_b = topology.BaseTopology(self.mol_b, self.ff) qlj_b, pot_b = bt_b.parameterize_nonbonded(self.ff.q_handle.params, self.ff.lj_handle.params) n_base_params = len( params ) // 2 # params is actually interpolated, so its 2x number of base params # qlj_c = np.mean([params[:n_base_params], params[n_base_params:]], axis=0) params_src = params[:n_base_params] params_dst = params[n_base_params:] # core testing np.testing.assert_array_equal(qlj_a[:6], params_src[:6]) np.testing.assert_array_equal(qlj_b[:6], params_dst[:6]) # test r-group in A np.testing.assert_array_equal(qlj_a[6], params_src[6]) np.testing.assert_array_equal(np.array([0, qlj_a[6][1], 0]), params_dst[6]) # test r-group in B np.testing.assert_array_equal(qlj_b[6:], params_dst[7:]) np.testing.assert_array_equal( np.array([[0, qlj_b[6][1], 0], [0, qlj_b[7][1], 0]]), params_src[7:])
num_host_atoms = host_coords.shape[0] 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],
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()