conda env create --name interchange-examples --file devtools/conda-envs/examples_env.yaml
conda activate interchange-examples
pip install -e .
cd examples/parameter_splitting
jupyter notebook parameter_splitting.ipynb
This example explains splitting a single parameter applied to two different parts of a topology into two separate parameters.
import openmm
from openff.toolkit import ForceField, Molecule
from openff.interchange import Interchange
from openff.interchange.components.potentials import Potential
from openff.interchange.models import PotentialKey, TopologyKey
We need a molecule that has two parameters that are the same in Sage, but could reasonably be different. We'll pick the two Carbon-Carbon bonds in propanol. Let's set up our Interchange, store a copy of the original as an OpenMM system, and make sure we know what propanol looks like:
propanol = Molecule.from_smiles("CCCO")
propanol.generate_conformers(n_conformers=1)
sage = ForceField("openff-2.0.0.offxml")
interchange = Interchange.from_smirnoff(sage, topology=[propanol])
original_openmm_system = interchange.to_openmm()
propanol.visualize(backend="rdkit")
Interchange stores parameters in PotentialHandler objects. A potential handler maps atom indices to parameters via a PotentialKey, which identifies the parameters in the original force field (see Tweaking and Inspecting Parameters). We can see the atom indices of all the bonds and the corresponding potential keys by inspecting the key_map attribute of the bonds potential handler:
interchange.handlers["Bonds"].key_map
By iterating over the topology, we can pick out the C-C bonds. We'll need the potential key they both use, so let's give it a name:
for bond in interchange.topology.bonds:
if all(atom.atomic_number == 6 for atom in bond.atoms):
atom_indices = tuple(
interchange.topology.atom_index(atom) for atom in bond.atoms
)
top_key = TopologyKey(atom_indices=atom_indices)
pot_key = interchange.handlers["Bonds"].key_map[top_key]
print(atom_indices, pot_key.__repr__())
Notice that the PotentialKey associated with each of the C-C bonds - atom indices (0, 1) and (1, 2) - is the same, in this case associated with SMIRKS pattern '[#6X4:1]-[#6X4:2]'. This means the same parameters have been applied to each. For the sake of this example, let's consider splitting these parameters into two types without re-running SMIRKS/SMARTS-based atom-typing. Let's increase the force constant of the C-C bond nearest the O atom by 5% (atom indices (1, 2)). This is scientifically unmotivated; randomly changing a single force constant will not (usually) improve a force field.
We'll start by cloning the existing C-C bond PotentialKey. The new ID can be anything as long as its unique, so let's choose something that makes the parameter's heritage clear without confusing ourselves with something that looks like a SMIRKS code:
pot_key_mod = PotentialKey(**pot_key.dict())
pot_key_mod.id = "[#6X4:1]-[#6X4:2]_MODIFIED"
(pot_key, pot_key_mod)
Looks good! Now we need to do the same thing with the parameters themselves. We can get the potential out by indexing with the original PontentialKey, copy it, and then adjust it's force constant $k$:
pot = interchange.handlers["Bonds"].potentials[pot_key]
pot_mod = Potential(**pot.dict())
pot_mod.parameters["k"] *= 1.05
(pot, pot_mod)
Perfect. Now we add the new potential to the handler. This won't apply it anywhere in the topology, but it'll give us something to apply. The .potentials attribute is just a regular Python dict mapping potential keys to potentials, so we can use the regular Python dict.update() method with our modified key and potential:
interchange.handlers["Bonds"].potentials.update({pot_key_mod: pot_mod})
The last step is to apply it to the topology somewhere. Interchange identifies places in the topology by tuples of atom indices. We already decided that we want to apply our new potential to the bond between atoms 1 and 2, so we define a TopologyKey to that effect and check that the bond already exists (and is what we expect):
top_key = TopologyKey(atom_indices=(1, 2))
assert top_key in interchange["Bonds"].key_map
interchange["Bonds"].key_map[top_key]
The key_map attribute is another regular dict that maps a TopologyKey to a PotentialKey. Connecting the topology key (atom indices) to a potential key goes exactly how you'd expect:
interchange.handlers["Bonds"].key_map[top_key] = pot_key_mod
To prove that we've done what we expected, let's export to OpenMM and use it's machinery to compare the force constants.
# Export to OpenMM
modified_openmm_system = interchange.to_openmm()
# Define the atom indices we care about
i = 1
j = 2
# Get the original exported k value
for force in original_openmm_system.getForces():
if type(force) == openmm.HarmonicBondForce:
for bond_idx in range(force.getNumBonds()):
if force.getBondParameters(bond_idx)[:2] == [i, j]:
original_k = force.getBondParameters(bond_idx)[3]
print(
f"K in the original system between atoms {i} and {j} is", original_k
)
# Get the exported k value in the modified system
for force in modified_openmm_system.getForces():
if type(force) == openmm.HarmonicBondForce:
for bond_idx in range(force.getNumBonds()):
if force.getBondParameters(bond_idx)[:2] == [i, j]:
modified_k = force.getBondParameters(bond_idx)[3]
print(
f"K in the modified system between atoms {i} and {j} is", modified_k
)
# Check that the modified k is 5% more than the original k
assert abs(modified_k / original_k - 1.05) < 1e-12
print(f"{modified_k}/{original_k} = {modified_k/original_k}")