Exemplo n.º 1
0
def grosberg_angle(sim_object, triplets, k=1.5, name="grosberg_angle"):
    """Adds stiffness according to the Grosberg paper.
    (Halverson, Jonathan D., et al. "Molecular dynamics simulation study of
     nonconcatenated ring polymers in a melt. I. Statics."
     The Journal of chemical physics 134 (2011): 204904.)

    Parameters are synchronized with normal stiffness

    If k is an array, it has to be of the length N.
    Xth value then specifies stiffness of the angle centered at
    monomer number X.
    Values for ends of the chain will be simply ignored.

    Parameters
    ----------

    k : float or N-long list of floats
        Synchronized with regular stiffness.
        Default value is very flexible, as in Grosberg paper.
        Default value maximizes entanglement length.

    """
    k = _to_array_1d(k, len(triplets))

    force = openmm.CustomAngleForce("GRk * kT * (1 - cos(theta - 3.141592))")

    force.name = name
    force.addGlobalParameter("kT", sim_object.kT)
    force.addPerAngleParameter("GRk")

    for triplet_idx, (p1, p2, p3) in enumerate(triplets):
        force.addAngle(p1, p2, p3, [k[triplet_idx]])

    return force
Exemplo n.º 2
0
def angle_force(sim_object,
                triplets,
                k=1.5,
                theta_0=np.pi,
                name="angle",
                override_checks=False):
    """Adds harmonic angle bonds. k specifies energy in kT at one radian
    If k is an array, it has to be of the length N.
    Xth value then specifies stiffness of the angle centered at
    monomer number X.
    Values for ends of the chain will be simply ignored.

    Parameters
    ----------

    k : float or list of length N
        Stiffness of the bond.
        If list, then determines the stiffness of the i-th triplet
        Potential is k * alpha^2 * 0.5 * kT
    
    theta_0 : float or list of length N 
              Equilibrium angle of the bond. By default it is np.pi. 

    override_checks: bool
        If True then do not check that no bonds are repeated.
        False by default.
    """

    # check for repeated triplets
    if not override_checks:
        _check_angle_bonds(triplets)

    k = _to_array_1d(k, len(triplets))
    theta_0 = _to_array_1d(theta_0, len(triplets))

    energy = "kT*angK * (theta - angT0) * (theta - angT0) * (0.5)"
    force = openmm.CustomAngleForce(energy)
    force.name = name

    force.addGlobalParameter("kT", sim_object.kT)
    force.addPerAngleParameter("angK")
    force.addPerAngleParameter("angT0")

    for triplet_idx, (p1, p2, p3) in enumerate(triplets):
        force.addAngle(int(p1), int(p2), int(p3),
                       (float(k[triplet_idx]), float(theta_0[triplet_idx])))

    return force
Exemplo n.º 3
0
def _get_internal_from_omm(atom_coords, bond_coords, angle_coords, torsion_coords):
    import copy
    #master system, will be used for all three
    sys = openmm.System()
    platform = openmm.Platform.getPlatformByName("Reference")
    for i in range(4):
        sys.addParticle(1.0*unit.amu)

    #first, the bond length:
    bond_sys = openmm.System()
    bond_sys.addParticle(1.0*unit.amu)
    bond_sys.addParticle(1.0*unit.amu)
    bond_force = openmm.CustomBondForce("r")
    bond_force.addBond(0, 1, [])
    bond_sys.addForce(bond_force)
    bond_integrator = openmm.VerletIntegrator(1*unit.femtoseconds)
    bond_context = openmm.Context(bond_sys, bond_integrator, platform)
    bond_context.setPositions([atom_coords, bond_coords])
    bond_state = bond_context.getState(getEnergy=True)
    r = bond_state.getPotentialEnergy()
    del bond_sys, bond_context, bond_integrator

    #now, the angle:
    angle_sys = copy.deepcopy(sys)
    angle_force = openmm.CustomAngleForce("theta")
    angle_force.addAngle(0,1,2,[])
    angle_sys.addForce(angle_force)
    angle_integrator = openmm.VerletIntegrator(1*unit.femtoseconds)
    angle_context = openmm.Context(angle_sys, angle_integrator, platform)
    angle_context.setPositions([atom_coords, bond_coords, angle_coords, torsion_coords])
    angle_state = angle_context.getState(getEnergy=True)
    theta = angle_state.getPotentialEnergy()
    del angle_sys, angle_context, angle_integrator

    #finally, the torsion:
    torsion_sys = copy.deepcopy(sys)
    torsion_force = openmm.CustomTorsionForce("theta")
    torsion_force.addTorsion(0,1,2,3,[])
    torsion_sys.addForce(torsion_force)
    torsion_integrator = openmm.VerletIntegrator(1*unit.femtoseconds)
    torsion_context = openmm.Context(torsion_sys, torsion_integrator, platform)
    torsion_context.setPositions([atom_coords, bond_coords, angle_coords, torsion_coords])
    torsion_state = torsion_context.getState(getEnergy=True)
    phi = torsion_state.getPotentialEnergy()
    del torsion_sys, torsion_context, torsion_integrator

    return r, theta, phi
Exemplo n.º 4
0
    def _add_angle_force_terms(self):
        core_energy_expression = '(K/2)*(theta-theta0)^2;'
        core_energy_expression += 'K = k*scale_factor;'  # linearly interpolate spring constant
        core_energy_expression += self.scaling_expression()

        # Create the force and add the relevant parameters
        custom_core_force = openmm.CustomAngleForce(core_energy_expression)
        custom_core_force.addPerAngleParameter('theta0')
        custom_core_force.addPerAngleParameter('k')
        custom_core_force.addPerAngleParameter('identifier')

        custom_core_force.addGlobalParameter('solute_scale', 1.0)
        custom_core_force.addGlobalParameter('inter_scale', 1.0)

        self._out_system.addForce(custom_core_force)
        self._out_system_forces[
            custom_core_force.__class__.__name__] = custom_core_force
Exemplo n.º 5
0
    def modify_harmonic_angles(self):
        """
        turn the harmonic angles into a custom angle force

        lambda_protocol :
        - 'lambda_MM_angles' : 1 -> 0
        - 'lambda_scale' : beta / beta0
        """
        self._alchemical_to_old_angles = {}
        angle_expression = 'lambda_MM_angles * lambda_scale * (k/2)*(theta-theta0)^2;'
        custom_angle_force = openmm.CustomAngleForce(angle_expression)

        #add the global params
        custom_angle_force.addGlobalParameter('lambda_MM_angles', 1.)
        custom_angle_force.addGlobalParameter('lambda_scale', 1.)

        #add the perangleparams
        custom_angle_force.addPerAngleParameter('theta0')
        custom_angle_force.addPerAngleParameter('k')

        #now to iterate over the angles.
        for idx in range(self._system_forces['HarmonicAngleForce'].getNumAngles()):
            p1, p2, p3, theta0, k = self._system_forces['HarmonicAngleForce'].getAngleParameters(idx)
            if self.is_in_alchemical_region({p1,p2,p3}):
                #first thing to do is to zero the force from the `HarmonicAngleForce`
                self._system_forces['HarmonicAngleForce'].setAngleParameters(idx, p1, p2, p3, theta0, k*0.0)
                self._endstate_system_forces['HarmonicAngleForce'].setAngleParameters(idx, p1, p2, p3, theta0, k*0.0)

                #then add it to the custom angle force
                custom_angle_idx = custom_angle_force.addAngle(p1, p2, p3, [theta0, k])

                #add to the alchemical bonds dict for bookkeeping
                self._alchemical_to_old_angles[custom_angle_idx] = idx

        #then add the custom bond force to the system
        if self._system_forces['HarmonicAngleForce'].usesPeriodicBoundaryConditions():
            custom_angle_force.setUsesPeriodicBoundaryConditions(True)
        self._system.addForce(custom_angle_force)
Exemplo n.º 6
0
def _add_extras(system, bonds, restricted_angles, torsions):
    # add the extra bonds
    if bonds:
        f = [f for f in system.getForces() if isinstance(f, mm.HarmonicBondForce)][0]
        for bond in bonds:
            f.addBond(bond.i, bond.j, bond.length, bond.force_constant)

    # add the extra restricted_angles
    if restricted_angles:
        # create the new force for restricted angles
        f = mm.CustomAngleForce(
            "0.5 * k_ra * (theta - theta0_ra)^2 / sin(theta * 3.1459 / 180)"
        )
        f.addPerAngleParameter("k_ra")
        f.addPerAngleParameter("theta0_ra")
        for angle in restricted_angles:
            f.addAngle(
                angle.i,
                angle.j,
                angle.k,
                (angle.force_constant, angle.angle),
            )
        system.addForce(f)

    # add the extra torsions
    if torsions:
        f = [f for f in system.getForces() if isinstance(f, mm.PeriodicTorsionForce)][0]
        for tors in torsions:
            f.addTorsion(
                tors.i,
                tors.j,
                tors.k,
                tors.l,
                tors.multiplicity,
                tors.phase,
                tors.energy,
            )
Exemplo n.º 7
0
    def _addAngleToSystem(self, syst, moleculeType, bondedTypes, atomBonds,
                          baseAtomIndex):

        degToRad = math.pi / 180
        for fields in moleculeType.angles:
            atoms = [int(x) - 1 for x in fields[:3]]
            types = tuple(bondedTypes[i] for i in atoms)
            if len(fields) >= 6:
                params = fields[4:]
            elif types in self.angleTypes:
                params = self.angleTypes[types][4:]
            elif types[::-1] in self.angleTypes:
                params = self.angleTypes[types[::-1]][4:]
            else:
                raise ValueError('No parameters specified for angle: ' +
                                 fields[0] + ', ' + fields[1] + ', ' +
                                 fields[2])

            #angles = mm.HarmonicAngleForce()
            theta = float(params[0]) * degToRad
            if int(fields[3]) == 2:
                gromosAngle = mm.CustomAngleForce(
                    '0.5*k*(cos(theta)-cos(theta0))^2')
                gromosAngle.addPerAngleParameter('theta0')
                gromosAngle.addPerAngleParameter('k')
                syst.addForce(gromosAngle)
                gromosAngle.addAngle(baseAtomIndex + atoms[0],
                                     baseAtomIndex + atoms[1],
                                     baseAtomIndex + atoms[2],
                                     [theta, float(params[1])])

            elif int(fields[3]) == 1:
                angles = mm.HarmonicAngleForce()
                angles.addAngle(baseAtomIndex + atoms[0],
                                baseAtomIndex + atoms[1],
                                baseAtomIndex + atoms[2], theta,
                                float(params[1]))
