BondCharge virtual sites are not overriding correctly

Describe the bug SMIRNOFF hierarchy rules are somewhat complicated for virtualsites. A virtualsite should only override another virtualsite if they match the same atoms[1], and have the same name and type.

To Reproduce

This code makes two vsites with the same name and type, which should apply to the same atoms, but in a different order:

from openff.toolkit.typing.engines.smirnoff import ForceField
from openff.toolkit.topology import Molecule
from openff.units import unit
from openmm import unit as omm_unit

mol = Molecule.from_mapped_smiles('[Cl:1][Na:2]')
ff = ForceField()
vsh = ff.get_parameter_handler('VirtualSites', {"version":0.3})

# Define two vsites that apply in opposite directions to the same atoms of the test molecule. 
# The vsites may collide (since they have the same name), but they should be 
# distinguishable in the output based on their `distance` values
vsh.add_parameter({'type':'BondCharge', 
                      'smirks':'[Cl:1][Na:2]',
                      'match':'once',
                      'charge_increment':[0.100, -0.100] * unit.elementary_charge,
                      'distance': [0.1010] * unit.angstrom,
                      'name':'aaa'})
vsh.add_parameter({'type':'BondCharge', 
                      'smirks':'[Na:1][Cl:2]',
                      'match':'once',
                      'charge_increment':[0.100, -0.100] * unit.elementary_charge,
                      'distance': [0.220] * unit.angstrom,
                      'name':'aaa'})

sys = ff.create_openmm_system(mol.to_topology())

# For each vsite, print out which particles are involved, and at which distance.
for vsite_idx in range(sys.getNumParticles()):
    if not(sys.isVirtualSite(vsite_idx)):
        continue
    vsite = sys.getVirtualSite(vsite_idx)
    vsite_parents = []
    for vsite_parent_idx in range(vsite.getNumParticles()):
        vsite_parent = vsite.getParticle(vsite_parent_idx)
        vsite_parents.append(vsite_parent)
    distance = vsite.getLocalPosition()[0].value_in_unit(omm_unit.angstrom)
    print(vsite_parents, distance)

[0, 1] -0.101

The specific problems shown in this code are:

1. The first vsite applies to atoms (0,1), and the second applies to atoms (1,0). Since the atoms are in a different order, I don’t think that one should overwrite the other, and so I’d expect both vsites should be in the output[1]. Instead only a single vsite is in the output.
2. If the second virtualsite (with distance 0.220) MIGHT overwrite the first (with distance 0.101), then I expect that the output should EITHER have one vsite with each distance, OR one vsite with 0.220. Instead we get just one vsite with 0.101.

[1] It’s not specifically mentioned in the SMIRNOFF spec, but I assume this means “one vsite will overwrite another if they match the same atoms in the same order”

Computing environment (please complete the following information):

• Mac OS
• OFFTK topology biopolymer refactor branch

[mattwthompson]

So we have agreement on what should happen. It might be harder to modify things in a way that realizes that.

@j-wags and I hacked on this for a while today. We successfully identified several places in the code that do not fix it and at least one place that might be the best place to attack. Below is an ugly diff tracking some places in the code we poked around in - right now we think the failure is a result of if same_atoms and same_vsite and diff_keys: evaluating to True inside of _reduce_virtual_particles_to_sites, and at the end of today we’re at a loss as to how it could be patched to get the behavior we want. It might take a more extensive rewrite to handle these cross-interactions, I don’t currently have the energy or familiarity to help past that.

diff --git a/openff/toolkit/typing/engines/smirnoff/parameters.py b/openff/toolkit/typing/engines/smirnoff/parameters.py
index e170bb89..98b33eda 100644
--- a/openff/toolkit/typing/engines/smirnoff/parameters.py
+++ b/openff/toolkit/typing/engines/smirnoff/parameters.py
@@ -5273,7 +5273,9 @@ class VirtualSiteHandler(_NonbondedHandler):
             # has the match setting, which ultimately decides which orientations
             # to include.
             if self.match == "once":
+                # print(f"{orientations=}")
                 key = self.transformed_dict_cls.key_transform(orientations[0])
+                # print(f"{key=}")
                 orientations = [key]
                 # else all matches wanted, so keep whatever was matched.
 
@@ -5349,6 +5351,7 @@ class VirtualSiteHandler(_NonbondedHandler):
             ref_key = self.transformed_dict_cls.key_transform(orientations[0])
             atoms = list([molecule.atoms[i] for i in ref_key])
             args = (atoms, orientations)
+            # _add_virtual_site will ONLY PROCESS the FIRST value in the list `orientations`.
             off_idx = super()._add_virtual_site(fn, *args, replace=replace)
             return off_idx
 
@@ -5799,11 +5802,25 @@ class VirtualSiteHandler(_NonbondedHandler):
                     )
                     diff_keys = key not in vsite_struct[KEY_LIST]
                     same_vsite = self._same_virtual_site_type(vs_i, vs_j)
+                    # print(f"before, {combined_orientations=}")
+                    # import ipdb; ipdb.set_trace()
 
+                    # This conditional should probably not evaluate to True while reproducing issue #1206. Each
+                    # individual thing seems reasonable enough but something about this doesn't capture the breaking
+                    # case, and it's not clear how all of the necessary information can be gathered and processed
+                    # right here in a way that fixes #1206 _and_ doesn't break other stuff.
                     if same_atoms and same_vsite and diff_keys:
-                        combined_orientations[i][KEY_LIST].append(key)
+                        # print(f"{combined_orientations=}")
+                        # combined_orientations[i][VSITE_TYPE].append(vs_i)
+                        # The last match in the atom_matches iterator takes precendence according to SMIRNOFF rules;
+                        # since only the first orientation is processed by `_add_virtual_site` later, one idea is to
+                        # simply override the other orientations with the one that' scurrently found. This might need
+                        # a match="once" conditional.
+                        combined_orientations[i][KEY_LIST] = [key]
+                        combined_orientations[i][VSITE_TYPE] = vs_i
                         found = True
 
+                    # print(f"after, {combined_orientations=}")
                     # Skip out early since there is no reason to keep
                     # searching since we will never add the same
                     # particle twice
@@ -5850,13 +5867,19 @@ class VirtualSiteHandler(_NonbondedHandler):
 
             top_mol = Topology.from_molecules([molecule])
             matches = self.find_matches(top_mol, expand_permutations=True)
-
+            # print("[k for k in matches.keys()]")
+            # print([k for k in matches.keys()])
+            # In reproducing issue #1206, this method call is where information about the " 'distance': [0.220]" virtual site
+            # type is lost.
             virtual_sites = self._reduce_virtual_particles_to_sites(matches)
+            # print("[v[1] for v in virtual_sites]")
+            # print([v[1] for v in virtual_sites])
 
             # Now handle the vsites for this molecule
             # This call batches the key tuples into a single list, in order
             # for the virtual site to represent multiple particles
             for vsite_type, orientations in virtual_sites:
+                # orientations = [orientations[0]]  # import ipdb; ipdb.set_trace()
                 vsite_type.add_virtual_site(molecule, orientations, replace=True)
 
     def _create_openmm_virtual_sites(self, system, force, topology, molecule):

[trevorgokey]

Yeah, when I was looking at that failing test I was scratching my head at first. There are likely a few more tests than can hammer things out. The tests use a very symmetric problem, and I wouldn’t be surprised if your clever NaCl representation tricks show up in other places.

In any case, I am glad you are sitting down and really trying to hammer out a thorough test suite. As you can see, I only wrote the “does it work, in theory” tests, and you are making sure it works in practice. Much appreciated!