def test_harmonic_angle(self, n_particles=64, n_angles=25, dim=3): """Randomly connect triples of particles, then validate the resulting HarmonicAngle force""" np.random.seed(125) x = self.get_random_coords(n_particles, dim) atom_idxs = np.arange(n_particles) params = np.random.rand(n_angles, 2).astype(np.float64) angle_idxs = [] for _ in range(n_angles): angle_idxs.append( np.random.choice(atom_idxs, size=3, replace=False)) angle_idxs = np.array(angle_idxs, dtype=np.int32) lamb = 0.0 box = np.eye(3) * 100 # specific to harmonic angle force relative_tolerance_at_precision = {np.float32: 2e-5, np.float64: 1e-9} for precision, rtol in relative_tolerance_at_precision.items(): test_potential = potentials.HarmonicAngle(angle_idxs) ref_potential = functools.partial(bonded.harmonic_angle, angle_idxs=angle_idxs) self.compare_forces(x, params, box, lamb, ref_potential, test_potential, rtol, precision=precision)
def get_harmonic_angle(n_atoms, n_bonds): atom_idxs = np.arange(n_atoms) params = np.random.rand(n_bonds, 2).astype(np.float64) bond_idxs = [] for _ in range(n_bonds): bond_idxs.append(np.random.choice(atom_idxs, size=3, replace=False)) bond_idxs = np.array(bond_idxs, dtype=np.int32) lamb_mult = np.random.randint(-5, 5, size=n_bonds, dtype=np.int32) lamb_offset = np.random.randint(-5, 5, size=n_bonds, dtype=np.int32) return potentials.HarmonicAngle(bond_idxs, lamb_mult, lamb_offset), params
def parameterize_harmonic_angle(self, ff_params): params, idxs = self.ff.ha_handle.partial_parameterize(ff_params, self.mol) return params, potentials.HarmonicAngle(idxs)
def combine_potentials(ff_handlers, guest_mol, host_system, precision): """ This function is responsible for figuring out how to take two separate hamiltonians and combining them into one sensible alchemical system. Parameters ---------- ff_handlers: list of forcefield handlers Small molecule forcefield handlers guest_mol: Chem.ROMol RDKit molecule host_system: openmm.System Host system to be deserialized precision: np.float32 or np.float64 Numerical precision of the functional form Returns ------- tuple Returns a list of lib.potentials objects, combined masses, and a list of their corresponding vjp_fns back into the forcefield """ host_potentials, host_masses = openmm_deserializer.deserialize_system( host_system, precision, cutoff=1.0) host_nb_bp = None combined_potentials = [] combined_vjp_fns = [] 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: combined_potentials.append(bp) combined_vjp_fns.append([]) guest_masses = np.array([a.GetMass() for a in guest_mol.GetAtoms()], dtype=np.float64) num_guest_atoms = len(guest_masses) num_host_atoms = len(host_masses) combined_masses = np.concatenate([host_masses, guest_masses]) for handle in ff_handlers: results = handle.parameterize(guest_mol) if isinstance(handle, bonded.HarmonicBondHandler): bond_idxs, (bond_params, vjp_fn) = results bond_idxs += num_host_atoms combined_potentials.append( potentials.HarmonicBond(bond_idxs, precision=precision).bind(bond_params)) combined_vjp_fns.append([(handle, vjp_fn)]) elif isinstance(handle, bonded.HarmonicAngleHandler): angle_idxs, (angle_params, vjp_fn) = results angle_idxs += num_host_atoms combined_potentials.append( potentials.HarmonicAngle( angle_idxs, precision=precision).bind(angle_params)) combined_vjp_fns.append([(handle, vjp_fn)]) elif isinstance(handle, bonded.ProperTorsionHandler): torsion_idxs, (torsion_params, vjp_fn) = results torsion_idxs += num_host_atoms combined_potentials.append( potentials.PeriodicTorsion( torsion_idxs, precision=precision).bind(torsion_params)) combined_vjp_fns.append([(handle, vjp_fn)]) elif isinstance(handle, bonded.ImproperTorsionHandler): torsion_idxs, (torsion_params, vjp_fn) = results torsion_idxs += num_host_atoms combined_potentials.append( potentials.PeriodicTorsion( torsion_idxs, precision=precision).bind(torsion_params)) combined_vjp_fns.append([(handle, vjp_fn)]) elif isinstance(handle, nonbonded.AM1CCCHandler): charge_handle = handle guest_charge_params, guest_charge_vjp_fn = results elif isinstance(handle, nonbonded.LennardJonesHandler): guest_lj_params, guest_lj_vjp_fn = results lj_handle = handle else: print("Warning: skipping handler", handle) pass # process nonbonded terms combined_nb_params, (charge_vjp_fn, lj_vjp_fn) = nonbonded_vjps( guest_charge_params, guest_charge_vjp_fn, guest_lj_params, guest_lj_vjp_fn, host_nb_bp.params) # these vjp_fns take in adjoints of combined_params and returns derivatives # appropriate to the underlying handler combined_vjp_fns.append([(charge_handle, charge_vjp_fn), (lj_handle, lj_vjp_fn)]) # tbd change scale 14 for electrostatics guest_exclusion_idxs, guest_scale_factors = nonbonded.generate_exclusion_idxs( guest_mol, scale12=1.0, scale13=1.0, scale14=0.5) # allow the ligand to be alchemically decoupled # a value of one indicates that we allow the atom to be adjusted by the lambda value guest_lambda_offset_idxs = np.ones(len(guest_masses), dtype=np.int32) # use same scale factors until we modify 1-4s for electrostatics guest_scale_factors = np.stack([guest_scale_factors, guest_scale_factors], axis=1) combined_lambda_offset_idxs = np.concatenate( [host_nb_bp.get_lambda_offset_idxs(), guest_lambda_offset_idxs]) combined_exclusion_idxs = np.concatenate([ host_nb_bp.get_exclusion_idxs(), guest_exclusion_idxs + num_host_atoms ]) combined_scales = np.concatenate( [host_nb_bp.get_scale_factors(), guest_scale_factors]) combined_beta = 2.0 combined_cutoff = 1.0 # nonbonded cutoff combined_potentials.append( potentials.Nonbonded( combined_exclusion_idxs, combined_scales, combined_lambda_offset_idxs, combined_beta, combined_cutoff, precision=precision, ).bind(combined_nb_params)) return combined_potentials, combined_masses, combined_vjp_fns
def from_rdkit(cls, mol, ff_handlers): """ Initialize a system from an RDKit ROMol. Parameters ---------- mol: Chem.ROMol RDKit ROMol. Should have graphical hydrogens in the topology. ff_handlers: list of forcefield handlers. openforcefield small molecule handlers. """ masses = np.array([a.GetMass() for a in mol.GetAtoms()], dtype=np.float64) bound_potentials = [] for handle in ff_handlers: results = handle.parameterize(mol) if isinstance(handle, bonded.HarmonicBondHandler): bond_params, bond_idxs = results bound_potentials.append(potentials.HarmonicBond(bond_idxs).bind(bond_params)) elif isinstance(handle, bonded.HarmonicAngleHandler): angle_params, angle_idxs = results bound_potentials.append(potentials.HarmonicAngle(angle_idxs).bind(angle_params)) elif isinstance(handle, bonded.ProperTorsionHandler): torsion_params, torsion_idxs = results bound_potentials.append(potentials.PeriodicTorsion(torsion_idxs).bind(torsion_params)) elif isinstance(handle, bonded.ImproperTorsionHandler): torsion_params, torsion_idxs = results bound_potentials.append(potentials.PeriodicTorsion(torsion_idxs).bind(torsion_params)) elif isinstance(handle, nonbonded.AM1CCCHandler): charge_handle = handle charge_params = results elif isinstance(handle, nonbonded.LennardJonesHandler): lj_params = results lj_handle = handle else: print("WARNING: skipping handler", handle) pass lambda_plane_idxs = np.zeros(len(masses), dtype=np.int32) lambda_offset_idxs = np.zeros(len(masses), dtype=np.int32) exclusion_idxs, scale_factors = nonbonded.generate_exclusion_idxs( mol, scale12=1.0, scale13=1.0, scale14=0.5 ) # use same scale factors until we modify 1-4s for electrostatics scale_factors = np.stack([scale_factors, scale_factors], axis=1) # (ytz) fix this later to not be so hard coded alpha = 2.0 # same as ewald alpha cutoff = 1.0 # nonbonded cutoff qlj_params = jnp.concatenate([ jnp.reshape(charge_params, (-1, 1)), jnp.reshape(lj_params, (-1, 2)) ], axis=1) bound_potentials.append(potentials.Nonbonded( exclusion_idxs, scale_factors, lambda_plane_idxs, lambda_offset_idxs, alpha, cutoff).bind(qlj_params)) return cls(masses, bound_potentials)
def deserialize_system(system, cutoff): """ Deserialize an OpenMM XML file Parameters ---------- system: openmm.System A system object to be deserialized Returns ------- list of lib.Potential, masses Note: We add a small epsilon (1e-3) to all zero eps values to prevent a singularity from occuring in the lennard jones derivatives """ masses = [] for p in range(system.getNumParticles()): masses.append(value(system.getParticleMass(p))) N = len(masses) # this should not be a dict since we may have more than one instance of a given # force. bps = [] for force in system.getForces(): if isinstance(force, mm.HarmonicBondForce): bond_idxs = [] bond_params = [] for b_idx in range(force.getNumBonds()): src_idx, dst_idx, length, k = force.getBondParameters(b_idx) length = value(length) k = value(k) bond_idxs.append([src_idx, dst_idx]) bond_params.append((k, length)) bond_idxs = np.array(bond_idxs, dtype=np.int32) bond_params = np.array(bond_params, dtype=np.float64) bps.append(potentials.HarmonicBond(bond_idxs).bind(bond_params)) if isinstance(force, mm.HarmonicAngleForce): angle_idxs = [] angle_params = [] for a_idx in range(force.getNumAngles()): src_idx, mid_idx, dst_idx, angle, k = force.getAngleParameters( a_idx) angle = value(angle) k = value(k) angle_idxs.append([src_idx, mid_idx, dst_idx]) angle_params.append((k, angle)) angle_idxs = np.array(angle_idxs, dtype=np.int32) angle_params = np.array(angle_params, dtype=np.float64) bps.append(potentials.HarmonicAngle(angle_idxs).bind(angle_params)) if isinstance(force, mm.PeriodicTorsionForce): torsion_idxs = [] torsion_params = [] for t_idx in range(force.getNumTorsions()): a_idx, b_idx, c_idx, d_idx, period, phase, k = force.getTorsionParameters( t_idx) phase = value(phase) k = value(k) torsion_params.append((k, phase, period)) torsion_idxs.append([a_idx, b_idx, c_idx, d_idx]) torsion_idxs = np.array(torsion_idxs, dtype=np.int32) torsion_params = np.array(torsion_params, dtype=np.float64) bps.append( potentials.PeriodicTorsion(torsion_idxs).bind(torsion_params)) if isinstance(force, mm.NonbondedForce): num_atoms = force.getNumParticles() charge_params = [] lj_params = [] for a_idx in range(num_atoms): charge, sig, eps = force.getParticleParameters(a_idx) charge = value(charge) * np.sqrt(constants.ONE_4PI_EPS0) sig = value(sig) eps = value(eps) # increment eps by 1e-3 if we have eps==0 to avoid a singularity in parameter derivatives # override default amber types # this doesn't work for water! # if eps == 0: # print("Warning: overriding eps by 1e-3 to avoid a singularity") # eps += 1e-3 # charge_params.append(charge_idx) charge_params.append(charge) lj_params.append((sig, eps)) charge_params = np.array(charge_params, dtype=np.float64) # print("Protein net charge:", np.sum(np.array(global_params)[charge_param_idxs])) lj_params = np.array(lj_params, dtype=np.float64) # 1 here means we fully remove the interaction # 1-2, 1-3 # scale_full = insert_parameters(1.0, 20) # 1-4, remove half of the interaction # scale_half = insert_parameters(0.5, 21) exclusion_idxs = [] scale_factors = [] all_sig = lj_params[:, 0] all_eps = lj_params[:, 1] # validate exclusions/exceptions to make sure they make sense for a_idx in range(force.getNumExceptions()): # tbd charge scale factors src, dst, new_cp, new_sig, new_eps = force.getExceptionParameters( a_idx) new_sig = value(new_sig) new_eps = value(new_eps) src_sig = all_sig[src] dst_sig = all_sig[dst] src_eps = all_eps[src] dst_eps = all_eps[dst] expected_sig = (src_sig + dst_sig) / 2 expected_eps = np.sqrt(src_eps * dst_eps) exclusion_idxs.append([src, dst]) # sanity check this (expected_eps can be zero), redo this thing # the lj_scale factor measures how much we *remove* if expected_eps == 0: if new_eps == 0: lj_scale_factor = 1 else: raise RuntimeError( "Divide by zero in epsilon calculation") else: lj_scale_factor = 1 - new_eps / expected_eps scale_factors.append(lj_scale_factor) # tbd fix charge_scale_factors using new_cp if new_eps != 0: np.testing.assert_almost_equal(expected_sig, new_sig) exclusion_idxs = np.array(exclusion_idxs, dtype=np.int32) lambda_plane_idxs = np.zeros(N, dtype=np.int32) lambda_offset_idxs = np.zeros(N, dtype=np.int32) # cutoff = 1000.0 nb_params = np.concatenate( [np.expand_dims(charge_params, axis=1), lj_params], axis=1) # optimizations nb_params[:, 1] = nb_params[:, 1] / 2 nb_params[:, 2] = np.sqrt(nb_params[:, 2]) beta = 2.0 # erfc correction # use the same scale factors for electrostatics and lj scale_factors = np.stack([scale_factors, scale_factors], axis=1) bps.append( potentials.Nonbonded(exclusion_idxs, scale_factors, lambda_plane_idxs, lambda_offset_idxs, beta, cutoff).bind(nb_params)) # nrg_fns.append(('Exclusions', (exclusion_idxs, scale_factors, es_scale_factors))) return bps, masses