Exemplo n.º 8
0
def angle_force(sim_object, triplets, k=1.5, theta_0=np.pi, name="angle"):
    """Adds harmonic angle bonds. k specifies energy in kT at one radian
    If k is an array, it has to be of the length N.
    Xth value then specifies stiffness of the angle centered at
    monomer number X.
    Values for ends of the chain will be simply ignored.

    Parameters
    ----------

    k : float or list of length N
        Stiffness of the bond.
        If list, then determines the stiffness of the i-th triplet
        Potential is k * alpha^2 * 0.5 * kT
    
    theta_0 : float or list of length N 
              Equilibrium angle of the bond. By default it is np.pi. 
              
        
    """

    k = _to_array_1d(k, len(triplets))
    theta_0 = _to_array_1d(theta_0, len(triplets))

    energy = "kT*angK * (theta - angT0) * (theta - angT0) * (0.5)"
    force = openmm.CustomAngleForce(energy)
    force.name = name

    force.addGlobalParameter("kT", sim_object.kT)
    force.addPerAngleParameter("angK")
    force.addPerAngleParameter("angT0")

    for triplet_idx, (p1, p2, p3) in enumerate(triplets):
        force.addAngle(p1, p2, p3, [k[triplet_idx], theta_0[triplet_idx]])

    return force
