From 9382d4435ec5c929e7c1c343c5de9f0fcc0a8847 Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Wed, 12 Jun 2024 20:40:59 +1000 Subject: [PATCH 01/11] Add lammps files created by tests to .gitignore --- .gitignore | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.gitignore b/.gitignore index 3787d7096..0213a2c78 100644 --- a/.gitignore +++ b/.gitignore @@ -119,6 +119,11 @@ md.log *stderr.txt *stdout.txt +# LAMMPS +log.lammps +out.lmp +tmp.in + # OS .DS_Store .DS_Store? From 04892290bae72be764bfe939c5ad101b75416f33 Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Wed, 12 Jun 2024 21:05:36 +1000 Subject: [PATCH 02/11] Implement equality checking between TopologyKeys and tuples of atom indices --- .../_tests/unit_tests/test_models.py | 56 +++++++++++++++++++ openff/interchange/models.py | 31 ++++++++-- setup.cfg | 2 +- 3 files changed, 84 insertions(+), 5 deletions(-) diff --git a/openff/interchange/_tests/unit_tests/test_models.py b/openff/interchange/_tests/unit_tests/test_models.py index 81f6609f9..935cfe444 100644 --- a/openff/interchange/_tests/unit_tests/test_models.py +++ b/openff/interchange/_tests/unit_tests/test_models.py @@ -1,4 +1,5 @@ from openff.interchange.models import ( + AngleKey, BondKey, ImproperTorsionKey, PotentialKey, @@ -110,3 +111,58 @@ 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), 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)) != (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), + 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) diff --git a/openff/interchange/models.py b/openff/interchange/models.py index f41aecc69..5965e94ab 100644 --- a/openff/interchange/models.py +++ b/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): """ diff --git a/setup.cfg b/setup.cfg index f72de44ca..80ea3b60f 100644 --- a/setup.cfg +++ b/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 From bfb8acf77d0e5fc2139d63876cbf2a32faecb89f Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Wed, 12 Jun 2024 21:14:45 +1000 Subject: [PATCH 03/11] Implement indexing on collections --- .../_tests/unit_tests/components/test_interchange.py | 7 +++++++ openff/interchange/components/potentials.py | 3 +++ 2 files changed, 10 insertions(+) diff --git a/openff/interchange/_tests/unit_tests/components/test_interchange.py b/openff/interchange/_tests/unit_tests/components/test_interchange.py index c61aa4e8f..4ad906543 100644 --- a/openff/interchange/_tests/unit_tests/components/test_interchange.py +++ b/openff/interchange/_tests/unit_tests/components/test_interchange.py @@ -48,6 +48,13 @@ 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"].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]) diff --git a/openff/interchange/components/potentials.py b/openff/interchange/components/potentials.py index 18ac37750..c806e78b8 100644 --- a/openff/interchange/components/potentials.py +++ b/openff/interchange/components/potentials.py @@ -316,3 +316,6 @@ def __getattr__(self, attr: str): return self.key_map else: return super().__getattribute__(attr) + + def __getitem__(self, key) -> Potential: + return self.potentials[self.key_map[key]] From d4db57cb981ca81ddd6577da4a9ece4657bfc755 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 12 Jun 2024 11:30:38 +0000 Subject: [PATCH 04/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../_tests/unit_tests/smirnoff/test_virtual_sites.py | 2 +- openff/interchange/components/_packmol.py | 8 ++++---- openff/interchange/interop/amber/export/_export.py | 2 +- openff/interchange/interop/openmm/__init__.py | 4 ++-- openff/interchange/smirnoff/_gromacs.py | 4 ++-- openff/interchange/smirnoff/_virtual_sites.py | 4 ++-- 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/openff/interchange/_tests/unit_tests/smirnoff/test_virtual_sites.py b/openff/interchange/_tests/unit_tests/smirnoff/test_virtual_sites.py index c27eda457..179733837 100644 --- a/openff/interchange/_tests/unit_tests/smirnoff/test_virtual_sites.py +++ b/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]], ), diff --git a/openff/interchange/components/_packmol.py b/openff/interchange/components/_packmol.py index eafa9798d..6ebf00af3 100644 --- a/openff/interchange/components/_packmol.py +++ b/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): diff --git a/openff/interchange/interop/amber/export/_export.py b/openff/interchange/interop/amber/export/_export.py index 008e88862..8239bce23 100644 --- a/openff/interchange/interop/amber/export/_export.py +++ b/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]) diff --git a/openff/interchange/interop/openmm/__init__.py b/openff/interchange/interop/openmm/__init__.py index fee99b29e..7e911b7d0 100644 --- a/openff/interchange/interop/openmm/__init__.py +++ b/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) diff --git a/openff/interchange/smirnoff/_gromacs.py b/openff/interchange/smirnoff/_gromacs.py index 2644d3fff..bb0e3f1a4 100644 --- a/openff/interchange/smirnoff/_gromacs.py +++ b/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 diff --git a/openff/interchange/smirnoff/_virtual_sites.py b/openff/interchange/smirnoff/_virtual_sites.py index eb596ad76..a130aedd6 100644 --- a/openff/interchange/smirnoff/_virtual_sites.py +++ b/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 From 9e59df536d654d2fe1ad244dfc46a39a010c0c04 Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Wed, 10 Jul 2024 15:15:09 +1000 Subject: [PATCH 05/11] Test equality of TopologyKeys with reversed indices --- openff/interchange/_tests/unit_tests/test_models.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/openff/interchange/_tests/unit_tests/test_models.py b/openff/interchange/_tests/unit_tests/test_models.py index 935cfe444..c8c012bfd 100644 --- a/openff/interchange/_tests/unit_tests/test_models.py +++ b/openff/interchange/_tests/unit_tests/test_models.py @@ -125,6 +125,7 @@ def test_bondkey_eq_hash(): 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) @@ -143,6 +144,7 @@ def test_anglekey_eq_hash(): """ 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) @@ -157,6 +159,7 @@ def test_torsionkey_eq_hash(): """ 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, From 569d9e4892a1fdf4037bec43540f50c18980aaf8 Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Wed, 10 Jul 2024 15:33:11 +1000 Subject: [PATCH 06/11] Indexing into a collection with a tuple looks up forward and reverse --- .../_tests/unit_tests/components/test_interchange.py | 1 + openff/interchange/components/potentials.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/openff/interchange/_tests/unit_tests/components/test_interchange.py b/openff/interchange/_tests/unit_tests/components/test_interchange.py index 4ad906543..fc5474973 100644 --- a/openff/interchange/_tests/unit_tests/components/test_interchange.py +++ b/openff/interchange/_tests/unit_tests/components/test_interchange.py @@ -52,6 +52,7 @@ def test_getitem(self, sage): 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]] ) diff --git a/openff/interchange/components/potentials.py b/openff/interchange/components/potentials.py index c806e78b8..7001f3afe 100644 --- a/openff/interchange/components/potentials.py +++ b/openff/interchange/components/potentials.py @@ -318,4 +318,11 @@ def __getattr__(self, attr: str): 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]] From 959d604bdb072823dbd8457bbb4b35edf83b6359 Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Thu, 11 Jul 2024 15:26:44 -0500 Subject: [PATCH 07/11] TST: Update test comparing tuple representation --- openff/interchange/_tests/unit_tests/test_models.py | 2 +- openff/interchange/models.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/openff/interchange/_tests/unit_tests/test_models.py b/openff/interchange/_tests/unit_tests/test_models.py index 935cfe444..81e90d996 100644 --- a/openff/interchange/_tests/unit_tests/test_models.py +++ b/openff/interchange/_tests/unit_tests/test_models.py @@ -130,7 +130,7 @@ def test_bondkey_eq_hash(): 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) + assert BondKey(atom_indices=(1, 3), bond_order=1.5) == ((1, 3), 1.5) def test_anglekey_eq_hash(): diff --git a/openff/interchange/models.py b/openff/interchange/models.py index beb2cc974..67e9a3c7f 100644 --- a/openff/interchange/models.py +++ b/openff/interchange/models.py @@ -76,7 +76,8 @@ 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)) + else: + return hash((tuple(self.atom_indices), self.bond_order)) def __eq__(self, other) -> bool: return super().__eq__(other) or ( From 4d3d81a8c3eb61e409cb9967508ac1f12b77fbe0 Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Fri, 12 Jul 2024 19:56:52 +1000 Subject: [PATCH 08/11] Add extended tuples and TopologyKeys to TopologyKey tests --- .../_tests/unit_tests/test_models.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/openff/interchange/_tests/unit_tests/test_models.py b/openff/interchange/_tests/unit_tests/test_models.py index 76f2dced9..617a66cb7 100644 --- a/openff/interchange/_tests/unit_tests/test_models.py +++ b/openff/interchange/_tests/unit_tests/test_models.py @@ -122,6 +122,8 @@ def test_bondkey_eq_hash(): 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) @@ -142,6 +144,8 @@ def test_anglekey_eq_hash(): 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) @@ -157,6 +161,12 @@ def test_torsionkey_eq_hash(): 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) @@ -169,3 +179,22 @@ def test_torsionkey_eq_hash(): 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, + ) From 28d2c609fe8b3db804ef6e3378372ac941056af7 Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Fri, 12 Jul 2024 20:00:06 +1000 Subject: [PATCH 09/11] Implement extended tuple TopologyKey equality and make equality testing more robust --- openff/interchange/models.py | 73 +++++++++++++++++++----------------- 1 file changed, 38 insertions(+), 35 deletions(-) diff --git a/openff/interchange/models.py b/openff/interchange/models.py index 67e9a3c7f..cd68240b8 100644 --- a/openff/interchange/models.py +++ b/openff/interchange/models.py @@ -45,11 +45,20 @@ class TopologyKey(_BaseModel, abc.ABC): description="The indices of the atoms occupied by this interaction", ) + def _tuple(self) -> tuple[Any, ...]: + """Tuple representation of this key.""" + return tuple(self.atom_indices) + def __hash__(self) -> int: - return hash(tuple(self.atom_indices)) + return hash(self._tuple()) def __eq__(self, other: Any) -> bool: - return self.__hash__() == other.__hash__() + if isinstance(other, tuple): + return self._tuple() == other + elif isinstance(other, TopologyKey): + return self._tuple() == other._tuple() + else: + return NotImplemented def __repr__(self) -> str: return f"{self.__class__.__name__} with atom indices {self.atom_indices}" @@ -73,16 +82,11 @@ class BondKey(TopologyKey): ), ) - def __hash__(self) -> int: + def _tuple(self) -> tuple[int, ...] | tuple[tuple[int, ...], float]: if self.bond_order is None: - return hash(tuple(self.atom_indices)) + return tuple(self.atom_indices) else: - 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 - ) + return (tuple(self.atom_indices), float(self.bond_order)) def __repr__(self) -> str: return ( @@ -100,11 +104,8 @@ class AngleKey(TopologyKey): description="The indices of the atoms occupied by this interaction", ) - def __hash__(self) -> int: - return hash(tuple(self.atom_indices)) - - def __eq__(self, other) -> bool: - return super().__eq__(other) or other == self.atom_indices + def _tuple(self) -> tuple[int, ...]: + return tuple(self.atom_indices) class ProperTorsionKey(TopologyKey): @@ -137,18 +138,22 @@ 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 _tuple( + self, + ) -> ( + tuple[()] + | tuple[int, int, int, int] + | tuple[ + tuple[int, int, int, int] | tuple[()], + int | None, + float | None, + float | None, + ] + ): + if self.mult is None and self.phase is None and self.bond_order is None: + return tuple(self.atom_indices) + else: + return (tuple(self.atom_indices), self.mult, self.phase, self.bond_order) def __repr__(self) -> str: return ( @@ -248,14 +253,12 @@ class VirtualSiteKey(TopologyKey): description="The `match` attribute of the associated virtual site type", ) - def __hash__(self) -> int: - return hash( - ( - self.orientation_atom_indices, - self.name, - self.type, - self.match, - ), + def _tuple(self) -> tuple[tuple[int, ...], str, str, str]: + return ( + self.orientation_atom_indices, + self.name, + self.type, + self.match, ) From c6c90132e933f9bb16a994d943f86abb9bc0defe Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Fri, 12 Jul 2024 20:01:17 +1000 Subject: [PATCH 10/11] Add examples to docstrings --- openff/interchange/models.py | 92 ++++++++++++++++++++++++++++++++++-- 1 file changed, 89 insertions(+), 3 deletions(-) diff --git a/openff/interchange/models.py b/openff/interchange/models.py index cd68240b8..ee07a8c7c 100644 --- a/openff/interchange/models.py +++ b/openff/interchange/models.py @@ -19,9 +19,12 @@ class TopologyKey(_BaseModel, abc.ABC): bond, but not the force constant or equilibrium bond length as determined by the force field. + Topology keys compare equal to (and hash the same as) tuples of their atom + indices as long as their other fields are `None`. + Examples -------- - Create a TopologyKey identifying some speicfic angle + Create a ``TopologyKey`` identifying some specific angle .. code-block:: pycon @@ -30,7 +33,7 @@ class TopologyKey(_BaseModel, abc.ABC): >>> this_angle TopologyKey with atom indices (2, 1, 3) - Create a TopologyKey indentifying just one atom + Create a ``TopologyKey`` indentifying just one atom .. code-block:: pycon @@ -38,6 +41,22 @@ class TopologyKey(_BaseModel, abc.ABC): >>> this_atom TopologyKey with atom indices (4,) + Compare a ``TopologyKey`` to a tuple containing the atom indices + + .. code-block:: pycon + + >>> key = TopologyKey(atom_indices=(0, 1)) + >>> key == (0, 1) + True + + Index into a dictionary with a tuple + + .. code-block:: pycon + + >>> d = {TopologyKey(atom_indices=(0, 1)): "some_bond"} + >>> d[0, 1] + 'some_bond' + """ # TODO: Swith to `pydantic.contuple` once 1.10.3 or 2.0.0 is released @@ -67,6 +86,22 @@ def __repr__(self) -> str: class BondKey(TopologyKey): """ A unique identifier of the atoms associated in a bond potential. + + Examples + -------- + Index into a dictionary with a tuple + + .. code-block:: pycon + + >>> d = { + ... BondKey(atom_indices=(0, 1)): "some_bond", + ... BondKey(atom_indices=(1, 2), bond_order=1.5): "some_other_bond", + ... } + >>> d[0, 1] + 'some_bond' + >>> d[(1, 2), 1.5] + 'some_other_bond' + """ atom_indices: tuple[int, int] = Field( @@ -98,6 +133,17 @@ def __repr__(self) -> str: class AngleKey(TopologyKey): """ A unique identifier of the atoms associated in an angle potential. + + Examples + -------- + Index into a dictionary with a tuple + + .. code-block:: pycon + + >>> d = {AngleKey(atom_indices=(0, 1, 2)): "some_angle"} + >>> d[0, 1, 2] + 'some_angle' + """ atom_indices: tuple[int, int, int] = Field( @@ -111,6 +157,25 @@ def _tuple(self) -> tuple[int, ...]: class ProperTorsionKey(TopologyKey): """ A unique identifier of the atoms associated in a proper torsion potential. + + Examples + -------- + Index into a dictionary with a tuple + + .. code-block:: pycon + + >>> d = { + ... ProperTorsionKey(atom_indices=(0, 1, 2, 3)): "torsion1", + ... ProperTorsionKey(atom_indices=(0, 1, 2, 3), mult=2): "torsion2", + ... ProperTorsionKey(atom_indices=(5, 6, 7, 8), mult=2, phase=0.78, bond_order=1.5): "torsion3", + ... } + >>> d[0, 1, 2, 3] + 'torsion1' + >>> d[(0, 1, 2, 3), 2, None, None] + 'torsion2' + >>> d[(5, 6, 7, 8), 2, 0.78, 1.5] + 'torsion3' + """ atom_indices: tuple[int, int, int, int] | tuple[()] = Field( @@ -168,6 +233,25 @@ class ImproperTorsionKey(ProperTorsionKey): A unique identifier of the atoms associated in an improper torsion potential. The central atom is the second atom in the `atom_indices` tuple, or accessible via `get_central_atom_index`. + + Examples + -------- + Index into a dictionary with a tuple + + .. code-block:: pycon + + >>> d = { + ... ImproperTorsionKey(atom_indices=(0, 1, 2, 3)): "torsion1", + ... ImproperTorsionKey(atom_indices=(0, 1, 2, 3), mult=2): "torsion2", + ... ImproperTorsionKey(atom_indices=(5, 6, 7, 8), mult=2, phase=0.78, bond_order=1.5): "torsion3", + ... } + >>> d[0, 1, 2, 3] + 'torsion1' + >>> d[(0, 1, 2, 3), 2, None, None] + 'torsion2' + >>> d[(5, 6, 7, 8), 2, 0.78, 1.5] + 'torsion3' + """ def get_central_atom_index(self) -> int: @@ -235,7 +319,9 @@ def __hash__(self) -> int: class VirtualSiteKey(TopologyKey): - """A unique identifier of a virtual site in the scope of a chemical topology.""" + """ + A unique identifier of a virtual site in the scope of a chemical topology. + """ # TODO: Overriding the attribute of a parent class is clumsy, but less grief than # having this not inherit from `TopologyKey`. It might be useful to just have From ed12deb4ee2357e9d9ff0103971517a6e9fbaff3 Mon Sep 17 00:00:00 2001 From: Josh Mitchell Date: Fri, 12 Jul 2024 20:19:52 +1000 Subject: [PATCH 11/11] Add docs for collection ergonomics --- docs/using/collections.md | 29 +++++++++++++++----- openff/interchange/components/interchange.py | 15 ++++++++++ 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/docs/using/collections.md b/docs/using/collections.md index 8bc142329..229e6752b 100644 --- a/docs/using/collections.md +++ b/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': , 'length': }, ...), @@ -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': , 'length': }, 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 diff --git a/openff/interchange/components/interchange.py b/openff/interchange/components/interchange.py index 75d070f78..ef13bd5cc 100644 --- a/openff/interchange/components/interchange.py +++ b/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())