nonbonded

Convert non-bonded potentials to OpenMM forces.

Modules:

• smee –

  Differentiably evaluate energies of molecules using SMIRNOFF force fields

Functions:

• convert_custom_vdw_potential –

  Converts an arbitrary vdW potential to OpenMM forces.

• convert_lj_potential –

  Convert a Lennard-Jones potential to an OpenMM force.

• convert_dexp_potential –

  Convert a DEXP potential to OpenMM forces.

• convert_coulomb_potential –

  Convert a Coulomb potential to an OpenMM force.

convert_custom_vdw_potential

convert_custom_vdw_potential(
    potential: TensorPotential,
    system: TensorSystem,
    energy_fn: str,
    mixing_fn: dict[str, str],
) -> tuple[CustomNonbondedForce, CustomBondForce]

Converts an arbitrary vdW potential to OpenMM forces.

The intermolecular interactions are described by a custom nonbonded force, while the intramolecular interactions are described by a custom bond force.

If the potential has custom mixing rules (i.e. exceptions), a lookup table will be used to store the parameters. Otherwise, the mixing rules will be applied directly in the energy function.

Parameters:

• potential (TensorPotential) –

  The potential to convert.

• system (TensorSystem) –

  The system the potential belongs to.

• energy_fn (str) –

  The energy function of the potential, written in OpenMM's custom energy function syntax.

• mixing_fn (dict[str, str]) –

  A dictionary of mixing rules for each parameter of the potential. The keys are the parameter names, and the values are the mixing rules.

  The mixing rules should be a single expression that can be evaluated using OpenMM's energy function syntax, and should not contain any assignments.

Examples:

For a Lennard-Jones potential using Lorentz-Berthelot mixing rules:

>>> energy_fn = "4*epsilon*x6*(x6 - 1.0);x6=x4*x2;x4=x2*x2;x2=x*x;x=sigma/r;"
>>> mixing_fn = {
...     "epsilon": "sqrt(epsilon1 * epsilon2)",
...     "sigma": "0.5 * (sigma1 + sigma2)",
... }

Source code in smee/converters/openmm/nonbonded.py

def convert_custom_vdw_potential(
    potential: smee.TensorPotential,
    system: smee.TensorSystem,
    energy_fn: str,
    mixing_fn: dict[str, str],
) -> tuple[openmm.CustomNonbondedForce, openmm.CustomBondForce]:
    """Converts an arbitrary vdW potential to OpenMM forces.

    The intermolecular interactions are described by a custom nonbonded force, while the
    intramolecular interactions are described by a custom bond force.

    If the potential has custom mixing rules (i.e. exceptions), a lookup table will be
    used to store the parameters. Otherwise, the mixing rules will be applied directly
    in the energy function.

    Args:
        potential: The potential to convert.
        system: The system the potential belongs to.
        energy_fn: The energy function of the potential, written in OpenMM's custom
            energy function syntax.
        mixing_fn: A dictionary of mixing rules for each parameter of the potential.
            The keys are the parameter names, and the values are the mixing rules.

            The mixing rules should be a single expression that can be evaluated using
            OpenMM's energy function syntax, and should not contain any assignments.

    Examples:
        For a Lennard-Jones potential using Lorentz-Berthelot mixing rules:

        >>> energy_fn = "4*epsilon*x6*(x6 - 1.0);x6=x4*x2;x4=x2*x2;x2=x*x;x=sigma/r;"
        >>> mixing_fn = {
        ...     "epsilon": "sqrt(epsilon1 * epsilon2)",
        ...     "sigma": "0.5 * (sigma1 + sigma2)",
        ... }
    """
    energy_fn = re.sub(r"\s+", "", energy_fn)
    mixing_fn = {k: re.sub(r"\s+", "", v) for k, v in mixing_fn.items()}

    used_parameters, used_attributes = _detect_parameters(
        potential, energy_fn, mixing_fn
    )
    requires_lookup = potential.exceptions is not None

    inter_force = _create_nonbonded_force(
        potential, system, openmm.CustomNonbondedForce
    )
    inter_force.setEnergyFunction(energy_fn)

    intra_force = openmm.CustomBondForce(
        _prepend_scale_to_energy_fn(energy_fn, _INTRA_SCALE_VAR)
    )
    intra_force.addPerBondParameter(_INTRA_SCALE_VAR)
    intra_force.setUsesPeriodicBoundaryConditions(system.is_periodic)

    for force in [inter_force, intra_force]:
        for attr in used_attributes:
            attr_unit = potential.attribute_units[potential.attribute_cols.index(attr)]
            attr_conv = (
                (1.0 * attr_unit)
                .to_openmm()
                .value_in_unit_system(openmm.unit.md_unit_system)
            )
            attr_idx = potential.attribute_cols.index(attr)
            attr_val = float(potential.attributes[attr_idx]) * attr_conv

            force.addGlobalParameter(attr, attr_val)

    if requires_lookup:
        _add_parameters_to_vdw_with_lookup(
            potential, system, energy_fn, mixing_fn, inter_force, intra_force
        )
    else:
        _add_parameters_to_vdw_without_lookup(
            potential,
            system,
            energy_fn,
            mixing_fn,
            inter_force,
            intra_force,
            used_parameters,
        )

    return inter_force, intra_force

convert_lj_potential

convert_lj_potential(
    potential: TensorPotential, system: TensorSystem
) -> (
    NonbondedForce
    | list[CustomNonbondedForce | CustomBondForce]
)