Exemplo n.º 9
0
def add_molecule_to_system(system,
                           molecule_system,
                           core_atoms,
                           variant,
                           atoms_to_exclude=[]):
    """
    Add the valence terms for the molecule from molecule_system.

    Parameters
    ----------
    system : simtk.openmm.System
       The system object to which the valence terms are to be added.
    molecule_system : simtk.openmm.System
       The system object from which core valence terms are to be taken.
    core_atoms : list of int
       The list of atom indices within molecule_system corresponding to core atoms.
    variant : int
       The variant index of this molecule if not a core fragment, or 0 if this is a core fragment and only core atoms are to be added.

    Returns
    -------
    mapping : dict of int
       mapping[index] is the atom index in `system` corresponding to atom `index` within `molecule_system`.

    """
    def _createCustomNonbondedForce(self,
                                    system,
                                    molecule_system,
                                    softcore_alpha=0.5,
                                    softcore_beta=12 * unit.angstrom**2):
        """
        Create alchemically-modified version of NonbondedForce.

        Parameters
        ----------
        system : simtk.openmm.System
            Alchemically-modified system being built.  This object will be modified.
        molecule_system : simtk.openmm.System
            Source molecule system to copy from.
        softcore_alpha : float, optional, default = 0.5
            Alchemical softcore parameter for Lennard-Jones.
        softcore_beta : simtk.unit.Quantity with units compatible with angstroms**2, optional, default = 12*angstrom**2
            Alchemical softcore parameter for electrostatics.

        TODO
        ----
        Try using a single, common "reff" effective softcore distance for both Lennard-Jones and Coulomb.

        """

        alchemical_atom_indices = self.ligand_atoms

        # Create a copy of the NonbondedForce to handle non-alchemical interactions.
        nonbonded_force = copy.deepcopy(reference_force)
        system.addForce(nonbonded_force)

        # Create CustomNonbondedForce objects to handle softcore interactions between alchemically-modified system and rest of system.

        # Create atom groups.
        natoms = system.getNumParticles()
        atomset1 = set(
            alchemical_atom_indices)  # only alchemically-modified atoms
        atomset2 = set(range(system.getNumParticles())
                       )  # all atoms, including alchemical region

        # CustomNonbondedForce energy expression.
        sterics_energy_expression = ""
        electrostatics_energy_expression = ""

        # Select functional form based on nonbonded method.
        method = reference_force.getNonbondedMethod()
        if method in [openmm.NonbondedForce.NoCutoff]:
            # soft-core Lennard-Jones
            sterics_energy_expression += "U_sterics = lambda_sterics*4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;"
            # soft-core Coulomb
            electrostatics_energy_expression += "U_electrostatics = ONE_4PI_EPS0*lambda_electrostatics*chargeprod/reff_electrostatics;"
        elif method in [
                openmm.NonbondedForce.CutoffPeriodic,
                openmm.NonbondedForce.CutoffNonPeriodic
        ]:
            # soft-core Lennard-Jones
            sterics_energy_expression += "U_sterics = lambda_sterics*4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;"
            # reaction-field electrostatics
            epsilon_solvent = reference_force.getReactionFieldDielectric()
            r_cutoff = reference_force.getCutoffDistance()
            electrostatics_energy_expression += "U_electrostatics = lambda_electrostatics*ONE_4PI_EPS0*chargeprod*(reff_electrostatics^(-1) + k_rf*reff_electrostatics^2 - c_rf);"
            k_rf = r_cutoff**(-3) * ((epsilon_solvent - 1) /
                                     (2 * epsilon_solvent + 1))
            c_rf = r_cutoff**(-1) * ((3 * epsilon_solvent) /
                                     (2 * epsilon_solvent + 1))
            electrostatics_energy_expression += "k_rf = %f;" % (
                k_rf / k_rf.in_unit_system(unit.md_unit_system).unit)
            electrostatics_energy_expression += "c_rf = %f;" % (
                c_rf / c_rf.in_unit_system(unit.md_unit_system).unit)
        elif method in [
                openmm.NonbondedForce.PME, openmm.NonbondedForce.Ewald
        ]:
            # soft-core Lennard-Jones
            sterics_energy_expression += "U_sterics = lambda_sterics*4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;"
            # Ewald direct-space electrostatics
            [alpha_ewald, nx, ny, nz] = reference_force.getPMEParameters()
            if alpha_ewald == 0.0:
                # If alpha is 0.0, alpha_ewald is computed by OpenMM from from the error tolerance.
                delta = reference_force.getEwaldErrorTolerance()
                r_cutoff = reference_force.getCutoffDistance()
                alpha_ewald = np.sqrt(-np.log(2 * delta)) / r_cutoff
            electrostatics_energy_expression += "U_electrostatics = lambda_electrostatics*ONE_4PI_EPS0*chargeprod*erfc(alpha_ewald*reff_electrostatics)/reff_electrostatics;"
            electrostatics_energy_expression += "alpha_ewald = %f;" % (
                alpha_ewald /
                alpha_ewald.in_unit_system(unit.md_unit_system).unit)
            # TODO: Handle reciprocal-space electrostatics
        else:
            raise Exception("Nonbonded method %s not supported yet." %
                            str(method))

        # Add additional definitions common to all methods.
        sterics_energy_expression += "reff_sterics = sigma*((softcore_alpha*(1.-lambda_sterics) + (r/sigma)^6))^(1/6);"  # effective softcore distance for sterics
        sterics_energy_expression += "softcore_alpha = %f;" % softcore_alpha
        electrostatics_energy_expression += "reff_electrostatics = sqrt(softcore_beta*(1.-lambda_electrostatics) + r^2);"  # effective softcore distance for electrostatics
        electrostatics_energy_expression += "softcore_beta = %f;" % (
            softcore_beta /
            softcore_beta.in_unit_system(unit.md_unit_system).unit)
        electrostatics_energy_expression += "ONE_4PI_EPS0 = %f;" % ONE_4PI_EPS0  # already in OpenMM units

        # Define mixing rules.
        sterics_mixing_rules = ""
        sterics_mixing_rules += "epsilon = sqrt(epsilon1*epsilon2);"  # mixing rule for epsilon
        sterics_mixing_rules += "sigma = 0.5*(sigma1 + sigma2);"  # mixing rule for sigma
        electrostatics_mixing_rules = ""
        electrostatics_mixing_rules += "chargeprod = charge1*charge2;"  # mixing rule for charges

        # Create CustomNonbondedForce to handle interactions between alchemically-modified atoms and rest of system.
        electrostatics_custom_nonbonded_force = openmm.CustomNonbondedForce(
            "U_electrostatics;" + electrostatics_energy_expression +
            electrostatics_mixing_rules)
        electrostatics_custom_nonbonded_force.addGlobalParameter(
            "lambda_electrostatics", 1.0)
        electrostatics_custom_nonbonded_force.addPerParticleParameter(
            "charge")  # partial charge
        sterics_custom_nonbonded_force = openmm.CustomNonbondedForce(
            "U_sterics;" + sterics_energy_expression + sterics_mixing_rules)
        sterics_custom_nonbonded_force.addGlobalParameter(
            "lambda_sterics", 1.0)
        sterics_custom_nonbonded_force.addPerParticleParameter(
            "sigma")  # Lennard-Jones sigma
        sterics_custom_nonbonded_force.addPerParticleParameter(
            "epsilon")  # Lennard-Jones epsilon

        # Set parameters to match reference force.
        sterics_custom_nonbonded_force.setUseSwitchingFunction(
            nonbonded_force.getUseSwitchingFunction())
        electrostatics_custom_nonbonded_force.setUseSwitchingFunction(False)
        sterics_custom_nonbonded_force.setCutoffDistance(
            nonbonded_force.getCutoffDistance())
        electrostatics_custom_nonbonded_force.setCutoffDistance(
            nonbonded_force.getCutoffDistance())
        sterics_custom_nonbonded_force.setSwitchingDistance(
            nonbonded_force.getSwitchingDistance())
        sterics_custom_nonbonded_force.setUseLongRangeCorrection(
            nonbonded_force.getUseDispersionCorrection())
        electrostatics_custom_nonbonded_force.setUseLongRangeCorrection(False)

        # Set periodicity and cutoff parameters corresponding to reference Force.
        if nonbonded_force.getNonbondedMethod() in [
                openmm.NonbondedForce.Ewald, openmm.NonbondedForce.PME,
                openmm.NonbondedForce.CutoffPeriodic
        ]:
            sterics_custom_nonbonded_force.setNonbondedMethod(
                openmm.CustomNonbondedForce.CutoffPeriodic)
            electrostatics_custom_nonbonded_force.setNonbondedMethod(
                openmm.CustomNonbondedForce.CutoffPeriodic)
        else:
            sterics_custom_nonbonded_force.setNonbondedMethod(
                nonbonded_force.getNonbondedMethod())
            electrostatics_custom_nonbonded_force.setNonbondedMethod(
                nonbonded_force.getNonbondedMethod())

        # Restrict interaction evaluation to be between alchemical atoms and rest of environment.
        # TODO: Exclude intra-alchemical region if we are separately handling that through a separate CustomNonbondedForce for decoupling.
        sterics_custom_nonbonded_force.addInteractionGroup(atomset1, atomset2)
        electrostatics_custom_nonbonded_force.addInteractionGroup(
            atomset1, atomset2)

        # Add custom forces.
        system.addForce(sterics_custom_nonbonded_force)
        system.addForce(electrostatics_custom_nonbonded_force)

        # Create CustomBondForce to handle exceptions for both kinds of interactions.
        custom_bond_force = openmm.CustomBondForce(
            "U_sterics + U_electrostatics;" + sterics_energy_expression +
            electrostatics_energy_expression)
        custom_bond_force.addGlobalParameter("lambda_electrostatics", 1.0)
        custom_bond_force.addGlobalParameter("lambda_sterics", 1.0)
        custom_bond_force.addPerBondParameter("chargeprod")  # charge product
        custom_bond_force.addPerBondParameter(
            "sigma")  # Lennard-Jones effective sigma
        custom_bond_force.addPerBondParameter(
            "epsilon")  # Lennard-Jones effective epsilon
        system.addForce(custom_bond_force)

        # Move NonbondedForce particle terms for alchemically-modified particles to CustomNonbondedForce.
        for particle_index in range(nonbonded_force.getNumParticles()):
            # Retrieve parameters.
            [charge, sigma,
             epsilon] = nonbonded_force.getParticleParameters(particle_index)
            # Add parameters to custom force handling interactions between alchemically-modified atoms and rest of system.
            sterics_custom_nonbonded_force.addParticle([sigma, epsilon])
            electrostatics_custom_nonbonded_force.addParticle([charge])
            # Turn off Lennard-Jones contribution from alchemically-modified particles.
            if particle_index in alchemical_atom_indices:
                nonbonded_force.setParticleParameters(particle_index,
                                                      0 * charge, sigma,
                                                      0 * epsilon)

        # Move NonbondedForce exception terms for alchemically-modified particles to CustomNonbondedForce/CustomBondForce.
        for exception_index in range(nonbonded_force.getNumExceptions()):
            # Retrieve parameters.
            [iatom, jatom, chargeprod, sigma,
             epsilon] = nonbonded_force.getExceptionParameters(exception_index)
            # Exclude this atom pair in CustomNonbondedForce.
            sterics_custom_nonbonded_force.addExclusion(iatom, jatom)
            electrostatics_custom_nonbonded_force.addExclusion(iatom, jatom)
            # Move exceptions involving alchemically-modified atoms to CustomBondForce.
            if self.annihilate_sterics and (iatom in alchemical_atom_indices
                                            ) and (jatom
                                                   in alchemical_atom_indices):
                # Add special CustomBondForce term to handle alchemically-modified Lennard-Jones exception.
                custom_bond_force.addBond(iatom, jatom,
                                          [chargeprod, sigma, epsilon])
                # Zero terms in NonbondedForce.
                nonbonded_force.setExceptionParameters(exception_index, iatom,
                                                       jatom, 0 * chargeprod,
                                                       sigma, 0 * epsilon)

        # TODO: Add back NonbondedForce terms for alchemical system needed in case of decoupling electrostatics or sterics via second CustomBondForce.
        # TODO: Also need to change current CustomBondForce to not alchemically disappear system.

        return

    # Build dict of forces.
    def create_force_dict(system):
        return {
            system.getForce(index).__class__.__name__: system.getForce(index)
            for index in range(system.getNumForces())
        }

    molecule_forces = create_force_dict(molecule_system)
    forces = create_force_dict(system)

    # Create Custom*Force classes if necessary.
    if 'CustomBondForce' not in forces:
        energy_expression = 'lambda*(K/2)*(r-length)^2;'
        energy_expression += 'lambda = (1-alchemical_lambda)*delta(variant) + alchemical_lambda*delta(variant-alchemical_variant);'
        custom_force = mm.CustomBondForce(energy_expression)
        custom_force.addGlobalParameter('alchemical_lambda', 0.0)
        custom_force.addGlobalParameter('alchemical_variant', 0.0)
        custom_force.addPerBondParameter('variant')
        custom_force.addPerBondParameter('length')
        custom_force.addPerBondParameter('K')
        system.addForce(custom_force)
        forces['CustomBondForce'] = custom_force
    if 'CustomAngleForce' not in forces:
        energy_expression = 'lambda*(K/2)*(theta-theta0)^2;'
        energy_expression += 'lambda = (1-alchemical_lambda)*delta(variant) + alchemical_lambda*delta(variant-alchemical_variant);'
        custom_force = mm.CustomAngleForce(energy_expression)
        custom_force.addGlobalParameter('alchemical_lambda', 0.0)
        custom_force.addGlobalParameter('alchemical_variant', 0.0)
        custom_force.addPerAngleParameter('variant')
        custom_force.addPerAngleParameter('theta0')
        custom_force.addPerAngleParameter('K')
        system.addForce(custom_force)
        forces['CustomAngleForce'] = custom_force
    if 'CustomTorsionForce' not in forces:
        energy_expression = 'lambda*K*(1+cos(periodicity*theta-phase));'
        energy_expression += 'lambda = (1-alchemical_lambda)*delta(variant) + alchemical_lambda*delta(variant-alchemical_variant);'
        custom_force = mm.CustomTorsionForce(energy_expression)
        custom_force.addGlobalParameter('alchemical_lambda', 0.0)
        custom_force.addGlobalParameter('alchemical_variant', 0.0)
        custom_force.addPerTorsionParameter('variant')
        custom_force.addPerTorsionParameter('periodicity')
        custom_force.addPerTorsionParameter('phase')
        custom_force.addPerTorsionParameter('K')
        system.addForce(custom_force)
        forces['CustomTorsionForce'] = custom_force
    if 'CustomNonbondedForce' not in forces:
        # DEBUG
        # TODO: Create proper CustomNonbondedForce here.
        energy_expression = "0.0;"
        custom_force = mm.CustomNonbondedForce(energy_expression)
        custom_force.addGlobalParameter('alchemical_lambda', 0.0)
        custom_force.addGlobalParameter('alchemical_variant', 0.0)
        custom_force.addPerParticleParameter('variant')
        custom_force.addPerParticleParameter('charge')
        custom_force.addPerParticleParameter('sigma')
        custom_force.addPerParticleParameter('epsilon')
        system.addForce(custom_force)
        forces['CustomNonbondedForce'] = custom_force

        # Add parameters for existing particles.
        for index in range(system.getNumParticles()):
            [charge, sigma,
             epsilon] = forces['NonbondedForce'].getParticleParameters(index)
            custom_force.addParticle([0, charge, sigma, epsilon])

    # Add particles to system.
    mapping = dict()  # mapping[index_in_molecule] = index_in_system
    for index_in_molecule in range(molecule_system.getNumParticles()):
        # Add all atoms, unless we're adding the core, in which case we just add core atoms.
        if (variant) or (index_in_molecule in core_atoms):
            # TODO: We may want to make masses lighter.
            index_in_system = system.addParticle(
                system.getParticleMass(index_in_molecule))
            mapping[index_in_molecule] = index_in_system

    # Constraints are not supported.
    # TODO: Later, consider supporting some constraints, such as those within core and those within variants.
    if (molecule_system.getNumConstraints() > 0):
        raise Exception(
            "Constraints are not supported for alchemically modified molecule."
        )

    # Process forces.
    # Valence terms involving only core atoms are created as Custom*Force classes where lambda=0 activates the "core" image and lambda=1 activates the "variant" image.
    for (force_name, force) in molecule_forces.iteritems():
        print force_name
        if force_name == 'HarmonicBondForce':
            for index in range(force.getNumBonds()):
                [atom_i, atom_j, length, K] = force.getBondParameters(index)
                if set([atom_i, atom_j]).issubset(core_atoms):
                    forces['CustomBondForce'].addBond(mapping[atom_i],
                                                      mapping[atom_j],
                                                      [variant, length, K])
                elif (variant):
                    forces[force_name].addBond(mapping[atom_i],
                                               mapping[atom_j], length, K)

        elif force_name == 'HarmonicAngleForce':
            for index in range(force.getNumAngles()):
                [atom_i, atom_j, atom_k, theta0,
                 K] = force.getAngleParameters(index)
                if set([atom_i, atom_j, atom_k]).issubset(core_atoms):
                    forces['CustomAngleForce'].addAngle(
                        mapping[atom_i], mapping[atom_j], mapping[atom_k],
                        [variant, theta0, K])
                elif (variant):
                    forces[force_name].addAngle(mapping[atom_i],
                                                mapping[atom_j],
                                                mapping[atom_k], theta0, K)

        elif force_name == 'PeriodicTorsionForce':
            for index in range(force.getNumTorsions()):
                [atom_i, atom_j, atom_k, atom_l, periodicity, phase,
                 K] = force.getTorsionParameters(index)
                if set([atom_i, atom_j, atom_k, atom_l]).issubset(core_atoms):
                    forces['CustomTorsionForce'].addTorsion(
                        mapping[atom_i], mapping[atom_j], mapping[atom_k],
                        mapping[atom_l], [variant, periodicity, phase, K])
                elif (variant):
                    forces[force_name].addTorsion(mapping[atom_i],
                                                  mapping[atom_j],
                                                  mapping[atom_k],
                                                  mapping[atom_l], periodicity,
                                                  phase, K)

        elif force_name == 'NonbondedForce':
            for index in range(force.getNumParticles()):
                # TODO: Nonbonded terms will have to be handled as CustomNonbondedForce terms.
                [charge, sigma, epsilon] = force.getParticleParameters(index)
                if set([index]).issubset(core_atoms):
                    forces[force_name].addParticle(0.0 * charge, sigma,
                                                   0.0 * epsilon)
                    forces['CustomNonbondedForce'].addParticle(
                        [variant, charge, sigma, epsilon])
                elif (variant):
                    forces[force_name].addParticle(0.0 * charge, sigma,
                                                   0.0 * epsilon)
                    forces['CustomNonbondedForce'].addParticle(
                        [variant, charge, sigma, epsilon])

            for index in range(force.getNumExceptions()):
                [atom_i, atom_j, chargeProd, sigma,
                 epsilon] = force.getExceptionParameters(index)
                if set([atom_i, atom_j]).issubset(core_atoms):
                    # TODO: Nonbonded exceptions will have to be handled as CustomBondForce terms.
                    forces[force_name].addException(
                        mapping[atom_i], mapping[atom_j],
                        0.0 * unit.elementary_charge**2, 1.0 * unit.angstrom,
                        0.0 * unit.kilocalories_per_mole)
                elif (variant):
                    forces[force_name].addException(mapping[atom_i],
                                                    mapping[atom_j],
                                                    chargeProd, sigma, epsilon)

        # TODO: Add GB force processing.

    # Add exclusions to previous variants and core.
    for atom_i in mapping.values():
        for atom_j in atoms_to_exclude:
            forces['NonbondedForce'].addException(
                atom_i, atom_j, 0.0 * unit.elementary_charge**2,
                1.0 * unit.angstrom, 0.0 * unit.kilocalories_per_mole)
            forces['CustomNonbondedForce'].addExclusion(atom_i, atom_j)

    print system.getNumParticles(), forces['NonbondedForce'].getNumParticles()

    return mapping
