Merge pull request #990 from openforcefield/collection-ergonomics

Collection ergonomics

.gitignore

@@ -119,6 +119,11 @@ md.log
 *stderr.txt
 *stdout.txt
 
+# LAMMPS
+log.lammps
+out.lmp
+tmp.in
+
 # OS
 .DS_Store
 .DS_Store?

devtools/conda-envs/dev_env.yaml

@@ -4,22 +4,22 @@ channels:
   - openeye
 dependencies:
   # Core
-  - python =3.10
+  - python =3.11
   - pip
   - numpy
   - pydantic =2
   - openmm
   # OpenFF stack
   - openff-toolkit ~=0.16
   - openff-interchange-base
-  # smirnoff-plugins =2024
+  - smirnoff-plugins =2024
   - openff-nagl
   - openff-nagl-models
   - ambertools =23
   # Optional features
-  - mbuild =0.17
-  - foyer >=0.12.1
-  - gmso =0.12
+  - mbuild ~=0.17
+  - foyer ~=0.12
+  - gmso ~=0.12
   # Testing
   - mdtraj
   - intermol
@@ -31,7 +31,7 @@ dependencies:
   - nbval
   # de-forcefields  # add back after smirnoff-plugins update
   # Drivers
-  - gromacs
+  - gromacs =2024
   - lammps >=2023.08.02
   - panedr
   # Typing

docs/using/collections.md

