Support importing (some) virtual sites in Interchange.from_openmm

[mattwthompson]

Description

Resolves #1063

Checklist

Stuff to do in later PRs

• Process LocalCoordinatesSites
• Handle getting positions of virtual sites mixes from OpenMM and SMIRNOFF sources
• Export virtual sites from OpenMM to GROMACS
• Refactor VirtualSiteType to be more explicit
• Consider storing virtual site and particle maps directly on Interchange objects
• Prototype subclassing Topology

[codecov]

Codecov Report

Attention: Patch coverage is 97.53695% with 5 lines in your changes missing coverage. Please review.

Project coverage is 93.79%. Comparing base (ed5e38b) to head (3ea2e99).
Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
openff/interchange/interop/_virtual_sites.py	92.59%	2 Missing ⚠️
...enff/interchange/interop/openmm/_import/_import.py	97.43%	1 Missing ⚠️
...terchange/interop/openmm/_import/_virtual_sites.py	96.55%	1 Missing ⚠️
openff/interchange/smirnoff/_gromacs.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1081      +/-   ##
==========================================
+ Coverage   93.60%   93.79%   +0.18%     
==========================================
  Files          70       71       +1     
  Lines        6039     6153     +114     
==========================================
+ Hits         5653     5771     +118     
+ Misses        386      382       -4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Yoshanuikabundi]

I tested this on a hand-built OpenMM Simulation featuring a single water molecule with made up parameters, and also on a single Tip4P molecule. Things mostly worked, and errors were mostly clear. I do have a few notes:

• export INTERCHANGE_EXPERIMENTAL=1 seems to be required, but this does not seem to be documented and the error message doesn't make this clear. This was the only time I was confused by an error message.
• At the moment, the number of atoms in the imported topology must match the number of particles in the system. I think we should consider if topologies matching the number of real atoms in the system should be accepted as well. This would allow OpenFF Topology objects to work when importing systems with supported virtual sites, as they currently do for systems without virtual sites. This makes particular sense given that it appears that virtual sites are stripped out of the OpenFF topology that's available from interchange.topology.
• The exported OpenMM topology via to_openmm_simulation() had its virtual sites in a separate chain and residue - I think it'd be better if they were in the same chain and residue as the parent atom.
• This is unrelated to the PR, but I tried a hand-built system without bonds (single water molecule with constraints, but no bonds force), and Interchange complained:

    136 # TODO: Does this run through the Interchange.box validator?
    137 interchange.box = _box_vectors
--> 139 if interchange.topology.n_bonds > len(interchange.collections["Bonds"].key_map):
    140     # There are probably missing (physics) bonds from rigid waters. The topological
    141     # bonds are probably processed correctly.
    142     _fill_in_rigid_water_bonds(interchange)
    144 return interchange

KeyError: 'Bonds'

Since this code currently attempts to account for bond parameters missing from the bond parameters, maybe we should account for this possibility?

A full review for such a big patch is gonna take me some time to work through, especially since I'm gonna need to get my head around how everything works, but those are my notes from some independent testing.

But yeah this is great! Things seem to be working pretty robustly once you come up with a system and topology that works - I was able to swap the topology of simple molecules out for a "real" openff topology after the Interchange was constructed and things kept working. Round trips work, energy minimization works from Interchange land, and visualizations work perfectly.

[Yoshanuikabundi]

Would it be reasonable to squish topology._particle_map and topology._molecule_virtual_site_map into atom metadata rather than monkey patching? I think the answer is probably "no", and I see this information is moved to the Interchange, and I've tested that Interchange.topology seems to be able to be swapped out for a true topology (with full molecular graphs etc).

[mattwthompson]

Possibly; I don't love the current design but it has the benefit of not interacting with anything the user might do in the public API

[mattwthompson]

#1201

[mattwthompson]

Thanks for the review

export INTERCHANGE_EXPERIMENTAL=1 seems to be required

This shouldn't be the case in the 0.4.x line, possibly your local setup is out of date. I've removed some obsolete stuff in tests now that Interchange.from_openmm is not "experimental."

I think we should consider if topologies matching the number of real atoms in the system should be accepted as well

Agree, but something for future work. I want to get this functionality over the finish line in its current scope as soon as I can

virtual sites in a separate chain and residue

Right, this was the result of #1051

we should account for this possibility

Probably yeah, topological edges being defined different ways is pretty frustrating to work with

[mattwthompson]

I think this is ready for another pass;

• Update experimental tag on examples, tests, experimental docs #1150 should have fixed some errant need for using INTERCHANGE_EXPERIMENTAL
• Based on Should from_openmm, when virtual sites are present, require virtual site positions? #1197, the first pass of this feature will require positions (if non-None) to be of length equal to the number of particles. I think supporting the case of being equal to the number of atoms is valid, if less common, and I can't think of a reason why it wouldn't work. But I want to push that off into a future changeset.

My hope is that this is more or less good to go in its current state, but if there are more changes needed before it finds a way into a release that's fine. OpenFE is a key stakeholder for this feature but I don't think there's significant time pressure from them.

I think everything else is either out of the original scope or something I forgot!