Exemplo n.º 10
0
    def redefine_angle(self,
                       topology,
                       residue,
                       atom1,
                       atom2,
                       atom3,
                       angle,
                       K=None,
                       group=1):
        """
        Changes the equilibrium value of a specified angle for integration within its original
        time scale. The difference between the original and the redefined angle potentials is
        evaluated at another time scale.

        Parameters
        ----------
            topology : openmm.Topology
                The topology corresponding to the original system.
            residue : str
                A name or regular expression to identify the residue which contains the redefined
                angle.
            atom1 : str
                A name or regular expression to identify the first atom that makes the angle.
            atom2 : str
                A name or regular expression to identify the second atom that makes the angle.
            atom3 : str
                A name or regular expression to identify the third atom that makes the angle.
            angle : unit.Quantity
                The redifined equilibrium angle value for integration at the shortest time scale.
            K : unit.Quantity, optional, default=None
                The harmonic force constant for the angle. If this is `None`, then the original
                value will be maintained.
            group : int, optional, default=1
                The force group with which the difference between the original and the redefined
                angle potentials must be evaluated.

        """
        resname = [atom.residue.name for atom in topology.atoms()]
        atom = [atom.name for atom in topology.atoms()]
        r_regex = re.compile(residue)
        a_regex = [re.compile(a) for a in [atom1, atom2, atom3]]

        def r_match(*args):
            return all(r_regex.match(resname[j]) for j in args)

        def a_match(*args):
            return all(a_regex[i].match(atom[j]) for i, j in enumerate(args))

        angle_list = []
        for force in self.getForces():
            if isinstance(force, openmm.HarmonicAngleForce):
                for index in range(force.getNumAngles()):
                    i, j, k, theta0, K0 = force.getAngleParameters(index)
                    if r_match(i, j, k) and (a_match(i, j, k)
                                             or a_match(k, j, i)):
                        force.setAngleParameters(index, i, j, k, angle,
                                                 K0 if K is None else K)
                        angle_list.append((i, j, k, theta0, K0))
        if angle_list and self._special_angle_force is None:
            new_force = openmm.CustomAngleForce(
                '0.5*(K0*(theta - t0)^2 - Kn*(theta - tn)^2)')
            new_force.addPerAngleParameter('t0')
            new_force.addPerAngleParameter('K0')
            new_force.addPerAngleParameter('tn')
            new_force.addPerAngleParameter('Kn')
            new_force.setForceGroup(group)
            self.addForce(new_force)
            self._special_angle_force = new_force
        for (i, j, k, theta0, K0) in angle_list:
            self._special_angle_force.addAngle(
                i, j, k, (theta0, K0, angle, K0 if K is None else K))
