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 e1e4dd8b..27d2727e 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 import json import warnings from collections.abc import Iterable @@ -13,112 +14,59 @@ from ..utils import ensure_filelike -class Transformation(GufeTokenizable): - _stateA: ChemicalSystem - _stateB: ChemicalSystem - _name: Optional[str] - _mapping: Optional[Union[ComponentMapping, list[ComponentMapping]]] +class TransformationBase(GufeTokenizable): _protocol: Protocol + _name: Optional[str] def __init__( self, - stateA: ChemicalSystem, - stateB: ChemicalSystem, protocol: Protocol, - mapping: Optional[Union[ComponentMapping, list[ComponentMapping], dict[str, ComponentMapping]]] = None, name: Optional[str] = None, ): - r"""Two chemical states with a method for estimating free energy difference - - Connects two :class:`.ChemicalSystem` objects, with directionality, - and relates this to a :class:`.Protocol` which will provide an estimate of - the free energy difference of moving between these systems. - Used as an edge of an :class:`.AlchemicalNetwork`. + """Transformation base class. Parameters ---------- - stateA, stateB: ChemicalSystem - The start (A) and end (B) states of the transformation - protocol: Protocol - The method used to estimate the free energy difference between states - A and B - mapping : Optional[Union[ComponentMapping, list[ComponentMapping]]] - the details of any transformations between :class:`.Component` \s of - the two states + protocol : Protocol + The sampling method to use for the transformation. name : str, optional - a human-readable tag for this transformation - """ - if isinstance(mapping, dict): - 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 - self._stateB = stateB - self._mapping = mapping - self._name = name + A human-readable name for this transformation. + """ self._protocol = protocol + self._name = name @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})" - - @property - def stateA(self) -> ChemicalSystem: - """The starting :class:`.ChemicalSystem` for the transformation.""" - return self._stateA - - @property - def stateB(self) -> ChemicalSystem: - """The ending :class:`.ChemicalSystem` for the transformation.""" - return self._stateB - - @property - def protocol(self) -> Protocol: - """The 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""" - return self._mapping - @property def name(self) -> Optional[str]: - """ - Optional identifier for the transformation; used as part of its hash. + """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 + transformations to the same :class:`.AlchemicalNetwork` to avoid deduplication. """ return self._name - def _to_dict(self) -> dict: - return { - "stateA": self.stateA, - "stateB": self.stateB, - "protocol": self.protocol, - "mapping": self.mapping, - "name": self.name, - } - @classmethod def _from_dict(cls, d: dict): return cls(**d) + @property + @abc.abstractmethod + def stateA(self) -> ChemicalSystem: + """The starting :class:`.ChemicalSystem` for the transformation.""" + raise NotImplementedError + + @property + @abc.abstractmethod + def stateB(self) -> ChemicalSystem: + """The ending :class:`.ChemicalSystem` for the transformation.""" + raise NotImplementedError + + @abc.abstractmethod def create( self, *, @@ -126,32 +74,37 @@ def create( name: Optional[str] = None, ) -> ProtocolDAG: """ - Returns a ``ProtocolDAG`` executing this ``Transformation.protocol``. + Returns a :class:`.ProtocolDAG` executing this ``Transformation.protocol``. """ - return self.protocol.create( - stateA=self.stateA, - stateB=self.stateB, - mapping=self.mapping, - extends=extends, - name=name, - transformation_key=self.key, - ) + 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. - def gather(self, protocol_dag_results: Iterable[ProtocolDAGResult]) -> ProtocolResult: """ - Gather multiple ``ProtocolDAGResult`` into a single ``ProtocolResult``. + return self._protocol + + def gather(self, protocol_dag_results: Iterable[ProtocolDAGResult]) -> ProtocolResult: + """Gather multiple :class:`.ProtocolDAGResult` \s into a single + :class:`.ProtocolResult`. Parameters ---------- protocol_dag_results : Iterable[ProtocolDAGResult] - The ``ProtocolDAGResult`` objects to assemble aggregate quantities - from. + The :class:`.ProtocolDAGResult` objects to assemble aggregate + quantities from. Returns ------- ProtocolResult - Aggregated results from many ``ProtocolDAGResult`` objects, all from - a given ``Protocol``. + Aggregated results from many :class:`.ProtocolDAGResult` objects, + all from a given :class:`.Protocol`. """ return self.protocol.gather(protocol_dag_results=protocol_dag_results) @@ -159,14 +112,14 @@ def gather(self, protocol_dag_results: Iterable[ProtocolDAGResult]) -> ProtocolR 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. + 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. + 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) @@ -178,7 +131,7 @@ def load(cls, file): Parameters ---------- file : Union[PathLike, FileLike] - a pathlike or filelike to read this transformation from + 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) @@ -186,19 +139,106 @@ def load(cls, file): return cls.from_dict(dct) -# we subclass `Transformation` here for typing simplicity -class NonTransformation(Transformation): - """A non-alchemical edge of an alchemical network. +class Transformation(TransformationBase): + _stateA: ChemicalSystem + _stateB: ChemicalSystem + _protocol: Protocol + _mapping: Optional[Union[ComponentMapping, list[ComponentMapping]]] + _name: Optional[str] + + def __init__( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + protocol: Protocol, + mapping: Optional[Union[ComponentMapping, list[ComponentMapping], dict[str, ComponentMapping]]] = None, + name: Optional[str] = None, + ): + """Two chemical states with a method for estimating the free energy + difference between them. + + Connects two :class:`.ChemicalSystem` objects, with directionality, and + relates these to a :class:`.Protocol` which will provide an estimate of + the free energy difference between these systems. Used as an edge of an + :class:`.AlchemicalNetwork`. + + Parameters + ---------- + stateA, stateB : ChemicalSystem + The start (A) and end (B) states of the transformation. + protocol : Protocol + The method used to estimate the free energy difference between + states A and B. + mapping : Optional[Union[ComponentMapping, list[ComponentMapping]]] + The details of any transformations between :class:`.Component` \s + of the two states. + name : str, optional + A human-readable name for this transformation. - A "transformation" that performs no transformation at all. - Technically a self-loop, or an edge with the same ``ChemicalSystem`` at - either end. + """ + if isinstance(mapping, dict): + warnings.warn( + ("mapping input as a dict is deprecated, " "instead use either a single Mapping or list"), + DeprecationWarning, + ) + mapping = list(mapping.values()) - Functionally used for applying a dynamics protocol to a ``ChemicalSystem`` - that performs no alchemical transformation at all. This allows e.g. - equilibrium MD to be performed on a ``ChemicalSystem`` as desired alongside - alchemical protocols between it and and other ``ChemicalSystem`` objects. - """ + self._stateA = stateA + self._stateB = stateB + self._protocol = protocol + self._mapping = mapping + self._name = name + + def __repr__(self): + return f"{self.__class__.__name__}(stateA={self.stateA}, stateB={self.stateB}, protocol={self.protocol}, name={self.name})" + + @property + def stateA(self) -> ChemicalSystem: + """The starting :class:`.ChemicalSystem` for the transformation.""" + return self._stateA + + @property + def stateB(self) -> ChemicalSystem: + """The ending :class:`.ChemicalSystem` for the transformation.""" + return self._stateB + + @property + def mapping(self) -> Optional[Union[ComponentMapping, list[ComponentMapping]]]: + """The mappings relevant for this Transformation""" + return self._mapping + + def _to_dict(self) -> dict: + return { + "stateA": self.stateA, + "stateB": self.stateB, + "protocol": self.protocol, + "mapping": self.mapping, + "name": self.name, + } + + def create( + self, + *, + extends: Optional[ProtocolDAGResult] = None, + name: Optional[str] = None, + ) -> ProtocolDAG: + """ + Returns a ``ProtocolDAG`` executing this ``Transformation.protocol``. + """ + return self.protocol.create( + stateA=self.stateA, + stateB=self.stateB, + mapping=self.mapping, + extends=extends, + name=name, + transformation_key=self.key, + ) + + +class NonTransformation(TransformationBase): + _system: ChemicalSystem + _protocol: Protocol + _name: Optional[str] def __init__( self, @@ -206,33 +246,59 @@ def __init__( protocol: Protocol, name: Optional[str] = None, ): + """A non-alchemical edge of an alchemical network. + + A "transformation" that performs no transformation at all. + Technically a self-loop, or an edge with the same ``ChemicalSystem`` at + either end. + + Functionally used for applying a dynamics protocol to a ``ChemicalSystem`` + that performs no alchemical transformation at all. This allows e.g. + equilibrium MD to be performed on a ``ChemicalSystem`` as desired alongside + alchemical protocols between it and and other ``ChemicalSystem`` objects. + + Parameters + ---------- + system : ChemicalSystem + 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 + A human-readable name for this transformation. + + """ self._system = system - self._name = name self._protocol = protocol + self._name = name + + def __repr__(self): + return f"{self.__class__.__name__}(system={self.system}, protocol={self.protocol}, name={self.name})" @property - def stateA(self): + def stateA(self) -> ChemicalSystem: + """The :class:`.ChemicalSystem` this ``NonTransformation`` samples. + + Synonomous with ``system`` attribute and identical to ``stateB``. + + """ return self._system @property - def stateB(self): + def stateB(self) -> ChemicalSystem: + """The :class:`.ChemicalSystem` this ``NonTransformation`` samples. + + Synonomous with ``system`` attribute and identical to ``stateA``. + + + """ return self._system @property def system(self) -> ChemicalSystem: + """The :class:`.ChemicalSystem` this "transformation" samples.""" return self._system - @property - def protocol(self): - """ - The protocol for sampling dynamics of the `ChemicalSystem`. - - Includes all details needed to perform required - simulations/calculations. - """ - return self._protocol - def _to_dict(self) -> dict: return { "system": self.system, @@ -240,10 +306,6 @@ def _to_dict(self) -> dict: "name": self.name, } - @classmethod - def _from_dict(cls, d: dict): - return cls(**d) - def create( self, *, @@ -251,8 +313,11 @@ def create( name: Optional[str] = None, ) -> ProtocolDAG: """ - Returns a ``ProtocolDAG`` executing this ``Transformation.protocol``. + Returns a ``ProtocolDAG`` executing this ``NonTransformation.protocol``. """ + # TODO: once we have an implicit component mapping concept, use this + # here instead of None to allow use of alchemical protocols with + # NonTransformations return self.protocol.create( stateA=self.system, stateB=self.system,