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?

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/smirnoff/test_virtual_sites.py

@@ -203,7 +203,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,61 @@ 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)) == (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)) == (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)) == (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)

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/potentials.py

@@ -316,3 +316,13 @@ def __getattr__(self, attr: str):
             return self.key_map
         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]]

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)

openff/interchange/models.py

@@ -57,7 +57,7 @@ class BondKey(TopologyKey):
     A unique identifier of the atoms associated in a bond potential.
     """
 
-    atom_indices: tuple[int, ...] = Field(
+    atom_indices: tuple[int, int] = Field(
         description="The indices of the atoms occupied by this interaction",
     )
 
@@ -71,8 +71,15 @@ class BondKey(TopologyKey):
     )
 
     def __hash__(self) -> int:
+        if self.bond_order is None:
+            return hash(tuple(self.atom_indices))
         return hash((tuple(self.atom_indices), self.bond_order))
 
+    def __eq__(self, other) -> bool:
+        return super().__eq__(other) or (
+            self.bond_order is None and other == self.atom_indices
+        )
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__} with atom indices {self.atom_indices}"
@@ -85,17 +92,20 @@ class AngleKey(TopologyKey):
     A unique identifier of the atoms associated in an angle potential.
     """
 
-    atom_indices: tuple[int, ...] = Field(
+    atom_indices: tuple[int, int, int] = Field(
         description="The indices of the atoms occupied by this interaction",
     )
 
+    def __eq__(self, other) -> bool:
+        return super().__eq__(other) or other == self.atom_indices
+
 
 class ProperTorsionKey(TopologyKey):
     """
     A unique identifier of the atoms associated in a proper torsion potential.
     """
 
-    atom_indices: tuple[int, ...] = Field(
+    atom_indices: tuple[int, int, int, int] | tuple[()] = Field(
         description="The indices of the atoms occupied by this interaction",
     )
 
@@ -121,8 +131,18 @@ class ProperTorsionKey(TopologyKey):
     )
 
     def __hash__(self) -> int:
+        if self.mult is None and self.bond_order is None and self.phase is None:
+            return hash(tuple(self.atom_indices))
         return hash((tuple(self.atom_indices), self.mult, self.bond_order, self.phase))
 
+    def __eq__(self, other) -> bool:
+        return super().__eq__(other) or (
+            self.mult is None
+            and self.bond_order is None
+            and self.phase is None
+            and other == tuple(self.atom_indices)
+        )
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__} with atom indices {self.atom_indices}"
@@ -154,13 +174,16 @@ class LibraryChargeTopologyKey(DefaultModel):
     this_atom_index: int
 
     @property
-    def atom_indices(self) -> tuple[int, ...]:
+    def atom_indices(self) -> tuple[int]:
         """Alias for `this_atom_index`."""
         return (self.this_atom_index,)
 
     def __hash__(self) -> int:
         return hash((self.this_atom_index,))
 
+    def __eq__(self, other) -> bool:
+        return super().__eq__(other) or other == self.this_atom_index
+
 
 class SingleAtomChargeTopologyKey(LibraryChargeTopologyKey):
     """

openff/interchange/smirnoff/_gromacs.py

@@ -758,8 +758,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))
         ):
 
             # these are molecule indices, whereas in the OpenMM function they are topology indices

openff/interchange/smirnoff/_virtual_sites.py

@@ -446,8 +446,8 @@ def _convert_local_coordinates(
     # rather than 0 degrees from the z-axis.
     vsite_positions = local_coordinate_frames[0] + d * (
         cos_theta * cos_phi * local_coordinate_frames[1]
-        + sin_theta * cos_phi * local_coordinate_frames[2]  # noqa
-        + sin_phi * local_coordinate_frames[3]  # noqa
+        + sin_theta * cos_phi * local_coordinate_frames[2]
+        + sin_phi * local_coordinate_frames[3]
     )
 
     return vsite_positions

setup.cfg

@@ -12,7 +12,7 @@ exclude_lines =
 
 [flake8]
 max-line-length = 119
-ignore = E203,B028
+ignore = E203,B028,W503
 per-file-ignores =
     openff/interchange/_tests/unit_tests/test_types.py:F821
     openff/interchange/**/__init__.py:F401