Exemplo n.º 11
0
def create_relative_alchemical_transformation(system, topology, positions, molecule1_indices_in_system, molecule1, molecule2,
                                              softcore_alpha=0.5, softcore_beta=12*unit.angstrom**2):
    """
    Create an OpenMM System object to handle the alchemical transformation from molecule1 to molecule2.

    system : simtk.openmm.System
       The system to be modified, already containing molecule1 whose atoms correspond to molecule1_indices_in_system.
    topology : simtk.openmm.app.Topology
       The topology object corresponding to system.
    positions : simtk.unit.Quantity of numpy array natoms x 3 compatible with units angstroms
       The positions array corresponding to system and topology.
    molecule1_indices_in_system : list of int
       Indices of molecule1 in system, with atoms in same order.
    molecule1 : openeye.oechem.OEMol
       Molecule already present in system, where the atom mapping is given by molecule1_indices_in_system.
    molecule2 : openeye.oechem.OEMol
       New molecule that molecule1 will be transformed into as lambda parameter goes from 0 -> 1.
    softcore_alpha : float, optional, default=0.5
       Softcore parameter for Lennard-Jones softening.
    softcore_beta : simtk.unit.Quantity with units compatible with angstrom**2
       Softcore parameter for Coulomb interaction softening.

    Returns
    -------
    system : simtk.openmm.System
       Modified version of system in which old system is recovered for global context paramaeter `lambda` = 0 and new molecule is substituted for `lambda` = 1.
    topology : system.openmm.Topology
       Topology corresponding to system.

    """

    # Copy molecules.
    molecule1 = oe.OEMol(molecule1)
    molecule2 = oe.OEMol(molecule2)

    # Normalize molecules.
    # TODO: May need to do more normalization here.
    oe.OEPerceiveChiral(molecule1)
    oe.OEPerceiveChiral(molecule2)

    # Make copies to not destroy original objects.
    import copy
    system = copy.deepcopy(system)
    topology = copy.deepcopy(topology)
    positions = copy.deepcopy(positions)

    # Create lists of corresponding atoms for common substructure and groups specific to molecules 1 and 2.
    atomexpr = oe.OEExprOpts_DefaultAtoms
    bondexpr = oe.OEExprOpts_BondOrder | oe.OEExprOpts_EqSingleDouble | oe.OEExprOpts_EqAromatic
    mcss = oe.OEMCSSearch(molecule1, atomexpr, bondexpr, oe.OEMCSType_Exhaustive)
    # This modifies scoring function to prefer keeping cycles complete.
    mcss.SetMCSFunc( oe.OEMCSMaxAtomsCompleteCycles() )
    # TODO: Set initial substructure size?
    # mcss.SetMinAtoms( some_number )
    # We only need one match.
    mcss.SetMaxMatches(1)
    # Determine common atoms in second molecule.
    matches = [ match for match in mcss.Match(molecule2, True) ]
    match = matches[0] # we only need the first match

    # Align common substructure of molecule2 with molecule1.
    overlay = True
    rmat  = oe.OEDoubleArray(9)
    trans = oe.OEDoubleArray(3)
    rms = oe.OERMSD(mcss.GetPattern(), molecule2, match, overlay, rmat, trans)
    if rms < 0.0:
        raise Exception("RMS overlay failure")
    oe.OERotate(molecule2, rmat)
    oe.OETranslate(molecule2, trans)

    # Make a list of the atoms in common, molecule1 only, and molecule2 only
    common1 = list() # list of atom indices in molecule1 that also appear in molecule2
    common2 = list() # list of atom indices in molecule2 that also appear in molecule1
    unique1 = list() # list of atom indices in molecule1 that DO NOT appear in molecule2
    unique2 = list() # list of atom indices in molecule2 that DO NOT appear in molecule1
    mapping1 = dict() # mapping of atoms in molecule1 to molecule2
    mapping2 = dict() # mapping of atoms in molecule2 to molecule1
    for matchpair in match.GetAtoms():
        index1 = matchpair.pattern.GetIdx()
        index2 = matchpair.target.GetIdx()
        mapping1[ index1 ] = index2
        mapping2[ index2 ] = index1
    all1 = frozenset(range(molecule1.NumAtoms()))
    all2 = frozenset(range(molecule2.NumAtoms()))
    common1 = frozenset(mapping1.keys())
    common2 = frozenset(mapping2.keys())
    unique1 = all1 - common1
    unique2 = all2 - common2

    # DEBUG
    print "list of atoms common to both molecules:"
    print "molecule1: %s" % str(common1)
    print "molecule2: %s" % str(common2)
    print "list of atoms unqiue to individual molecules:"
    print "molecule1: %s" % str(unique1)
    print "molecule2: %s" % str(unique2)
    print "MAPPING FROM MOLECULE1 TO MOLECULE2"
    for atom1 in mapping1.keys():
        atom2 = mapping1[atom1]
        print "%5d => %5d" % (atom1, atom2)

    # Create OpenMM Topology and System objects for given molecules using GAFF/AM1-BCC.
    # NOTE: This must generate the same forcefield parameters as occur in `system`.
    [system1, topology1, positions1] = generate_openmm_system(molecule1)
    [system2, topology2, positions2] = generate_openmm_system(molecule2)

    #
    # Start building combined OpenMM System object.
    #

    molecule1_atoms = [ atom for atom in molecule1.GetAtoms() ]
    molecule2_atoms = [ atom for atom in molecule2.GetAtoms() ]

    molecule2_indices_in_system = dict()

    # Build mapping of common substructure for molecule 2.
    for atom2 in common2:
        molecule2_indices_in_system[atom2] = molecule1_indices_in_system[mapping2[atom2]]

    # Find residue for molecule1.
    residue = None
    for atom in topology.atoms():
        if atom.index in molecule1_indices_in_system:
            residue = atom.residue
            break

    # Handle additional particles.
    print "Adding particles from system2..."
    for atom2 in unique2:
        atom = molecule2_atoms[atom2]
        name = atom.GetName()
        atomic_number = atom.GetAtomicNum()
        element = app.Element.getByAtomicNumber(atomic_number)
        mass = system2.getParticleMass(atom2)
        print [name, element, mass]
        index = system.addParticle(mass)
        molecule2_indices_in_system[atom2] = index

        # TODO: Add new atoms to topology object as well.
        topology.addAtom(name, element, residue)

    # Turn molecule2_indices_in_system into list
    molecule2_indices_in_system = [ molecule2_indices_in_system[atom2] for atom2 in range(molecule2.NumAtoms()) ]

    print "Atom mappings into System object"
    print "molecule1: %s" % str(molecule1_indices_in_system)
    print "molecule2: %s" % str(molecule2_indices_in_system)

    # Handle constraints.
    # TODO: What happens if constraints change? Raise Exception then.
    print "Adding constraints from system2..."
    for index in range(system2.getNumConstraints()):
        # Extract constraint distance from system2.
        [atom2_i, atom2_j, distance] = system.getConstraintParameters(index)
        # Map atoms from system2 into system.
        atom_i = molecule2_indices_in_system[atom2_i]
        atom_j = molecule2_indices_in_system[atom2_j]
        # Add constraint to system.
        system.addConstraint(atom_i, atom_j, distance)

    # Create new positions array.
    natoms = positions.shape[0] + len(unique2) # new number of atoms
    positions = unit.Quantity(np.resize(positions/positions.unit, [natoms,3]), positions.unit)
    for atom2 in unique2:
        (x, y, z) = molecule2.GetCoords(molecule2_atoms[atom2])
        index = molecule2_indices_in_system[atom2]
        positions[index,0] = x * unit.angstrom
        positions[index,1] = y * unit.angstrom
        positions[index,2] = z * unit.angstrom

    # Build a list of Force objects in system.
    forces = [ system.getForce(index) for index in range(system.getNumForces()) ]
    forces1 = { system1.getForce(index).__class__.__name__ : system1.getForce(index) for index in range(system1.getNumForces()) }
    forces2 = { system2.getForce(index).__class__.__name__ : system2.getForce(index) for index in range(system2.getNumForces()) }

    # Process forces.
    for force in forces:
        # Get force name.
        force_name = force.__class__.__name__
        force1 = forces1[force_name]
        force2 = forces2[force_name]
        print force_name
        if force_name == 'HarmonicBondForce':
            #
            # Process HarmonicBondForce
            #

            # Create index of bonds in system, system1, and system2.
            def unique(*args):
                if args[0] > args[-1]:
                    return tuple(reversed(args))
                else:
                    return tuple(args)

            def index_bonds(force):
                bonds = dict()
                for index in range(force.getNumBonds()):
                    [atom_i, atom_j, length, K] = force.getBondParameters(index)
                    key = unique(atom_i, atom_j) # unique tuple, possibly in reverse order
                    bonds[key] = index
                return bonds

            bonds  = index_bonds(force)   # index of bonds for system
            bonds1 = index_bonds(force1)  # index of bonds for system1
            bonds2 = index_bonds(force2)  # index of bonds for system2

            # Find bonds that are unique to each molecule.
            print "Finding bonds unique to each molecule..."
            unique_bonds1 = [ bonds1[atoms] for atoms in bonds1 if not set(atoms).issubset(common1) ]
            unique_bonds2 = [ bonds2[atoms] for atoms in bonds2 if not set(atoms).issubset(common2) ]

            # Build list of bonds shared among all molecules.
            print "Building a list of shared bonds..."
            shared_bonds = list()
            for atoms2 in bonds2:
                if set(atoms2).issubset(common2):
                    atoms  = tuple(molecule2_indices_in_system[atom2] for atom2 in atoms2)
                    atoms1 = tuple(mapping2[atom2] for atom2 in atoms2)
                    # Find bond index terms.
                    index  = bonds[unique(*atoms)]
                    index1 = bonds1[unique(*atoms1)]
                    index2 = bonds2[unique(*atoms2)]
                    # Store.
                    shared_bonds.append( (index, index1, index2) )

            # Add bonds that are unique to molecule2.
            print "Adding bonds unique to molecule2..."
            for index2 in unique_bonds2:
                [atom2_i, atom2_j, length2, K2] = force2.getBondParameters(index2)
                atom_i = molecule2_indices_in_system[atom2_i]
                atom_j = molecule2_indices_in_system[atom2_j]
                force.addBond(atom_i, atom_j, length2, K2)

            # Create a CustomBondForce to handle interpolated bond parameters.
            print "Creating CustomBondForce..."
            energy_expression  = '(K/2)*(r-length)^2;'
            energy_expression += 'K = (1-lambda)*K1 + lambda*K2;' # linearly interpolate spring constant
            energy_expression += 'length = (1-lambda)*length1 + lambda*length2;' # linearly interpolate bond length
            custom_force = mm.CustomBondForce(energy_expression)
            custom_force.addGlobalParameter('lambda', 0.0)
            custom_force.addPerBondParameter('length1') # molecule1 bond length
            custom_force.addPerBondParameter('K1') # molecule1 spring constant
            custom_force.addPerBondParameter('length2') # molecule2 bond length
            custom_force.addPerBondParameter('K2') # molecule2 spring constant
            system.addForce(custom_force)

            # Process bonds that are shared by molecule1 and molecule2.
            print "Translating shared bonds to CustomBondForce..."
            for (index, index1, index2) in shared_bonds:
                # Zero out standard bond force.
                [atom_i, atom_j, length, K] = force.getBondParameters(index)
                force.setBondParameters(index, atom_i, atom_j, length, K*0.0)
                # Create interpolated bond parameters.
                [atom1_i, atom1_j, length1, K1] = force1.getBondParameters(index1)
                [atom2_i, atom2_j, length2, K2] = force2.getBondParameters(index2)
                custom_force.addBond(atom_i, atom_j, [length1, K1, length2, K2])

        if force_name == 'HarmonicAngleForce':
            #
            # Process HarmonicAngleForce
            #

            # Create index of angles in system, system1, and system2.
            def unique(*args):
                if args[0] > args[-1]:
                    return tuple(reversed(args))
                else:
                    return tuple(args)

            def index_angles(force):
                angles = dict()
                for index in range(force.getNumAngles()):
                    [atom_i, atom_j, atom_k, angle, K] = force.getAngleParameters(index)
                    key = unique(atom_i, atom_j, atom_k) # unique tuple, possibly in reverse order
                    angles[key] = index
                return angles

            angles  = index_angles(force)   # index of angles for system
            angles1 = index_angles(force1)  # index of angles for system1
            angles2 = index_angles(force2)  # index of angles for system2

            # Find angles that are unique to each molecule.
            print "Finding angles unique to each molecule..."
            unique_angles1 = [ angles1[atoms] for atoms in angles1 if not set(atoms).issubset(common1) ]
            unique_angles2 = [ angles2[atoms] for atoms in angles2 if not set(atoms).issubset(common2) ]

            # Build list of angles shared among all molecules.
            print "Building a list of shared angles..."
            shared_angles = list()
            for atoms2 in angles2:
                if set(atoms2).issubset(common2):
                    atoms  = tuple(molecule2_indices_in_system[atom2] for atom2 in atoms2)
                    atoms1 = tuple(mapping2[atom2] for atom2 in atoms2)
                    # Find angle index terms.
                    index  = angles[unique(*atoms)]
                    index1 = angles1[unique(*atoms1)]
                    index2 = angles2[unique(*atoms2)]
                    # Store.
                    shared_angles.append( (index, index1, index2) )

            # Add angles that are unique to molecule2.
            print "Adding angles unique to molecule2..."
            for index2 in unique_angles2:
                [atom2_i, atom2_j, atom2_k, theta2, K2] = force2.getAngleParameters(index2)
                atom_i = molecule2_indices_in_system[atom2_i]
                atom_j = molecule2_indices_in_system[atom2_j]
                atom_k = molecule2_indices_in_system[atom2_k]
                force.addAngle(atom_i, atom_j, atom_k, theta2, K2)

            # Create a CustomAngleForce to handle interpolated angle parameters.
            print "Creating CustomAngleForce..."
            energy_expression  = '(K/2)*(theta-theta0)^2;'
            energy_expression += 'K = (1-lambda)*K_1 + lambda*K_2;' # linearly interpolate spring constant
            energy_expression += 'theta0 = (1-lambda)*theta0_1 + lambda*theta0_2;' # linearly interpolate equilibrium angle
            custom_force = mm.CustomAngleForce(energy_expression)
            custom_force.addGlobalParameter('lambda', 0.0)
            custom_force.addPerAngleParameter('theta0_1') # molecule1 equilibrium angle
            custom_force.addPerAngleParameter('K_1') # molecule1 spring constant
            custom_force.addPerAngleParameter('theta0_2') # molecule2 equilibrium angle
            custom_force.addPerAngleParameter('K_2') # molecule2 spring constant
            system.addForce(custom_force)

            # Process angles that are shared by molecule1 and molecule2.
            print "Translating shared angles to CustomAngleForce..."
            for (index, index1, index2) in shared_angles:
                # Zero out standard angle force.
                [atom_i, atom_j, atom_k, theta0, K] = force.getAngleParameters(index)
                force.setAngleParameters(index, atom_i, atom_j, atom_k, theta0, K*0.0)
                # Create interpolated angle parameters.
                [atom1_i, atom1_j, atom1_k, theta1, K1] = force1.getAngleParameters(index1)
                [atom2_i, atom2_j, atom2_k, theta2, K2] = force2.getAngleParameters(index2)
                custom_force.addAngle(atom_i, atom_j, atom_k, [theta1, K1, theta2, K2])

        if force_name == 'PeriodicTorsionForce':
            #
            # Process PeriodicTorsionForce
            # TODO: Match up periodicities and deal with multiple terms per torsion
            #

            # Create index of torsions in system, system1, and system2.
            def unique(*args):
                if args[0] > args[-1]:
                    return tuple(reversed(args))
                else:
                    return tuple(args)

            def index_torsions(force):
                torsions = dict()
                for index in range(force.getNumTorsions()):
                    [atom_i, atom_j, atom_k, atom_l, periodicity, phase, K] = force.getTorsionParameters(index)
                    key = unique(atom_i, atom_j, atom_k, atom_l) # unique tuple, possibly in reverse order
                    torsions[key] = index
                return torsions

            torsions  = index_torsions(force)   # index of torsions for system
            torsions1 = index_torsions(force1)  # index of torsions for system1
            torsions2 = index_torsions(force2)  # index of torsions for system2

            # Find torsions that are unique to each molecule.
            print "Finding torsions unique to each molecule..."
            unique_torsions1 = [ torsions1[atoms] for atoms in torsions1 if not set(atoms).issubset(common1) ]
            unique_torsions2 = [ torsions2[atoms] for atoms in torsions2 if not set(atoms).issubset(common2) ]

            # Build list of torsions shared among all molecules.
            print "Building a list of shared torsions..."
            shared_torsions = list()
            for atoms2 in torsions2:
                if set(atoms2).issubset(common2):
                    atoms  = tuple(molecule2_indices_in_system[atom2] for atom2 in atoms2)
                    atoms1 = tuple(mapping2[atom2] for atom2 in atoms2)
                    # Find torsion index terms.
                    try:
                        index  = torsions[unique(*atoms)]
                        index1 = torsions1[unique(*atoms1)]
                        index2 = torsions2[unique(*atoms2)]
                    except Exception as e:
                        print e
                        print "torsions :  %s" % str(unique(*atoms))
                        print "torsions1:  %s" % str(unique(*atoms1))
                        print "torsions2:  %s" % str(unique(*atoms2))
                        raise Exception("Error occurred in building a list of torsions common to all molecules.")

                    # Store.
                    shared_torsions.append( (index, index1, index2) )

            # Add torsions that are unique to molecule2.
            print "Adding torsions unique to molecule2..."
            for index2 in unique_torsions2:
                [atom2_i, atom2_j, atom2_k, atom2_l, periodicity2, phase2, K2] = force2.getTorsionParameters(index2)
                atom_i = molecule2_indices_in_system[atom2_i]
                atom_j = molecule2_indices_in_system[atom2_j]
                atom_k = molecule2_indices_in_system[atom2_k]
                atom_l = molecule2_indices_in_system[atom2_l]
                force.addTorsion(atom_i, atom_j, atom_k, atom_l, periodicity2, phase2, K2)

            # Create a CustomTorsionForce to handle interpolated torsion parameters.
            print "Creating CustomTorsionForce..."
            energy_expression  = '(1-lambda)*U1 + lambda*U2;'
            energy_expression += 'U1 = K1*(1+cos(periodicity1*theta-phase1));'
            energy_expression += 'U2 = K2*(1+cos(periodicity2*theta-phase2));'
            custom_force = mm.CustomTorsionForce(energy_expression)
            custom_force.addGlobalParameter('lambda', 0.0)
            custom_force.addPerTorsionParameter('periodicity1') # molecule1 periodicity
            custom_force.addPerTorsionParameter('phase1') # molecule1 phase
            custom_force.addPerTorsionParameter('K1') # molecule1 spring constant
            custom_force.addPerTorsionParameter('periodicity2') # molecule2 periodicity
            custom_force.addPerTorsionParameter('phase2') # molecule2 phase
            custom_force.addPerTorsionParameter('K2') # molecule2 spring constant
            system.addForce(custom_force)

            # Process torsions that are shared by molecule1 and molecule2.
            print "Translating shared torsions to CustomTorsionForce..."
            for (index, index1, index2) in shared_torsions:
                # Zero out standard torsion force.
                [atom_i, atom_j, atom_k, atom_l, periodicity, phase, K] = force.getTorsionParameters(index)
                force.setTorsionParameters(index, atom_i, atom_j, atom_k, atom_l, periodicity, phase, K*0.0)
                # Create interpolated torsion parameters.
                [atom1_i, atom1_j, atom1_k, atom1_l, periodicity1, phase1, K1] = force1.getTorsionParameters(index1)
                [atom2_i, atom2_j, atom2_k, atom2_l, periodicity2, phase2, K2] = force2.getTorsionParameters(index2)
                custom_force.addTorsion(atom_i, atom_j, atom_k, atom_l, [periodicity1, phase1, K1, periodicity2, phase2, K2])

        if force_name == 'NonbondedForce':
            #
            # Process NonbondedForce
            #

            # Add nonbonded entries for molecule2 to ensure total number of particle entries is correct.
            for atom in unique2:
                [charge, sigma, epsilon] = force2.getParticleParameters(atom)
                force.addParticle(charge, sigma, epsilon)

            # Zero out nonbonded entries for molecule1.
            for atom in molecule1_indices_in_system:
                [charge, sigma, epsilon] = force.getParticleParameters(atom)
                force.setParticleParameters(atom, 0*charge, sigma, 0*epsilon)
            # Zero out nonbonded entries for molecule2.
            for atom in molecule2_indices_in_system:
                [charge, sigma, epsilon] = force.getParticleParameters(atom)
                force.setParticleParameters(atom, 0*charge, sigma, 0*epsilon)

            # Create index of exceptions in system, system1, and system2.
            def unique(*args):
                if args[0] > args[-1]:
                    return tuple(reversed(args))
                else:
                    return tuple(args)

            def index_exceptions(force):
                exceptions = dict()
                for index in range(force.getNumExceptions()):
                    [atom_i, atom_j, chargeProd, sigma, epsilon] = force.getExceptionParameters(index)
                    key = unique(atom_i, atom_j) # unique tuple, possibly in reverse order
                    exceptions[key] = index
                return exceptions

            exceptions  = index_exceptions(force)   # index of exceptions for system
            exceptions1 = index_exceptions(force1)  # index of exceptions for system1
            exceptions2 = index_exceptions(force2)  # index of exceptions for system2

            # Find exceptions that are unique to each molecule.
            print "Finding exceptions unique to each molecule..."
            unique_exceptions1 = [ exceptions1[atoms] for atoms in exceptions1 if not set(atoms).issubset(common1) ]
            unique_exceptions2 = [ exceptions2[atoms] for atoms in exceptions2 if not set(atoms).issubset(common2) ]

            # Build list of exceptions shared among all molecules.
            print "Building a list of shared exceptions..."
            shared_exceptions = list()
            for atoms2 in exceptions2:
                if set(atoms2).issubset(common2):
                    atoms  = tuple(molecule2_indices_in_system[atom2] for atom2 in atoms2)
                    atoms1 = tuple(mapping2[atom2] for atom2 in atoms2)
                    # Find exception index terms.
                    index  = exceptions[unique(*atoms)]
                    index1 = exceptions1[unique(*atoms1)]
                    index2 = exceptions2[unique(*atoms2)]
                    # Store.
                    shared_exceptions.append( (index, index1, index2) )

            # Add exceptions that are unique to molecule2.
            print "Adding exceptions unique to molecule2..."
            for index2 in unique_exceptions2:
                [atom2_i, atom2_j, chargeProd, sigma, epsilon] = force2.getExceptionParameters(index2)
                atom_i = molecule2_indices_in_system[atom2_i]
                atom_j = molecule2_indices_in_system[atom2_j]
                force.addException(atom_i, atom_j, chargeProd, sigma, epsilon)

            # Create list of alchemically modified atoms in system.
            alchemical_atom_indices = list(set(molecule1_indices_in_system).union(set(molecule2_indices_in_system)))

            # Create atom groups.
            natoms = system.getNumParticles()
            atomset1 = set(alchemical_atom_indices) # only alchemically-modified atoms
            atomset2 = set(range(system.getNumParticles())) # all atoms, including alchemical region

            # CustomNonbondedForce energy expression.
            sterics_energy_expression = ""
            electrostatics_energy_expression = ""

            # Create a CustomNonbondedForce to handle alchemically interpolated nonbonded parameters.
            # Select functional form based on nonbonded method.
            method = force.getNonbondedMethod()
            if method in [mm.NonbondedForce.NoCutoff]:
                # soft-core Lennard-Jones
                sterics_energy_expression += "U_sterics = 4*epsilon*x*(x-1.0); x1 = (sigma/reff_sterics)^6;"
                # soft-core Coulomb
                electrostatics_energy_expression += "U_electrostatics = ONE_4PI_EPS0*chargeprod/reff_electrostatics;"
            elif method in [mm.NonbondedForce.CutoffPeriodic, mm.NonbondedForce.CutoffNonPeriodic]:
                # soft-core Lennard-Jones
                sterics_energy_expression += "U_sterics = 4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;"
                # reaction-field electrostatics
                epsilon_solvent = force.getReactionFieldDielectric()
                r_cutoff = force.getCutoffDistance()
                electrostatics_energy_expression += "U_electrostatics = ONE_4PI_EPS0*chargeprod*(reff_electrostatics^(-1) + k_rf*reff_electrostatics^2 - c_rf);"
                k_rf = r_cutoff**(-3) * ((epsilon_solvent - 1) / (2*epsilon_solvent + 1))
                c_rf = r_cutoff**(-1) * ((3*epsilon_solvent) / (2*epsilon_solvent + 1))
                electrostatics_energy_expression += "k_rf = %f;" % (k_rf / k_rf.in_unit_system(unit.md_unit_system).unit)
                electrostatics_energy_expression += "c_rf = %f;" % (c_rf / c_rf.in_unit_system(unit.md_unit_system).unit)
            elif method in [mm.NonbondedForce.PME, mm.NonbondedForce.Ewald]:
                # soft-core Lennard-Jones
                sterics_energy_expression += "U_sterics = 4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;"
                # Ewald direct-space electrostatics
                [alpha_ewald, nx, ny, nz] = force.getPMEParameters()
                if alpha_ewald == 0.0:
                    # If alpha is 0.0, alpha_ewald is computed by OpenMM from from the error tolerance.
                    delta = force.getEwaldErrorTolerance()
                    r_cutoff = force.getCutoffDistance()
                    alpha_ewald = np.sqrt(-np.log(2*delta)) / r_cutoff
                electrostatics_energy_expression += "U_electrostatics = ONE_4PI_EPS0*chargeprod*erfc(alpha_ewald*reff_electrostatics)/reff_electrostatics;"
                electrostatics_energy_expression += "alpha_ewald = %f;" % (alpha_ewald / alpha_ewald.in_unit_system(unit.md_unit_system).unit)
                # TODO: Handle reciprocal-space electrostatics
            else:
                raise Exception("Nonbonded method %s not supported yet." % str(method))

            # Add additional definitions common to all methods.
            sterics_energy_expression += "epsilon = (1-lambda)*epsilonA + lambda*epsilonB;" #interpolation
            sterics_energy_expression += "reff_sterics = sigma*((softcore_alpha*lambda_alpha + (r/sigma)^6))^(1/6);" # effective softcore distance for sterics
            sterics_energy_expression += "softcore_alpha = %f;" % softcore_alpha
            # TODO: We may have to ensure that softcore_degree is 1 if we are close to an alchemically-eliminated endpoint.
            sterics_energy_expression += "lambda_alpha = lambda*(1-lambda);"
            electrostatics_energy_expression += "chargeProd = (1-lambda)*chargeProdA + lambda*chargeProdB;" #interpolation
            electrostatics_energy_expression += "reff_electrostatics = sqrt(softcore_beta*lambda_beta + r^2);" # effective softcore distance for electrostatics
            electrostatics_energy_expression += "softcore_beta = %f;" % (softcore_beta / softcore_beta.in_unit_system(unit.md_unit_system).unit)
            electrostatics_energy_expression += "ONE_4PI_EPS0 = %f;" % ONE_4PI_EPS0 # already in OpenMM units
            # TODO: We may have to ensure that softcore_degree is 1 if we are close to an alchemically-eliminated endpoint.
            sterics_energy_expression += "lambda_beta = lambda*(1-lambda);"

            # Define mixing rules.
            sterics_mixing_rules = ""
            sterics_mixing_rules += "epsilonA = sqrt(epsilonA1*epsilonA2);" # mixing rule for epsilon
            sterics_mixing_rules += "epsilonB = sqrt(epsilonB1*epsilonB2);" # mixing rule for epsilon
            sterics_mixing_rules += "sigmaA = 0.5*(sigmaA1 + sigmaA2);" # mixing rule for sigma
            sterics_mixing_rules += "sigmaB = 0.5*(sigmaB1 + sigmaB2);" # mixing rule for sigma
            electrostatics_mixing_rules = ""
            electrostatics_mixing_rules += "chargeprodA = chargeA1*chargeA2;" # mixing rule for charges
            electrostatics_mixing_rules += "chargeprodB = chargeB1*chargeB2;" # mixing rule for charges

            # Create CustomNonbondedForce to handle interactions between alchemically-modified atoms and rest of system.
            electrostatics_custom_nonbonded_force = mm.CustomNonbondedForce("U_electrostatics;" + electrostatics_energy_expression + electrostatics_mixing_rules)
            electrostatics_custom_nonbonded_force.addGlobalParameter("lambda", 0.0);
            electrostatics_custom_nonbonded_force.addPerParticleParameter("chargeA") # partial charge initial
            electrostatics_custom_nonbonded_force.addPerParticleParameter("chargeB") # partial charge final
            sterics_custom_nonbonded_force = mm.CustomNonbondedForce("U_sterics;" + sterics_energy_expression + sterics_mixing_rules)
            sterics_custom_nonbonded_force.addGlobalParameter("lambda", 0.0);
            sterics_custom_nonbonded_force.addPerParticleParameter("sigmaA") # Lennard-Jones sigma initial
            sterics_custom_nonbonded_force.addPerParticleParameter("epsilonA") # Lennard-Jones epsilon initial
            sterics_custom_nonbonded_force.addPerParticleParameter("sigmaB") # Lennard-Jones sigma final
            sterics_custom_nonbonded_force.addPerParticleParameter("epsilonB") # Lennard-Jones epsilon final

            # Restrict interaction evaluation to be between alchemical atoms and rest of environment.
            # TODO: Exclude intra-alchemical region if we are separately handling that through a separate CustomNonbondedForce for decoupling.
            sterics_custom_nonbonded_force.addInteractionGroup(atomset1, atomset2)
            electrostatics_custom_nonbonded_force.addInteractionGroup(atomset1, atomset2)

            # Add exclusions between unique parts of molecule1 and molecule2 so they do not interact.
            print "Add exclusions between unique parts of molecule1 and molecule2 that should not interact..."
            for atom1_i in unique1:
                for atom2_j in unique2:
                    atom_i = molecule1_indices_in_system[atom1_i]
                    atom_j = molecule2_indices_in_system[atom2_j]
                    electrostatics_custom_nonbonded_force.addExclusion(atom_i, atom_j)
                    sterics_custom_nonbonded_force.addExclusion(atom_i, atom_j)

            # Add custom forces to system.
            system.addForce(sterics_custom_nonbonded_force)
            system.addForce(electrostatics_custom_nonbonded_force)

            # Create CustomBondForce to handle exceptions for both kinds of interactions.
            #custom_bond_force = mm.CustomBondForce("U_sterics + U_electrostatics;" + sterics_energy_expression + electrostatics_energy_expression)
            #custom_bond_force.addGlobalParameter("lambda", 0.0);
            #custom_bond_force.addPerBondParameter("chargeprodA") # charge product
            #custom_bond_force.addPerBondParameter("sigmaA") # Lennard-Jones effective sigma
            #custom_bond_force.addPerBondParameter("epsilonA") # Lennard-Jones effective epsilon
            #custom_bond_force.addPerBondParameter("chargeprodB") # charge product
            #custom_bond_force.addPerBondParameter("sigmaB") # Lennard-Jones effective sigma
            #custom_bond_force.addPerBondParameter("epsilonB") # Lennard-Jones effective epsilon
            #system.addForce(custom_bond_force)

            # Copy over all Nonbonded parameters for normal atoms to Custom*Force objects.
            for particle_index in range(force.getNumParticles()):
                # Retrieve parameters.
                [charge, sigma, epsilon] = force.getParticleParameters(particle_index)
                # Add parameters to custom force handling interactions between alchemically-modified atoms and rest of system.
                sterics_custom_nonbonded_force.addParticle([sigma, epsilon, sigma, epsilon])
                electrostatics_custom_nonbonded_force.addParticle([charge, charge])

            # Copy over parameters for common substructure.
            for atom1 in common1:
                atom2 = mapping1[atom1] # index into system2
                index = molecule1_indices_in_system[atom1] # index into system
                [charge1, sigma1, epsilon1] = force1.getParticleParameters(atom1)
                [charge2, sigma2, epsilon2] = force2.getParticleParameters(atom2)
                sterics_custom_nonbonded_force.setParticleParameters(index, [sigma1, epsilon1, sigma2, epsilon2])
                electrostatics_custom_nonbonded_force.setParticleParameters(index, [charge1, charge2])

            # Copy over parameters for molecule1 unique atoms.
            for atom1 in unique1:
                index = molecule1_indices_in_system[atom1] # index into system
                [charge1, sigma1, epsilon1] = force1.getParticleParameters(atom1)
                sterics_custom_nonbonded_force.setParticleParameters(index, [sigma1, epsilon1, sigma1, 0*epsilon1])
                electrostatics_custom_nonbonded_force.setParticleParameters(index, [charge1, 0*charge1])

            # Copy over parameters for molecule2 unique atoms.
            for atom2 in unique2:
                index = molecule2_indices_in_system[atom2] # index into system
                [charge2, sigma2, epsilon2] = force2.getParticleParameters(atom2)
                sterics_custom_nonbonded_force.setParticleParameters(index, [sigma2, 0*epsilon2, sigma2, epsilon2])
                electrostatics_custom_nonbonded_force.setParticleParameters(index, [0*charge2, charge2])

        else:
            #raise Exception("Force type %s unknown." % force_name)
            pass

    return [system, topology, positions]
