<commit_before>#!/usr/bin/env python
"""
Testbed for creating a merged topology representation with continuous lambda parameters for alchemical transformations.

Notes
-----

In this scheme, the substructure that shares common atom types is assigned to the shared "core" that is not perturbed.

Context parameters
------------------
* `alchemical_variant` - index (0, 1, 2...) of variant currently selected to be active
* `alchemical_lambda` - alchemical parameter that interpolates between 0 (core only) and 1 (variant `alchemical_varid` is fully chemically present)

"""

import gaff2xml.openeye
import openeye.oechem as oe
import simtk.openmm as mm
from simtk import unit
import simtk.openmm.app as app
import numpy as np
import copy

def create_molecule(iupac_name):
    molecule = gaff2xml.openeye.iupac_to_oemol(iupac_name)
    molecule = gaff2xml.openeye.get_charges(molecule, max_confs=1)
    #import openeye.oeomega as om
    #omega = om.OEOmega()
    #omega.SetMaxConfs(1)
    #omega(molecule)
    return molecule

def generate_openmm_system(molecule):
    """
    Generate OpenMM system, topology, and positions for a given molecule.

    Parameters
    ----------
    molecule : openeye.oechem.OEMol
       The molecule for which parameters are to be generated.

    Returns
    -------
    system : simtk.openmm.System
       The OpenMM system for the molecule.
    topology : simtk.openmm.app.Topology
       The topology object corresponding to the molecule.
    positions : simtk.unit.Quantity of dimension natoms x 3 with units compatible with length
       The atomic positions of the molecule.

    TODO
    ----
    * Add option to not generate charges.

    """

    trajs, ffxmls = gaff2xml.openeye.oemols_to_ffxml([molecule])
    ff = app.ForceField(ffxmls)
    # Get OpenMM Topology.
    topology = trajs[0].top.to_openmm()
    # Create OpenMM System object.
    system = ff.createSystem(topology)
    # Create positions.
    natoms = molecule.NumAtoms()
    positions = unit.Quantity(np.zeros([natoms,3], np.float32), unit.angstrom)
    for atom in molecule.GetAtoms():
        (x, y, z) = molecule.GetCoords(atom)
        index = atom.GetIdx()
        positions[index,0] = x * unit.angstrom
        positions[index,1] = y * unit.angstrom
        positions[index,2] = z * unit.angstrom
    # Return results.
    return [system, topology, positions]

ONE_4PI_EPS0 = 138.935456 # OpenMM constant for Coulomb interactions (openmm/platforms/reference/include/SimTKOpenMMRealType.h) in OpenMM units

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`.

    """

    # 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)<commit_msg>Starting to insert code to handle CustomNonbondedForces.<commit_after>#!/usr/bin/env python
"""
Testbed for creating a merged topology representation with continuous lambda parameters for alchemical transformations.

Notes
-----

In this scheme, the substructure that shares common atom types is assigned to the shared "core" that is not perturbed.

Context parameters
------------------
* `alchemical_variant` - index (0, 1, 2...) of variant currently selected to be active
* `alchemical_lambda` - alchemical parameter that interpolates between 0 (core only) and 1 (variant `alchemical_varid` is fully chemically present)

"""

import gaff2xml.openeye
import openeye.oechem as oe
import simtk.openmm as mm
from simtk import unit
import simtk.openmm.app as app
import numpy as np
import copy

def create_molecule(iupac_name):
    molecule = gaff2xml.openeye.iupac_to_oemol(iupac_name)
    molecule = gaff2xml.openeye.get_charges(molecule, max_confs=1)
    #import openeye.oeomega as om
    #omega = om.OEOmega()
    #omega.SetMaxConfs(1)
    #omega(molecule)
    return molecule

def generate_openmm_system(molecule):
    """
    Generate OpenMM system, topology, and positions for a given molecule.

    Parameters
    ----------
    molecule : openeye.oechem.OEMol
       The molecule for which parameters are to be generated.

    Returns
    -------
    system : simtk.openmm.System
       The OpenMM system for the molecule.
    topology : simtk.openmm.app.Topology
       The topology object corresponding to the molecule.
    positions : simtk.unit.Quantity of dimension natoms x 3 with units compatible with length
       The atomic positions of the molecule.

    TODO
    ----
    * Add option to not generate charges.

    """

    trajs, ffxmls = gaff2xml.openeye.oemols_to_ffxml([molecule])
    ff = app.ForceField(ffxmls)
    # Get OpenMM Topology.
    topology = trajs[0].top.to_openmm()
    # Create OpenMM System object.
    system = ff.createSystem(topology)
    # Create positions.
    natoms = molecule.NumAtoms()
    positions = unit.Quantity(np.zeros([natoms,3], np.float32), unit.angstrom)
    for atom in molecule.GetAtoms():
        (x, y, z) = molecule.GetCoords(atom)
        index = atom.GetIdx()
        positions[index,0] = x * unit.angstrom
        positions[index,1] = y * unit.angstrom
        positions[index,2] = z * unit.angstrom
    # Return results.
    return [system, topology, positions]

ONE_4PI_EPS0 = 138.935456 # OpenMM constant for Coulomb interactions (openmm/platforms/reference/include/SimTKOpenMMRealType.h) in OpenMM units

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 _createCustomyNonbondedForce(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)
{'fraction_non_alphanumeric': 0.08378286436312388, 'fraction_numerical': 0.007970063663313409, 'mean_word_length': 3.5286091549295775, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 14}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3362760', 'n_tokens_mistral': 6281, 'n_tokens_neox': 5865, 'n_words': 1423}