@@ -91,13 +91,23 @@ SMIRNOFFImproperTorsionCollection(type='ImproperTorsions',
 
 ```
 
+We can also access a collection by indexing directly into the Interchange itself:
+
+```pycon
+>>> interchange['ImproperTorsions']  # doctest: +NORMALIZE_WHITESPACE,+ELLIPSIS
+SMIRNOFFImproperTorsionCollection(type='ImproperTorsions',
+    expression='k*(1+cos(periodicity*theta-phase))',
+    key_map={},
+    potentials={})
+```
+
 In the bond collection for example, each pair of bonded atoms maps to one of two
 potential keys, one for the carbon-carbon bond, and the other for the
 carbon-hydrogen bonds. It's clear from the SMIRKS codes that atoms 0 and 1 are
 the carbon atoms, and atoms 2 through 7 are the hydrogens:
 
 ```pycon
->>> interchange.collections['Bonds'].key_map  # doctest: +NORMALIZE_WHITESPACE,+ELLIPSIS
+>>> interchange['Bonds'].key_map  # doctest: +NORMALIZE_WHITESPACE,+ELLIPSIS
 {TopologyKey(atom_indices=(0, 1), ...): PotentialKey(id='[#6X4:1]-[#6X4:2]', ...),
  TopologyKey(atom_indices=(0, 2), ...): PotentialKey(id='[#6X4:1]-[#1:2]', ...),
  TopologyKey(atom_indices=(0, 3), ...): PotentialKey(id='[#6X4:1]-[#1:2]', ...),
@@ -117,7 +127,7 @@ The bond collection also maps the two potential keys to the appropriate `Potenti
 Here we can read off the force constant and length:
 
 ```pycon
->>> interchange.collections['Bonds'].potentials  # doctest: +NORMALIZE_WHITESPACE,+ELLIPSIS
+>>> interchange['Bonds'].potentials  # doctest: +NORMALIZE_WHITESPACE,+ELLIPSIS
 {PotentialKey(id='[#6X4:1]-[#6X4:2]', ...):
      Potential(parameters={'k': <Quantity(529.242972, 'kilocalorie / angstrom ** 2 / mole')>,
                            'length': <Quantity(1.52190126, 'angstrom')>}, ...),
@@ -127,16 +137,21 @@ Here we can read off the force constant and length:
 
 ```
 
-We can even modify a value here, export the new interchange, and see that all of
-the bonds have been updated:
+Any `TopologyKey` that only specifies atom indices can be accessed by indexing directly into the `Collection` with those atom indices:
+
+```pycon
+>>> interchange['Bonds'][0,1]  # doctest: +NORMALIZE_WHITESPACE,+ELLIPSIS
+Potential(parameters={'k': <Quantity(529.242972, 'kilocalorie / angstrom ** 2 / mole')>, 'length': <Quantity(1.52190126, 'angstrom')>}, map_key=None)
+
+```
+
+We can even modify a value here, export the new interchange, and see that all of the bonds have been updated:
 
 ```pycon
 >>> from openff.interchange.models import TopologyKey
 >>> from openff.units import unit
 >>> # Get the potential from the first C-H bond
->>> top_key = TopologyKey(atom_indices=(0, 2))
->>> pot_key = interchange.collections['Bonds'].key_map[top_key]
->>> potential = interchange.collections['Bonds'].potentials[pot_key]
+>>> potential = interchange['Bonds'][0, 2]
 >>> # Modify the potential
 >>> potential.parameters['length'] = 3.1415926 * unit.nanometer
 >>> # Write out the modified interchange to a GROMACS .top file

openff/interchange/_tests/unit_tests/components/test_interchange.py

@@ -48,6 +48,14 @@ def test_getitem(self, sage):
         with pytest.raises(LookupError, match="Could not find"):
             out["CMAPs"]
 
+        first_bondkey = next(iter(out["Bonds"].key_map))
+        idx_a, idx_b = first_bondkey.atom_indices
+        assert (
+            out["Bonds"][idx_a, idx_b]
+            == out["Bonds"][idx_b, idx_a]
+            == out["Bonds"].potentials[out["Bonds"].key_map[first_bondkey]]
+        )
+
     def test_get_parameters(self, sage):
         mol = Molecule.from_smiles("CCO")
         out = Interchange.from_smirnoff(force_field=sage, topology=[mol])

openff/interchange/_tests/unit_tests/interop/openmm/test_virtual_sites.py

@@ -330,7 +330,6 @@ def test_tip5p_num_exceptions(self, water, tip5p, combine, n_molecules):
                     # Safeguard against some of the behavior seen in #919
                     for index in range(num_exceptions):
                         p1, p2, *_ = force.getExceptionParameters(index)
-                        print(p1, p2)
 
                         if sorted([p1, p2]) == [0, 3]:
                             raise Exception(

openff/interchange/_tests/unit_tests/smirnoff/test_virtual_sites.py

@@ -205,7 +205,7 @@ def generate_v_site_coordinates(
                 (0, 1, 2, 3),
                 (
                     VirtualSiteMocking.sp2_conformer()[0]
-                    + Quantity(  # noqa
+                    + Quantity(
                         numpy.array(
                             [[1.0, numpy.sqrt(2), 1.0], [1.0, -numpy.sqrt(2), -1.0]],
                         ),

openff/interchange/_tests/unit_tests/test_models.py

@@ -1,4 +1,5 @@
 from openff.interchange.models import (
+    AngleKey,
     BondKey,
     ImproperTorsionKey,
     PotentialKey,
@@ -110,3 +111,90 @@ def test_reprs():
     assert "blah" in repr(potential_key)
     assert "mult 2" in repr(potential_key)
     assert "bond order 1.111" in repr(potential_key)
+
+
+def test_bondkey_eq_hash():
+    """
+    When __eq__ is true, the hashes must be equal.
+
+    The converse is not required in Python; hash collisions between unequal
+    objects are allowed and will be handled according to __eq__, possibly
+    with a small runtime cost.
+    """
+
+    assert BondKey(atom_indices=(1, 3)) == BondKey(atom_indices=(1, 3))
+    assert BondKey(atom_indices=(1, 3)) == TopologyKey(atom_indices=(1, 3))
+    assert BondKey(atom_indices=(1, 3)) == (1, 3)
+    assert hash(BondKey(atom_indices=(1, 3))) == hash((1, 3))
+    assert BondKey(atom_indices=(1, 3)) != (1, 4)
+    assert BondKey(atom_indices=(1, 3)) != (3, 1)
+    assert BondKey(atom_indices=(1, 3), bond_order=None) == (1, 3)
+    assert hash(BondKey(atom_indices=(1, 3), bond_order=None)) == hash((1, 3))
+    assert BondKey(atom_indices=(1, 3), bond_order=None) != ((1, 3), None)
+    assert BondKey(atom_indices=(1, 3), bond_order=1.5) != (1, 3)
+    assert BondKey(atom_indices=(1, 3), bond_order=1.5) != ((1, 3), None)
+    assert BondKey(atom_indices=(1, 3), bond_order=1.5) == ((1, 3), 1.5)
+
+
+def test_anglekey_eq_hash():
+    """
+    When __eq__ is true, the hashes must be equal.
+
+    The converse is not required in Python; hash collisions between unequal
+    objects are allowed and will be handled according to __eq__, possibly
+    with a small runtime cost.
+    """
+    assert AngleKey(atom_indices=(1, 3, 16)) == AngleKey(atom_indices=(1, 3, 16))
+    assert AngleKey(atom_indices=(1, 3, 16)) == TopologyKey(atom_indices=(1, 3, 16))
+    assert AngleKey(atom_indices=(1, 3, 16)) == (1, 3, 16)
+    assert hash(AngleKey(atom_indices=(1, 3, 16))) == hash((1, 3, 16))
+    assert AngleKey(atom_indices=(1, 3, 16)) != (16, 3, 1)
+    assert AngleKey(atom_indices=(1, 3, 16)) != (1, 3)
+    assert AngleKey(atom_indices=(1, 3, 16)) != (1, 3, 15)
+
+
+def test_torsionkey_eq_hash():
+    """
+    When __eq__ is true, the hashes must be equal.
+
+    The converse is not required in Python; hash collisions between unequal
+    objects are allowed and will be handled according to __eq__, possibly
+    with a small runtime cost.
+    """
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4)) == ProperTorsionKey(
+        atom_indices=(1, 2, 3, 4),
+    )
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4)) == TopologyKey(
+        atom_indices=(1, 2, 3, 4),
+    )
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4)) == (1, 2, 3, 4)
+    assert hash(ProperTorsionKey(atom_indices=(1, 2, 3, 4))) == hash((1, 2, 3, 4))
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4)) != (4, 3, 2, 1)
+    assert ProperTorsionKey(
+        atom_indices=(1, 2, 3, 4),
+        mult=None,
+        phase=None,
+        bond_order=None,
+    ) == (1, 2, 3, 4)
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4), mult=0) != (1, 2, 3, 4)
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4), phase=1.5) != (1, 2, 3, 4)
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4), bond_order=1.5) != (1, 2, 3, 4)
+
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4), mult=0) == (
+        (1, 2, 3, 4),
+        0,
+        None,
+        None,
+    )
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4), phase=1.5) == (
+        (1, 2, 3, 4),
+        None,
+        1.5,
+        None,
+    )
+    assert ProperTorsionKey(atom_indices=(1, 2, 3, 4), bond_order=1.5) == (
+        (1, 2, 3, 4),
+        None,
+        None,
+        1.5,
+    )