Exemplo n.º 12
0
def apply_dat_restraint(system,
                        restraint,
                        phase,
                        window_number,
                        flat_bottom=False,
                        force_group=None):
    """A utility function which takes in pAPRika restraints and applies the
    restraints to an OpenMM System object.

    Parameters
    ----------
    system : :class:`openmm.System`
        The system object to add the positional restraints to.
    restraint : list
        List of pAPRika defined restraints
    phase : str
        Phase of calculation ("attach", "pull" or "release")
    window_number : int
        The corresponding window number of the current phase
    flat_bottom : bool, optional
        Specify whether the restraint is a flat bottom potential
    force_group : int, optional
        The force group to add the positional restraints to.

    """

    from simtk import openmm, unit

    assert phase in {"attach", "pull", "release"}

    if flat_bottom and phase == "attach" and restraint.mask3:
        flat_bottom_force = openmm.CustomAngleForce(
            "step(-(theta - theta_0)) * k * (theta - theta_0)^2")
        # If theta is greater than theta_0, then the argument to step is negative,
        # which means the force is off.
        flat_bottom_force.addPerAngleParameter("k")
        flat_bottom_force.addPerAngleParameter("theta_0")

        theta_0 = 91.0 * unit.degrees
        k = (restraint.phase[phase]["force_constants"][window_number] *
             unit.kilocalories_per_mole / unit.radian**2)
        flat_bottom_force.addAngle(
            restraint.index1[0],
            restraint.index2[0],
            restraint.index3[0],
            [k, theta_0],
        )
        system.addForce(flat_bottom_force)
        if force_group:
            flat_bottom_force.setForceGroup(force_group)

        return

    elif flat_bottom and phase == "attach" and not restraint.mask3:
        flat_bottom_force = openmm.CustomBondForce(
            "step((r - r_0)) * k * (r - r_0)^2")
        # If x is greater than x_0, then the argument to step is positive, which means
        # the force is on.
        flat_bottom_force.addPerBondParameter("k")
        flat_bottom_force.addPerBondParameter("r_0")

        r_0 = restraint.phase[phase]["targets"][window_number] * unit.angstrom
        k = (restraint.phase[phase]["force_constants"][window_number] *
             unit.kilocalories_per_mole / unit.radian**2)
        flat_bottom_force.addBond(
            restraint.index1[0],
            restraint.index2[0],
            [k, r_0],
        )
        system.addForce(flat_bottom_force)
        if force_group:
            flat_bottom_force.setForceGroup(force_group)

        return

    elif flat_bottom and phase == "pull":
        return
    elif flat_bottom and phase == "release":
        return

    if restraint.mask2 and not restraint.mask3:
        if not restraint.group1 and not restraint.group2:
            bond_restraint = openmm.CustomBondForce("k * (r - r_0)^2")
            bond_restraint.addPerBondParameter("k")
            bond_restraint.addPerBondParameter("r_0")

            r_0 = restraint.phase[phase]["targets"][
                window_number] * unit.angstroms
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.angstrom**2)
            bond_restraint.addBond(restraint.index1[0], restraint.index2[0],
                                   [k, r_0])
            system.addForce(bond_restraint)
        else:
            bond_restraint = openmm.CustomCentroidBondForce(
                2, "k * (distance(g1, g2) - r_0)^2")
            bond_restraint.addPerBondParameter("k")
            bond_restraint.addPerBondParameter("r_0")
            r_0 = restraint.phase[phase]["targets"][
                window_number] * unit.angstroms
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.angstrom**2)
            g1 = bond_restraint.addGroup(restraint.index1)
            g2 = bond_restraint.addGroup(restraint.index2)
            bond_restraint.addBond([g1, g2], [k, r_0])
            system.addForce(bond_restraint)

        if force_group:
            bond_restraint.setForceGroup(force_group)

    elif restraint.mask3 and not restraint.mask4:
        if not restraint.group1 and not restraint.group2 and not restraint.group3:
            angle_restraint = openmm.CustomAngleForce(
                "k * (theta - theta_0)^2")
            angle_restraint.addPerAngleParameter("k")
            angle_restraint.addPerAngleParameter("theta_0")

            theta_0 = restraint.phase[phase]["targets"][
                window_number] * unit.degrees
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.radian**2)
            angle_restraint.addAngle(
                restraint.index1[0],
                restraint.index2[0],
                restraint.index3[0],
                [k, theta_0],
            )
            system.addForce(angle_restraint)
        else:
            # Probably needs openmm.CustomCentroidAngleForce (?)
            raise NotImplementedError
        if force_group:
            angle_restraint.setForceGroup(force_group)

    elif restraint.mask4:
        if (not restraint.group1 and not restraint.group2
                and not restraint.group3 and not restraint.group4):
            dihedral_restraint = openmm.CustomTorsionForce(
                f"k * min(min(abs(theta - theta_0), abs(theta - theta_0 + 2 * "
                f"{_PI_})), abs(theta - theta_0 - 2 * {_PI_}))^2")
            dihedral_restraint.addPerTorsionParameter("k")
            dihedral_restraint.addPerTorsionParameter("theta_0")

            theta_0 = restraint.phase[phase]["targets"][
                window_number] * unit.degrees
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.radian**2)
            dihedral_restraint.addTorsion(
                restraint.index1[0],
                restraint.index2[0],
                restraint.index3[0],
                restraint.index4[0],
                [k, theta_0],
            )
            system.addForce(dihedral_restraint)
        else:
            # Probably needs openmm.CustomCentroidTorsionForce (?)
            raise NotImplementedError
        if force_group:
            dihedral_restraint.setForceGroup(force_group)