Convert a Lennard-Jones potential to an OpenMM force.

If the potential has custom mixing rules (i.e. exceptions), the interactions will be split into an inter- and intra-molecular force.

Source code in smee/converters/openmm/nonbonded.py

@smee.converters.openmm.potential_converter(
    smee.PotentialType.VDW, smee.EnergyFn.VDW_LJ
)
def convert_lj_potential(
    potential: smee.TensorPotential, system: smee.TensorSystem
) -> openmm.NonbondedForce | list[openmm.CustomNonbondedForce | openmm.CustomBondForce]:
    """Convert a Lennard-Jones potential to an OpenMM force.

    If the potential has custom mixing rules (i.e. exceptions), the interactions will
    be split into an inter- and intra-molecular force.
    """
    energy_fn = "4*epsilon*x6*(x6 - 1.0);x6=x4*x2;x4=x2*x2;x2=x*x;x=sigma/r;"
    mixing_fn = {
        "epsilon": "sqrt(epsilon1 * epsilon2)",
        "sigma": "0.5 * (sigma1 + sigma2)",
    }

    if potential.exceptions is not None:
        return list(
            convert_custom_vdw_potential(potential, system, energy_fn, mixing_fn)
        )

    force = _create_nonbonded_force(potential, system)

    idx_offset = 0

    for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
        parameter_map = topology.parameters[potential.type]
        parameters = parameter_map.assignment_matrix @ potential.parameters.detach()

        for _ in range(n_copies):
            for epsilon, sigma in parameters:
                force.addParticle(0.0, sigma * _ANGSTROM, epsilon * _KCAL_PER_MOL)

            for index, (i, j) in enumerate(parameter_map.exclusions):
                scale = potential.attributes[parameter_map.exclusion_scale_idxs[index]]

                eps_i, sig_i = parameters[i, :]
                eps_j, sig_j = parameters[j, :]

                eps, sig = smee.potentials.nonbonded.lorentz_berthelot(
                    eps_i, eps_j, sig_i, sig_j
                )

                force.addException(
                    i + idx_offset,
                    j + idx_offset,
                    0.0,
                    float(sig) * _ANGSTROM,
                    float(eps * scale) * _KCAL_PER_MOL,
                )

            idx_offset += topology.n_particles

    return force

convert_dexp_potential

convert_dexp_potential(
    potential: TensorPotential, system: TensorSystem
) -> tuple[CustomNonbondedForce, CustomBondForce]

Convert a DEXP potential to OpenMM forces.

The intermolcular interactions are described by a custom nonbonded force, while the intramolecular interactions are described by a custom bond force.

If the potential has custom mixing rules (i.e. exceptions), a lookup table will be used to store the parameters. Otherwise, the mixing rules will be applied directly in the energy function.

Source code in smee/converters/openmm/nonbonded.py

@smee.converters.openmm.potential_converter(
    smee.PotentialType.VDW, smee.EnergyFn.VDW_DEXP
)
def convert_dexp_potential(
    potential: smee.TensorPotential, system: smee.TensorSystem
) -> tuple[openmm.CustomNonbondedForce, openmm.CustomBondForce]:
    """Convert a DEXP potential to OpenMM forces.

    The intermolcular interactions are described by a custom nonbonded force, while the
    intramolecular interactions are described by a custom bond force.

    If the potential has custom mixing rules (i.e. exceptions), a lookup table will be
    used to store the parameters. Otherwise, the mixing rules will be applied directly
    in the energy function.
    """
    energy_fn = (
        "epsilon * (repulsion - attraction);"
        "repulsion  = beta  / (alpha - beta) * exp(alpha * (1 - x));"
        "attraction = alpha / (alpha - beta) * exp(beta  * (1 - x));"
        "x = r / r_min;"
    )
    mixing_fn = {
        "epsilon": "sqrt(epsilon1 * epsilon2)",
        "r_min": "0.5 * (r_min1 + r_min2)",
    }

    return convert_custom_vdw_potential(potential, system, energy_fn, mixing_fn)

convert_coulomb_potential

convert_coulomb_potential(
    potential: TensorPotential, system: TensorSystem
) -> NonbondedForce

Convert a Coulomb potential to an OpenMM force.

Source code in smee/converters/openmm/nonbonded.py

@smee.converters.openmm.potential_converter(
    smee.PotentialType.ELECTROSTATICS, smee.EnergyFn.COULOMB
)
def convert_coulomb_potential(
    potential: smee.TensorPotential, system: smee.TensorSystem
) -> openmm.NonbondedForce:
    """Convert a Coulomb potential to an OpenMM force."""
    force = _create_nonbonded_force(potential, system)

    idx_offset = 0

    for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
        parameter_map = topology.parameters[potential.type]
        parameters = parameter_map.assignment_matrix @ potential.parameters.detach()

        for _ in range(n_copies):
            for charge in parameters:
                force.addParticle(
                    charge.detach() * openmm.unit.elementary_charge,
                    1.0 * _ANGSTROM,
                    0.0 * _KCAL_PER_MOL,
                )

            for index, (i, j) in enumerate(parameter_map.exclusions):
                q_i, q_j = parameters[i], parameters[j]
                q = q_i * q_j

                scale = potential.attributes[parameter_map.exclusion_scale_idxs[index]]

                force.addException(
                    i + idx_offset,
                    j + idx_offset,
                    scale * q,
                    1.0,
                    0.0,
                )

            idx_offset += topology.n_particles

    return force