openff/interchange/components/_packmol.py

@@ -134,17 +134,17 @@ def _validate_inputs(
     """
     if (
         box_vectors is None
-        and mass_density is None  # noqa: W503
-        and (solute is None or solute.box_vectors is None)  # noqa: W503
+        and mass_density is None
+        and (solute is None or solute.box_vectors is None)
     ):
         raise PACKMOLValueError(
             "One of `box_vectors`, `mass_density`, or"
-            + " `solute.box_vectors` must be specified.",  # noqa: W503
+            + " `solute.box_vectors` must be specified.",
         )
     if box_vectors is not None and mass_density is not None:
         raise PACKMOLValueError(
             "`box_vectors` and `mass_density` cannot be specified together;"
-            + " choose one or the other.",  # noqa: W503
+            + " choose one or the other.",
         )
 
     if box_vectors is not None and box_vectors.shape != (3, 3):

openff/interchange/components/interchange.py

@@ -56,6 +56,21 @@ class Interchange(_BaseModel):
 
     .. warning :: This object is in an early and experimental state and unsuitable for production.
     .. warning :: This API is experimental and subject to change.
+
+    Examples
+    --------
+    Create an ``Interchange`` from an OpenFF ``ForceField`` and ``Molecule``
+
+    >>> from openff.toolkit import ForceField, Molecule
+    >>> sage = ForceField("openff-2.2.0.offxml")
+    >>> top = Molecule.from_smiles("CCC").to_topology()
+    >>> interchange = sage.create_interchange(top)
+
+    Get the parameters for the bond between atoms 0 and 1
+
+    >>> interchange["Bonds"][0, 1]
+    Potential(...)
+
     """
 
     collections: _AnnotatedCollections = Field(dict())

openff/interchange/components/potentials.py

@@ -419,6 +419,16 @@ def __getattr__(self, attr: str):
         else:
             return super().__getattribute__(attr)
 
+    def __getitem__(self, key) -> Potential:
+        if (
+            isinstance(key, tuple)
+            and key not in self.key_map
+            and tuple(reversed(key)) in self.key_map
+        ):
+            return self.potentials[self.key_map[tuple(reversed(key))]]
+
+        return self.potentials[self.key_map[key]]
+
 
 def validate_collections(
     v: Any,

openff/interchange/interop/amber/export/_export.py

@@ -640,7 +640,7 @@ def to_prmtop(interchange: "Interchange", file_path: Path | str):
         prmtop.write("%FLAG ANGLE_FORCE_CONSTANT\n" "%FORMAT(5E16.8)\n")
         angle_k = [
             interchange["Angles"].potentials[key].parameters["k"].m_as(kcal_mol_rad2)
-            / 2  # noqa
+            / 2
             for key in potential_key_to_angle_type_mapping
         ]
         text_blob = "".join([f"{val:16.8E}" for val in angle_k])

openff/interchange/interop/openmm/__init__.py

@@ -204,8 +204,8 @@ def _is_water(molecule: Molecule) -> bool:
         # TODO: This should only skip rigid waters, even though HMR or flexible water is questionable
         if (
             (hydrogen_atom.atomic_number == 1)
-            and (heavy_atom.atomic_number != 1)  # noqa: W503
-            and not (_is_water(hydrogen_atom.molecule))  # noqa: W503
+            and (heavy_atom.atomic_number != 1)
+            and not (_is_water(hydrogen_atom.molecule))
         ):
 
             hydrogen_index = interchange.topology.atom_index(hydrogen_atom)