From fb4a68823ec32baf625a31335ab889be891a836e Mon Sep 17 00:00:00 2001 From: David Dotson Date: Wed, 24 Apr 2024 22:55:42 -0700 Subject: [PATCH 1/9] Make `Transformation` and `NonTransformation` subclass `TransformationBase` In some cases it can be awkward for `NonTransformation` to be a subclass of `Transformation`, such as in `alchemiscale`, for cases where `NonTransformation` should be handled very differently. Switching to a shared, abstract base class for `Transformation` and `NonTransformation` simplifies this. --- gufe/transformations/transformation.py | 189 ++++++++++++++----------- 1 file changed, 106 insertions(+), 83 deletions(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index 58a898be..25562cc8 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -1,6 +1,7 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/gufe +import abc from typing import Optional, Iterable, Union import json import warnings @@ -13,7 +14,95 @@ from ..mapping import ComponentMapping -class Transformation(GufeTokenizable): +class TransformationBase(GufeTokenizable): + """Transformation base class. + + """ + + @classmethod + def _defaults(cls): + return super()._defaults() + + @property + def name(self) -> Optional[str]: + """ + Optional identifier for the transformation; used as part of its hash. + + Set this to a unique value if adding multiple, otherwise identical + transformations to the same :class:`AlchemicalNetwork` to avoid + deduplication. + """ + return self._name + + @classmethod + def _from_dict(cls, d: dict): + return cls(**d) + + @abc.abstractmethod + def create( + self, + *, + extends: Optional[ProtocolDAGResult] = None, + name: Optional[str] = None, + ) -> ProtocolDAG: + """ + Returns a ``ProtocolDAG`` executing this ``Transformation.protocol``. + """ + raise NotImplementedError + + def gather( + self, protocol_dag_results: Iterable[ProtocolDAGResult] + ) -> ProtocolResult: + """ + Gather multiple ``ProtocolDAGResult`` into a single ``ProtocolResult``. + + Parameters + ---------- + protocol_dag_results : Iterable[ProtocolDAGResult] + The ``ProtocolDAGResult`` objects to assemble aggregate quantities + from. + + Returns + ------- + ProtocolResult + Aggregated results from many ``ProtocolDAGResult`` objects, all from + a given ``Protocol``. + + """ + return self.protocol.gather(protocol_dag_results=protocol_dag_results) + + def dump(self, file): + """Dump this Transformation to a JSON file. + + Note that this is not space-efficient: for example, any + ``Component`` which is used in both ``ChemicalSystem`` objects will be + represented twice in the JSON output. + + Parameters + ---------- + file : Union[PathLike, FileLike] + a pathlike of filelike to save this transformation to. + """ + with ensure_filelike(file, mode='w') as f: + json.dump(self.to_dict(), f, cls=JSON_HANDLER.encoder, + sort_keys=True) + + @classmethod + def load(cls, file): + """Create a Transformation from a JSON file. + + Parameters + ---------- + file : Union[PathLike, FileLike] + a pathlike or filelike to read this transformation from + """ + with ensure_filelike(file, mode='r') as f: + dct = json.load(f, cls=JSON_HANDLER.decoder) + + return cls.from_dict(dct) + + +class Transformation(TransformationBase): _stateA: ChemicalSystem _stateB: ChemicalSystem _name: Optional[str] @@ -56,18 +145,14 @@ def __init__( self._stateA = stateA self._stateB = stateB + self._protocol = protocol self._mapping = mapping self._name = name - self._protocol = protocol - - @classmethod - def _defaults(cls): - return super()._defaults() - def __repr__(self): - return f"{self.__class__.__name__}(stateA={self.stateA}, "\ - f"stateB={self.stateB}, protocol={self.protocol})" + attrs = ['name', 'stateA', 'stateB', 'protocol', 'mapping'] + content = ", ".join([f"{i}={getattr(self, i)}" for i in attrs]) + return f"{self.__class__.__name__}({content})" @property def stateA(self) -> ChemicalSystem: @@ -95,17 +180,6 @@ def mapping(self) -> Optional[Union[ComponentMapping, list[ComponentMapping]]]: """The mappings relevant for this Transformation""" return self._mapping - @property - def name(self) -> Optional[str]: - """ - Optional identifier for the transformation; used as part of its hash. - - Set this to a unique value if adding multiple, otherwise identical - transformations to the same :class:`AlchemicalNetwork` to avoid - deduplication. - """ - return self._name - def _to_dict(self) -> dict: return { "stateA": self.stateA, @@ -115,10 +189,6 @@ def _to_dict(self) -> dict: "name": self.name, } - @classmethod - def _from_dict(cls, d: dict): - return cls(**d) - def create( self, *, @@ -137,60 +207,8 @@ def create( transformation_key=self.key, ) - def gather( - self, protocol_dag_results: Iterable[ProtocolDAGResult] - ) -> ProtocolResult: - """ - Gather multiple ``ProtocolDAGResult`` into a single ``ProtocolResult``. - - Parameters - ---------- - protocol_dag_results : Iterable[ProtocolDAGResult] - The ``ProtocolDAGResult`` objects to assemble aggregate quantities - from. - - Returns - ------- - ProtocolResult - Aggregated results from many ``ProtocolDAGResult`` objects, all from - a given ``Protocol``. - - """ - return self.protocol.gather(protocol_dag_results=protocol_dag_results) - - def dump(self, file): - """Dump this Transformation to a JSON file. - - Note that this is not space-efficient: for example, any - ``Component`` which is used in both ``ChemicalSystem`` objects will be - represented twice in the JSON output. - - Parameters - ---------- - file : Union[PathLike, FileLike] - a pathlike of filelike to save this transformation to. - """ - with ensure_filelike(file, mode='w') as f: - json.dump(self.to_dict(), f, cls=JSON_HANDLER.encoder, - sort_keys=True) - - @classmethod - def load(cls, file): - """Create a Transformation from a JSON file. - - Parameters - ---------- - file : Union[PathLike, FileLike] - a pathlike or filelike to read this transformation from - """ - with ensure_filelike(file, mode='r') as f: - dct = json.load(f, cls=JSON_HANDLER.decoder) - - return cls.from_dict(dct) - -# we subclass `Transformation` here for typing simplicity -class NonTransformation(Transformation): +class NonTransformation(TransformationBase): """A non-alchemical edge of an alchemical network. A "transformation" that performs no transformation at all. @@ -211,8 +229,17 @@ def __init__( ): self._system = system - self._name = name self._protocol = protocol + self._name = name + + def __repr__(self): + return f"{self.__class__.__name__}(system={self.system}, "\ + f"protocol={self.protocol})" + + def __repr__(self): + attrs = ['name', 'system', 'protocol'] + content = ", ".join([f"{i}={getattr(self, i)}" for i in attrs]) + return f"{self.__class__.__name__}({content})" @property def stateA(self): @@ -243,10 +270,6 @@ def _to_dict(self) -> dict: "name": self.name, } - @classmethod - def _from_dict(cls, d: dict): - return cls(**d) - def create( self, *, @@ -254,7 +277,7 @@ def create( name: Optional[str] = None, ) -> ProtocolDAG: """ - Returns a ``ProtocolDAG`` executing this ``Transformation.protocol``. + Returns a ``ProtocolDAG`` executing this ``NonTransformation.protocol``. """ return self.protocol.create( stateA=self.system, From 0366d0e6f7b2bda1e95c966b036547d607847dbc Mon Sep 17 00:00:00 2001 From: David Dotson Date: Wed, 24 Apr 2024 23:02:54 -0700 Subject: [PATCH 2/9] Address mypy complaints --- gufe/transformations/transformation.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index 25562cc8..e018fdc5 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -18,6 +18,13 @@ class TransformationBase(GufeTokenizable): """Transformation base class. """ + def __init__( + self, + protocol: Protocol, + name: Optional[str] = None, + ): + self._protocol = protocol + self._name = name @classmethod def _defaults(cls): @@ -232,10 +239,6 @@ def __init__( self._protocol = protocol self._name = name - def __repr__(self): - return f"{self.__class__.__name__}(system={self.system}, "\ - f"protocol={self.protocol})" - def __repr__(self): attrs = ['name', 'system', 'protocol'] content = ", ".join([f"{i}={getattr(self, i)}" for i in attrs]) From e07909cc980c55139c42515005a3c515ea7074ce Mon Sep 17 00:00:00 2001 From: David Dotson Date: Wed, 4 Dec 2024 23:10:03 -0700 Subject: [PATCH 3/9] Updated repr tests for Transformation, NonTransformation --- gufe/tests/test_transformation.py | 4 ++-- gufe/transformations/transformation.py | 6 ++---- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/gufe/tests/test_transformation.py b/gufe/tests/test_transformation.py index 445ff651..5d4c1f74 100644 --- a/gufe/tests/test_transformation.py +++ b/gufe/tests/test_transformation.py @@ -34,7 +34,7 @@ def complex_equilibrium(solvated_complex): class TestTransformation(GufeTokenizableTestsMixin): cls = Transformation - repr = "Transformation(stateA=ChemicalSystem(name=, components={'ligand': SmallMoleculeComponent(name=toluene), 'solvent': SolventComponent(name=O, K+, Cl-)}), stateB=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=)" + repr = "Transformation(stateA=ChemicalSystem(name=, components={'ligand': SmallMoleculeComponent(name=toluene), 'solvent': SolventComponent(name=O, K+, Cl-)}), stateB=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=, name=None)" @pytest.fixture def instance(self, absolute_transformation): @@ -144,7 +144,7 @@ def test_deprecation_warning_on_dict_mapping(self, solvated_ligand, solvated_com class TestNonTransformation(GufeTokenizableTestsMixin): cls = NonTransformation - repr = "NonTransformation(stateA=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), stateB=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=)" + repr = "NonTransformation(system=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=, name=None)" @pytest.fixture def instance(self, complex_equilibrium): diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index e8df7a3b..b76a2dbc 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -174,7 +174,7 @@ def __init__( self._name = name def __repr__(self): - return f"{self.__class__.__name__}(stateA={self.stateA}, " f"stateB={self.stateB}, protocol={self.protocol})" + return f"{self.__class__.__name__}(stateA={self.stateA}, stateB={self.stateB}, protocol={self.protocol}, name={self.name})" @property def stateA(self) -> ChemicalSystem: @@ -265,9 +265,7 @@ def __init__( self._name = name def __repr__(self): - attrs = ['name', 'system', 'protocol'] - content = ", ".join([f"{i}={getattr(self, i)}" for i in attrs]) - return f"{self.__class__.__name__}({content})" + return f"{self.__class__.__name__}(system={self.system}, protocol={self.protocol}, name={self.name})" @property def stateA(self) -> ChemicalSystem: From 1597ec457840f9bb57105dc45f841431194918c1 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Wed, 4 Dec 2024 23:18:21 -0700 Subject: [PATCH 4/9] Moved protocol property to TranformationBase --- gufe/transformations/transformation.py | 32 ++++++++++---------------- 1 file changed, 12 insertions(+), 20 deletions(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index b76a2dbc..85bc00ec 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -75,6 +75,18 @@ def create( """ raise NotImplementedError + @property + def protocol(self) -> Protocol: + """The :class:`.Protocol` used to perform the transformation. + + The protocol estimates the free energy differences between ``stateA`` + and ``stateB`` :class:`.ChemicalSystem` objects. It includes all + details needed to perform required simulations/calculations and encodes + the alchemical or non-alchemical pathway used. + + """ + return self._protocol + def gather(self, protocol_dag_results: Iterable[ProtocolDAGResult]) -> ProtocolResult: """Gather multiple :class:`.ProtocolDAGResult` \s into a single :class:`.ProtocolResult`. @@ -186,17 +198,6 @@ def stateB(self) -> ChemicalSystem: """The ending :class:`.ChemicalSystem` for the transformation.""" return self._stateB - @property - def protocol(self) -> Protocol: - """The :class:`.Protocol` used to perform the transformation. - - This protocol estimates the free energy differences between ``stateA`` - and ``stateB`` :class:`.ChemicalSystem` objects. It includes all details - needed to perform required simulations/calculations and encodes the - alchemical pathway used. - """ - return self._protocol - @property def mapping(self) -> Optional[Union[ComponentMapping, list[ComponentMapping]]]: """The mappings relevant for this Transformation""" @@ -290,15 +291,6 @@ def system(self) -> ChemicalSystem: """The :class:`.ChemicalSystem` this "transformation" samples.""" return self._system - @property - def protocol(self): - """The :class:`.Protocol` for sampling dynamics of the ``system``. - - Includes all details needed to perform required - simulations/calculations. - """ - return self._protocol - def _to_dict(self) -> dict: return { "system": self.system, From 1c9daf055a193f84ce6486b2fbb852e6e116fddc Mon Sep 17 00:00:00 2001 From: "David L. Dotson" Date: Mon, 9 Dec 2024 22:29:44 -0700 Subject: [PATCH 5/9] Update gufe/transformations/transformation.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> --- gufe/transformations/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index 85bc00ec..bd0fa534 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -253,7 +253,7 @@ def __init__( Parameters ---------- system : ChemicalSystem - The (identical) end states of the "transformation" to be sampled + The system to be sampled, acting as both the starting and end state of the ``NonTransformation``. protocol : Protocol The sampling method to use on the ``system`` name : str, optional From b250faccd18cfbc0b3b27e58da7707738a6bdf7c Mon Sep 17 00:00:00 2001 From: "David L. Dotson" Date: Mon, 9 Dec 2024 22:36:08 -0700 Subject: [PATCH 6/9] Apply suggestions from code review Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> --- gufe/transformations/transformation.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index bd0fa534..9ce74cfa 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -270,18 +270,19 @@ def __repr__(self): @property def stateA(self) -> ChemicalSystem: - """The :class:`.ChemicalSystem` this "transformation" samples. + """The :class:`.ChemicalSystem` this ``NonTransformation`` samples. - Synonomous with ``system`` attribute. + Synonomous with ``system`` attribute and identical to `stateB`. """ return self._system @property def stateB(self) -> ChemicalSystem: - """The :class:`.ChemicalSystem` this "transformation" samples. + """The :class:`.ChemicalSystem` this ``NonTransformation`` samples. + + Synonomous with ``system`` attribute and identical to `stateA`. - Synonomous with ``system`` attribute. """ return self._system From 9a9f77b873314f3c289b7f1cdc91770795933f51 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 9 Dec 2024 22:37:54 -0700 Subject: [PATCH 7/9] Add attribute types to TransformationBase, NonTransformation on @atravitz suggestion --- gufe/transformations/transformation.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index 9ce74cfa..49f009b2 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -15,6 +15,9 @@ class TransformationBase(GufeTokenizable): + _protocol: Protocol + _name: Optional[str] + def __init__( self, protocol: Protocol, @@ -139,9 +142,9 @@ def load(cls, file): class Transformation(TransformationBase): _stateA: ChemicalSystem _stateB: ChemicalSystem - _name: Optional[str] - _mapping: Optional[Union[ComponentMapping, list[ComponentMapping]]] _protocol: Protocol + _mapping: Optional[Union[ComponentMapping, list[ComponentMapping]]] + _name: Optional[str] def __init__( self, @@ -232,6 +235,9 @@ def create( class NonTransformation(TransformationBase): + _system: ChemicalSystem + _protocol: Protocol + _name: Optional[str] def __init__( self, From 51ca9232c29ff68f8b9a383597fa8cf5e4e3baaa Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 9 Dec 2024 22:38:40 -0700 Subject: [PATCH 8/9] Backticks fix --- gufe/transformations/transformation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index 49f009b2..f48e92a6 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -278,7 +278,7 @@ def __repr__(self): def stateA(self) -> ChemicalSystem: """The :class:`.ChemicalSystem` this ``NonTransformation`` samples. - Synonomous with ``system`` attribute and identical to `stateB`. + Synonomous with ``system`` attribute and identical to ``stateB``. """ return self._system @@ -287,7 +287,7 @@ def stateA(self) -> ChemicalSystem: def stateB(self) -> ChemicalSystem: """The :class:`.ChemicalSystem` this ``NonTransformation`` samples. - Synonomous with ``system`` attribute and identical to `stateA`. + Synonomous with ``system`` attribute and identical to ``stateA``. """ From 45583b0c643ac5b7366f55a069fc84766d8ed9fb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 10 Dec 2024 05:39:09 +0000 Subject: [PATCH 9/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- gufe/transformations/transformation.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index f48e92a6..27d2727e 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -31,7 +31,7 @@ def __init__( The sampling method to use for the transformation. name : str, optional A human-readable name for this transformation. - + """ self._protocol = protocol self._name = name @@ -177,9 +177,10 @@ def __init__( """ if isinstance(mapping, dict): - warnings.warn(("mapping input as a dict is deprecated, " - "instead use either a single Mapping or list"), - DeprecationWarning) + warnings.warn( + ("mapping input as a dict is deprecated, " "instead use either a single Mapping or list"), + DeprecationWarning, + ) mapping = list(mapping.values()) self._stateA = stateA