`Interchange.combine(...).to_openmm` fails when mixing non-SMIRNOFF source with SMIRNOFF source with NAGL charges

[mattwthompson]

Description

Because of #1393, causing a key mismatch, the following workflow fails

• Create an Interchange from a non-SMIRNOFF source
• Create another Interchange with some partial charges assigned by NAGL
• Combine them
• Export the result to OpenMM

Reproduction

from openff.toolkit import ForceField, Quantity
from openff.interchange import Interchange
from openff.interchange._tests import MoleculeWithConformer
from openff.interchange.models import SingleAtomChargeTopologyKey


topology = MoleculeWithConformer.from_smiles("CCO").to_topology()
topology.box_vectors = Quantity([4, 4, 4], "nanometer")

interchange = ForceField("openff-2.3.0-rc1.offxml").create_interchange(topology)

combined = Interchange.from_openmm(
    system=interchange.to_openmm_system(),
    topology=interchange.topology,
    positions=None,
).combine(interchange.model_copy(deep=True))

try:
    combined.to_openmm_system()
except KeyError as error:
    print(
        [
            key
            for key in combined["Electrostatics"].key_map
            if type(key) is SingleAtomChargeTopologyKey and key.this_atom_index == 9
        ]
    )

    raise error

Output

$ python test_openmm_roundtrip_nagl.py
/Users/mattthompson/software/openff-interchange/openff/interchange/interop/openmm/_import/_import.py:121: MissingPositionsWarning: Nothing was passed to the `positions` argument, are you sure you don't want to pass positions?
  warnings.warn(
/Users/mattthompson/software/openff-interchange/openff/interchange/operations/_combine.py:104: InterchangeCombinationWarning: Interchange object combination is complex and may produce strange results outside of use cases it has been tested in. Use with caution and thoroughly validate results!
  warnings.warn(
/Users/mattthompson/software/openff-interchange/openff/interchange/operations/_combine.py:84: InterchangeCombinationWarning: Found electrostatics 1-4 scaling factors of 5/6 with slightly different rounding (0.833333 and 0.8333333333). This likely stems from OpenFF using more digits in rounding 1/1.2. The value of 0.8333333333 will be used, which may or may not introduce small errors.
  warnings.warn(
/Users/mattthompson/software/openff-interchange/openff/interchange/operations/_combine.py:132: UserWarning: One Interchange object has collection with name ImproperTorsions not found in the other Interchange object, but it has now been added.
  warnings.warn(
/Users/mattthompson/software/openff-interchange/openff/interchange/operations/_combine.py:183: InterchangeCombinationWarning: Setting positions to None because one or both objects added together were missing positions.
  warnings.warn(
[SingleAtomChargeTopologyKey(this_atom_index=9, extras={'handler': 'NAGLChargesHandler', 'partial_charge_method': 'openff-gnn-am1bcc-1.0.0.pt'})]
Traceback (most recent call last):
  File "/Users/mattthompson/software/openff-interchange/openff/interchange/interop/openmm/_nonbonded.py", line 390, in _create_single_nonbonded_force
    partial_charge = partial_charges[top_key].m_as("elementary_charge")
                     ~~~~~~~~~~~~~~~^^^^^^^^^
KeyError: TopologyKey with atom indices (9,)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/mattthompson/software/openff-interchange/test_openmm_roundtrip_nagl.py", line 29, in <module>
    raise error
  File "/Users/mattthompson/software/openff-interchange/test_openmm_roundtrip_nagl.py", line 19, in <module>
    combined.to_openmm_system()
  File "/Users/mattthompson/software/openff-interchange/openff/interchange/components/interchange.py", line 568, in to_openmm_system
    return _to_openmm_system(
           ^^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/mamba/envs/openff-interchange-dev/lib/python3.11/site-packages/openff/utilities/utilities.py", line 79, in wrapper
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/openff-interchange/openff/interchange/interop/openmm/__init__.py", line 95, in to_openmm_system
    particle_map = _process_nonbonded_forces(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/openff-interchange/openff/interchange/interop/openmm/_nonbonded.py", line 153, in _process_nonbonded_forces
    _func(
  File "/Users/mattthompson/software/openff-interchange/openff/interchange/interop/openmm/_nonbonded.py", line 399, in _create_single_nonbonded_force
    partial_charge = partial_charges[other_top_key].m_as("elementary_charge")
                     ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
KeyError: SingleAtomChargeTopologyKey(this_atom_index=9, extras={})

The output is long, but the core of the issue is that SingleAtomChargeTopologyKey stores metadata of the partial charge source but the OpenMM code is written to assume a key can be created with just the relevant atom index.

This is not a problem with pure SMIRNOFF sources because SMIRNOFFElectrostaticsCollection.get_charge_array uses a different code path.

Software versions

• Which operating system and version are you using?
• How did you install Interchange?
• What is the output of running conda list?