Exemplo n.º 13
0
def apply_openmm_restraints(system,
                            restraint,
                            window,
                            flat_bottom=False,
                            ForceGroup=None):
    if window[0] == "a":
        phase = "attach"
    elif window[0] == "p":
        phase = "pull"
    elif window[0] == "r":
        phase = "release"
    window_number = int(window[1:])

    if flat_bottom and phase == "attach" and restraint.mask3:
        flat_bottom_force = openmm.CustomAngleForce(
            'step(-(theta - theta_0)) * k * (theta - theta_0)^2')
        # If theta is greater than theta_0, then the argument to step is negative, which means the force is off.
        flat_bottom_force.addPerAngleParameter("k")
        flat_bottom_force.addPerAngleParameter("theta_0")

        theta_0 = 91.0 * unit.degrees
        k = (restraint.phase[phase]["force_constants"][window_number] *
             unit.kilocalories_per_mole / unit.radian**2)
        flat_bottom_force.addAngle(
            restraint.index1[0],
            restraint.index2[0],
            restraint.index3[0],
            [k, theta_0],
        )
        system.addForce(flat_bottom_force)
        if ForceGroup:
            flat_bottom_force.setForceGroup(ForceGroup)

        return system
    elif flat_bottom and phase == "attach" and not restraint.mask3:
        flat_bottom_force = openmm.CustomBondForce(
            'step((r - r_0)) * k * (r - r_0)^2')
        # If x is greater than x_0, then the argument to step is positive, which means the force is on.
        flat_bottom_force.addPerBondParameter("k")
        flat_bottom_force.addPerBondParameter("r_0")

        r_0 = restraint.phase[phase]["targets"][window_number] * unit.angstrom
        k = (restraint.phase[phase]["force_constants"][window_number] *
             unit.kilocalories_per_mole / unit.radian**2)
        flat_bottom_force.addBond(
            restraint.index1[0],
            restraint.index2[0],
            [k, r_0],
        )
        system.addForce(flat_bottom_force)
        if ForceGroup:
            flat_bottom_force.setForceGroup(ForceGroup)

        return system

    elif flat_bottom and phase == "pull":
        return system
    elif flat_bottom and phase == "release":
        return system

    if restraint.mask2 and not restraint.mask3:
        if not restraint.group1 and not restraint.group2:
            bond_restraint = openmm.CustomBondForce("k * (r - r_0)^2")
            bond_restraint.addPerBondParameter("k")
            bond_restraint.addPerBondParameter("r_0")

            r_0 = restraint.phase[phase]["targets"][
                window_number] * unit.angstroms
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.angstrom**2)
            bond_restraint.addBond(restraint.index1[0], restraint.index2[0],
                                   [k, r_0])
            system.addForce(bond_restraint)
        else:
            bond_restraint = openmm.CustomCentroidBondForce(
                2, "k * (distance(g1, g2) - r_0)^2")
            bond_restraint.addPerBondParameter("k")
            bond_restraint.addPerBondParameter("r_0")
            r_0 = restraint.phase[phase]["targets"][
                window_number] * unit.angstroms
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.angstrom**2)
            g1 = bond_restraint.addGroup(restraint.index1)
            g2 = bond_restraint.addGroup(restraint.index2)
            bond_restraint.addBond([g1, g2], [k, r_0])
            system.addForce(bond_restraint)

        if ForceGroup:
            bond_restraint.setForceGroup(ForceGroup)

    elif restraint.mask3 and not restraint.mask4:
        if not restraint.group1 and not restraint.group2 and not restraint.group3:
            angle_restraint = openmm.CustomAngleForce(
                "k * (theta - theta_0)^2")
            angle_restraint.addPerAngleParameter("k")
            angle_restraint.addPerAngleParameter("theta_0")

            theta_0 = restraint.phase[phase]["targets"][
                window_number] * unit.degrees
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.radian**2)
            angle_restraint.addAngle(
                restraint.index1[0],
                restraint.index2[0],
                restraint.index3[0],
                [k, theta_0],
            )
            system.addForce(angle_restraint)
        else:
            # Probably needs openmm.CustomCentroidAngleForce (?)
            raise NotImplementedError
        if ForceGroup:
            angle_restraint.setForceGroup(ForceGroup)

    elif restraint.mask4:
        if (not restraint.group1 and not restraint.group2
                and not restraint.group3 and not restraint.group4):
            dihedral_restraint = openmm.CustomTorsionForce(
                f"k * min(min(abs(theta - theta_0), abs(theta - theta_0 + 2 * {_PI_})), abs(theta - theta_0 - 2 * {_PI_}))^2"
            )
            dihedral_restraint.addPerTorsionParameter("k")
            dihedral_restraint.addPerTorsionParameter("theta_0")

            theta_0 = restraint.phase[phase]["targets"][
                window_number] * unit.degrees
            k = (restraint.phase[phase]["force_constants"][window_number] *
                 unit.kilocalories_per_mole / unit.radian**2)
            dihedral_restraint.addTorsion(
                restraint.index1[0],
                restraint.index2[0],
                restraint.index3[0],
                restraint.index4[0],
                [k, theta_0],
            )
            system.addForce(dihedral_restraint)
        else:
            # Probably needs openmm.CustomCentroidTorsionForce (?)
            raise NotImplementedError
        if ForceGroup:
            dihedral_restraint.setForceGroup(ForceGroup)
    return system
Exemplo n.º 14
0
def setup_openmm_restraints(system, restraint, phase, window):
    """
    Add particle restraints with OpenMM.
    """

    # http://docs.openmm.org/7.1.0/api-c++/generated/OpenMM.CustomExternalForce.html
    # It's possible we might need to use `periodicdistance`.

    if (restraint.mask1 is not None and restraint.mask2 is not None
            and restraint.mask3 is None and restraint.mask4 is None):

        if restraint.group1 is False and restraint.group2 is False:
            bond_restraint = mm.CustomBondForce("k * (r - r_0)^2")
            bond_restraint.addPerBondParameter("k")
            bond_restraint.addPerBondParameter("r_0")

            r_0 = restraint.phase[phase]["targets"][window] * unit.angstrom
            k = (restraint.phase[phase]["force_constants"][window] *
                 unit.kilocalorie_per_mole / unit.angstrom**2)
            bond_restraint.addBond(restraint.index1[0], restraint.index2[0],
                                   [k, r_0])
            bond_restraint.setForceGroup(1)
            system.addForce(bond_restraint)
            log.debug(
                "Added bond restraint between {} and {} with target value = "
                "{} and force constant = {}".format(restraint.mask1,
                                                    restraint.mask2, r_0, k))
        elif restraint.group1 is True or restraint.group2 is True:
            # http://docs.openmm.org/7.0.0/api-python/generated/simtk.openmm.openmm.CustomManyParticleForce.html
            # http://getyank.org/development/_modules/yank/restraints.html
            bond_restraint = mm.CustomCentroidBondForce(
                2, "k * (distance(g1, g2) - r_0)^2")
            bond_restraint.addPerBondParameter("k")
            bond_restraint.addPerBondParameter("r_0")

            r_0 = restraint.phase[phase]["targets"][window] * unit.angstrom
            k = (restraint.phase[phase]["force_constants"][window] *
                 unit.kilocalorie_per_mole / unit.angstrom**2)
            g1 = bond_restraint.addGroup(restraint.index1)
            g2 = bond_restraint.addGroup(restraint.index2)
            bond_restraint.addBond([g1, g2], [k, r_0])
            bond_restraint.setForceGroup(1)
            system.addForce(bond_restraint)
            log.debug(
                "Added bond restraint between {} and {} with target value = "
                "{} and force constant = {}".format(restraint.mask1,
                                                    restraint.mask2, r_0, k))
    else:
        log.error("Unable to add bond restraint...")
        log.debug("restraint.index1 = {}".format(restraint.index1))
        log.debug("restraint.index2 = {}".format(restraint.index2))
        raise Exception("Unable to add bond restraint...")

    if (restraint.mask1 is not None and restraint.mask2 is not None
            and restraint.mask3 is not None and restraint.mask4 is None):
        if (restraint.group1 is not False and restraint.group2 is not False
                and restraint.group3 is not False):
            log.error("Unable to add a group angle restraint...")
            log.debug("restraint.index1 = {}".format(restraint.index1))
            log.debug("restraint.index2 = {}".format(restraint.index2))
            log.debug("restraint.index3 = {}".format(restraint.index3))
            raise Exception("Unable to add a group angle restraint...")

        angle_restraint = mm.CustomAngleForce("k * (theta - theta_0)^2")
        angle_restraint.addPerAngleParameter("k")
        angle_restraint.addPerAngleParameter("theta_0")

        log.debug("Setting an angle restraint in degrees using a "
                  "force constant in kcal per mol rad**2...")
        theta_0 = restraint.phase[phase]["targets"][window] * unit.degrees
        k = (restraint.phase[phase]["force_constants"][window] *
             unit.kilocalorie_per_mole / unit.radian**2)
        angle_restraint.addAngle(restraint.index1[0], restraint.index2[0],
                                 restraint.index3[0], [k, theta_0])
        system.addForce(angle_restraint)

    if (restraint.mask1 is not None and restraint.mask2 is not None
            and restraint.mask3 is not None and restraint.mask4 is not None):
        if (restraint.group1 is not False and restraint.group2 is not False
                and restraint.group3 is not False
                and restraint.group4 is not False):
            log.error("Unable to add a group dihedral restraint...")
            log.debug("restraint.index1 = {}".format(restraint.index1))
            log.debug("restraint.index2 = {}".format(restraint.index2))
            log.debug("restraint.index3 = {}".format(restraint.index3))
            log.debug("restraint.index4 = {}".format(restraint.index4))
            raise Exception("Unable to add a group dihedral restraint...")

        dihedral_restraint = mm.CustomTorsionForce("k * (theta - theta_0)^2")
        dihedral_restraint.addPerTorsionParameter("k")
        dihedral_restraint.addPerTorsionParameter("theta_0")

        log.debug("Setting a torsion restraint in degrees using a "
                  "force constant in kcal per mol rad**2...")
        theta_0 = restraint.phase[phase]["targets"][window] * unit.degrees
        k = (restraint.phase[phase]["force_constants"][window] *
             unit.kilocalorie_per_mole / unit.radian**2)
        dihedral_restraint.addTorsion(
            restraint.index1[0],
            restraint.index2[0],
            restraint.index3[0],
            restraint.index4[0],
            [k, theta_0],
        )
        system.addForce(dihedral_restraint)

    return system