From 8c598638257af6e4f50988238d8130f49d8f65aa Mon Sep 17 00:00:00 2001 From: riesben Date: Sat, 31 Aug 2024 22:02:25 +0200 Subject: [PATCH 1/3] initial drafting --- gufe/__init__.py | 10 +- gufe/{network.py => alchemical_network.py} | 0 gufe/mapping/__init__.py | 7 - gufe/protocols/protocol.py | 4 +- gufe/setup/__init__.py | 0 gufe/setup/alchemical_network_planner.py | 16 + gufe/setup/chemical_system_generator.py | 25 ++ gufe/setup/network_planning/__init__.py | 0 .../network_planning/atom_based/__init__.py | 0 .../atom_based}/atom_mapper.py | 2 +- .../atom_based}/atom_mapping.py | 2 +- .../atom_based/ligand_atom_mapping.py} | 4 +- .../atom_based}/ligandnetwork.py | 4 +- .../network_planning/component_mapper.py | 0 .../network_planning/component_mapping.py} | 0 .../component_mapping_scorer.py | 37 ++ gufe/setup/network_planning/network_plan.py | 318 ++++++++++++++++++ .../setup/network_planning/network_planner.py | 0 gufe/transformations/transformation.py | 2 +- 19 files changed, 407 insertions(+), 24 deletions(-) rename gufe/{network.py => alchemical_network.py} (100%) delete mode 100644 gufe/mapping/__init__.py create mode 100644 gufe/setup/__init__.py create mode 100644 gufe/setup/alchemical_network_planner.py create mode 100644 gufe/setup/chemical_system_generator.py create mode 100644 gufe/setup/network_planning/__init__.py create mode 100644 gufe/setup/network_planning/atom_based/__init__.py rename gufe/{mapping => setup/network_planning/atom_based}/atom_mapper.py (95%) rename gufe/{mapping => setup/network_planning/atom_based}/atom_mapping.py (95%) rename gufe/{mapping/ligandatommapping.py => setup/network_planning/atom_based/ligand_atom_mapping.py} (98%) rename gufe/{ => setup/network_planning/atom_based}/ligandnetwork.py (99%) create mode 100644 gufe/setup/network_planning/component_mapper.py rename gufe/{mapping/componentmapping.py => setup/network_planning/component_mapping.py} (100%) create mode 100644 gufe/setup/network_planning/component_mapping_scorer.py create mode 100644 gufe/setup/network_planning/network_plan.py create mode 100644 gufe/setup/network_planning/network_planner.py diff --git a/gufe/__init__.py b/gufe/__init__.py index 10e12d57..891f059a 100644 --- a/gufe/__init__.py +++ b/gufe/__init__.py @@ -16,12 +16,6 @@ from .chemicalsystem import ChemicalSystem -from .mapping import ( - ComponentMapping, # how individual Components relate - AtomMapping, AtomMapper, # more specific to atom based components - LigandAtomMapping, -) - from .settings import Settings from .protocols import ( @@ -36,7 +30,7 @@ from .transformations import Transformation, NonTransformation -from .network import AlchemicalNetwork -from .ligandnetwork import LigandNetwork +from .alchemical_network import AlchemicalNetwork +from gufe.setup.network_planning.atom_based.ligandnetwork import LigandNetwork __version__ = version("gufe") diff --git a/gufe/network.py b/gufe/alchemical_network.py similarity index 100% rename from gufe/network.py rename to gufe/alchemical_network.py diff --git a/gufe/mapping/__init__.py b/gufe/mapping/__init__.py deleted file mode 100644 index bd6be51b..00000000 --- a/gufe/mapping/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# This code is part of gufe and is licensed under the MIT license. -# For details, see https://github.com/OpenFreeEnergy/gufe -"""Defining the relationship between different components""" -from .componentmapping import ComponentMapping -from .atom_mapping import AtomMapping -from .atom_mapper import AtomMapper -from .ligandatommapping import LigandAtomMapping diff --git a/gufe/protocols/protocol.py b/gufe/protocols/protocol.py index bfda10c8..6edab43b 100644 --- a/gufe/protocols/protocol.py +++ b/gufe/protocols/protocol.py @@ -10,10 +10,10 @@ from openff.units import Quantity import warnings -from ..settings import Settings, SettingsBaseModel +from ..settings import Settings from ..tokenization import GufeTokenizable, GufeKey from ..chemicalsystem import ChemicalSystem -from ..mapping import ComponentMapping +from gufe.mapping import ComponentMapping from .protocoldag import ProtocolDAG, ProtocolDAGResult from .protocolunit import ProtocolUnit diff --git a/gufe/setup/__init__.py b/gufe/setup/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/gufe/setup/alchemical_network_planner.py b/gufe/setup/alchemical_network_planner.py new file mode 100644 index 00000000..1e7bb1de --- /dev/null +++ b/gufe/setup/alchemical_network_planner.py @@ -0,0 +1,16 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +import abc +from typing import Iterable +from .. import AlchemicalNetwork + + +class AlchemicalNetworkPlanner(abc.ABC): + """ + this abstract class defines the interface for the alchemical Network Planners. + """ + + @abc.abstractmethod + def __call__(self, *args, **kwargs) -> AlchemicalNetwork: + raise NotImplementedError() diff --git a/gufe/setup/chemical_system_generator.py b/gufe/setup/chemical_system_generator.py new file mode 100644 index 00000000..8b238f7a --- /dev/null +++ b/gufe/setup/chemical_system_generator.py @@ -0,0 +1,25 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +import abc +from enum import Enum + +from typing import Iterable +from gufe import ChemicalSystem + +# Todo: connect to protocols - use this for labels? + +class RFEComponentLabels(str, Enum): + PROTEIN = "protein" + LIGAND = "ligand" + SOLVENT = "solvent" + COFACTOR = "cofactor" + + +class AbstractChemicalSystemGenerator(abc.ABC): + """ + this abstract class defines the interface for the chemical system generators. + """ + + @abc.abstractmethod + def __call__(self, *args, **kwargs) -> Iterable[ChemicalSystem]: + raise NotImplementedError() \ No newline at end of file diff --git a/gufe/setup/network_planning/__init__.py b/gufe/setup/network_planning/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/gufe/setup/network_planning/atom_based/__init__.py b/gufe/setup/network_planning/atom_based/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/gufe/mapping/atom_mapper.py b/gufe/setup/network_planning/atom_based/atom_mapper.py similarity index 95% rename from gufe/mapping/atom_mapper.py rename to gufe/setup/network_planning/atom_based/atom_mapper.py index 0fe1c9cd..f33ef7c3 100644 --- a/gufe/mapping/atom_mapper.py +++ b/gufe/setup/network_planning/atom_based/atom_mapper.py @@ -4,7 +4,7 @@ from collections.abc import Iterator import gufe -from ..tokenization import GufeTokenizable +from gufe.tokenization import GufeTokenizable from .atom_mapping import AtomMapping diff --git a/gufe/mapping/atom_mapping.py b/gufe/setup/network_planning/atom_based/atom_mapping.py similarity index 95% rename from gufe/mapping/atom_mapping.py rename to gufe/setup/network_planning/atom_based/atom_mapping.py index 58c617f3..03cb0ac5 100644 --- a/gufe/mapping/atom_mapping.py +++ b/gufe/setup/network_planning/atom_based/atom_mapping.py @@ -6,7 +6,7 @@ import gufe -from .componentmapping import ComponentMapping +from gufe.setup.network_planning.component_mapping import ComponentMapping class AtomMapping(ComponentMapping, abc.ABC): diff --git a/gufe/mapping/ligandatommapping.py b/gufe/setup/network_planning/atom_based/ligand_atom_mapping.py similarity index 98% rename from gufe/mapping/ligandatommapping.py rename to gufe/setup/network_planning/atom_based/ligand_atom_mapping.py index 9a5b198b..0ecf87ee 100644 --- a/gufe/mapping/ligandatommapping.py +++ b/gufe/setup/network_planning/atom_based/ligand_atom_mapping.py @@ -9,8 +9,8 @@ from gufe.components import SmallMoleculeComponent from gufe.visualization.mapping_visualization import draw_mapping -from . import AtomMapping -from ..tokenization import JSON_HANDLER +from gufe.mapping import AtomMapping +from gufe.tokenization import JSON_HANDLER class LigandAtomMapping(AtomMapping): diff --git a/gufe/ligandnetwork.py b/gufe/setup/network_planning/atom_based/ligandnetwork.py similarity index 99% rename from gufe/ligandnetwork.py rename to gufe/setup/network_planning/atom_based/ligandnetwork.py index 177aa61e..62751570 100644 --- a/gufe/ligandnetwork.py +++ b/gufe/setup/network_planning/atom_based/ligandnetwork.py @@ -9,8 +9,8 @@ import gufe from gufe import SmallMoleculeComponent -from .mapping import LigandAtomMapping -from .tokenization import GufeTokenizable +from gufe.mapping import LigandAtomMapping +from gufe.tokenization import GufeTokenizable class LigandNetwork(GufeTokenizable): diff --git a/gufe/setup/network_planning/component_mapper.py b/gufe/setup/network_planning/component_mapper.py new file mode 100644 index 00000000..e69de29b diff --git a/gufe/mapping/componentmapping.py b/gufe/setup/network_planning/component_mapping.py similarity index 100% rename from gufe/mapping/componentmapping.py rename to gufe/setup/network_planning/component_mapping.py diff --git a/gufe/setup/network_planning/component_mapping_scorer.py b/gufe/setup/network_planning/component_mapping_scorer.py new file mode 100644 index 00000000..a842f498 --- /dev/null +++ b/gufe/setup/network_planning/component_mapping_scorer.py @@ -0,0 +1,37 @@ +# This code is part of kartograf and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/gufe + +import abc +from gufe.tokenization import GufeTokenizable + +from gufe.mapping import AtomMapping + + +class ComponentMappingScorer(GufeTokenizable): + """A generic class for scoring Atom mappings. + this class can be used for example to build graph algorithm based networks. + Implementations of this class can require an arbitrary and non-standardised + number of input arguments to create. + Implementations of this class provide the :meth:`.get_score` method + """ + + def __call__(self, mapping: ComponentMappingScorer) -> float: + return self.get_score(mapping) + + @abc.abstractmethod + def get_score(self, mapping: ComponentMappingScorer) -> float: + """ calculate the score for an :class:`.AtomMapping` + the scoring function returns a value between 0 and 1. + a value close to 1.0 indicates a small change, a score close to zero indicates a large cost/change. + Parameters + ---------- + mapping: AtomMapping + the mapping to be scored + args + kwargs + Returns + ------- + float + a value between [0,1] where zero is a very bad score and one a very good one. + """ + pass \ No newline at end of file diff --git a/gufe/setup/network_planning/network_plan.py b/gufe/setup/network_planning/network_plan.py new file mode 100644 index 00000000..d60e84fe --- /dev/null +++ b/gufe/setup/network_planning/network_plan.py @@ -0,0 +1,318 @@ +# This code is part of gufe and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/gufe +from __future__ import annotations + +from itertools import chain +import json +import networkx as nx +from typing import FrozenSet, Iterable, Optional +import gufe + +from gufe import SmallMoleculeComponent +from gufe.mapping import LigandAtomMapping +from gufe.tokenization import GufeTokenizable + + +class NetworkPlan(GufeTokenizable): + """A directed graph connecting many ligands according to their atom mapping + + Parameters + ---------- + edges : Iterable[LigandAtomMapping] + edges for this network + nodes : Iterable[SmallMoleculeComponent] + nodes for this network + """ + def __init__( + self, + edges: Iterable[LigandAtomMapping], + nodes: Optional[Iterable[SmallMoleculeComponent]] = None + ): + if nodes is None: + nodes = [] + + self._edges = frozenset(edges) + edge_nodes = set(chain.from_iterable((e.componentA, e.componentB) + for e in edges)) + self._nodes = frozenset(edge_nodes) | frozenset(nodes) + self._graph = None + + @classmethod + def _defaults(cls): + return {} + + def _to_dict(self) -> dict: + return {'graphml': self.to_graphml()} + + @classmethod + def _from_dict(cls, dct: dict): + return cls.from_graphml(dct['graphml']) + + @property + def graph(self) -> nx.MultiDiGraph: + """NetworkX graph for this network + + This graph will have :class:`.ChemicalSystem` objects as nodes and + :class:`.Transformation` objects as directed edges + """ + if self._graph is None: + graph = nx.MultiDiGraph() + # set iterator order depends on PYTHONHASHSEED, sorting ensures + # reproducibility + for node in sorted(self._nodes): + graph.add_node(node) + for edge in sorted(self._edges): + graph.add_edge(edge.componentA, edge.componentB, object=edge, + **edge.annotations) + + self._graph = nx.freeze(graph) + + return self._graph + + @property + def edges(self) -> FrozenSet[LigandAtomMapping]: + """A read-only view of the edges of the Network""" + return self._edges + + @property + def nodes(self) -> FrozenSet[SmallMoleculeComponent]: + """A read-only view of the nodes of the Network""" + return self._nodes + + def _serializable_graph(self) -> nx.Graph: + """ + Create NetworkX graph with serializable attribute representations. + + This enables us to use easily use different serialization + approaches. + """ + # sorting ensures that we always preserve order in files, so two + # identical networks will show no changes if you diff their + # serialized versions + sorted_nodes = sorted(self.nodes, key=lambda m: (m.smiles, m.name)) + mol_to_label = {mol: f"mol{num}" + for num, mol in enumerate(sorted_nodes)} + + edge_data = sorted([ + ( + mol_to_label[edge.componentA], + mol_to_label[edge.componentB], + json.dumps(list(edge.componentA_to_componentB.items())) + ) + for edge in self.edges + ]) + + # from here, we just build the graph + serializable_graph = nx.MultiDiGraph() + for mol, label in mol_to_label.items(): + serializable_graph.add_node(label, + moldict=json.dumps(mol.to_dict(), + sort_keys=True)) + + for molA, molB, mapping in edge_data: + serializable_graph.add_edge(molA, molB, mapping=mapping) + + return serializable_graph + + @classmethod + def _from_serializable_graph(cls, graph: nx.Graph): + """Create network from NetworkX graph with serializable attributes. + + This is the inverse of ``_serializable_graph``. + """ + label_to_mol = {node: SmallMoleculeComponent.from_dict(json.loads(d)) + for node, d in graph.nodes(data='moldict')} + + edges = [ + LigandAtomMapping(componentA=label_to_mol[node1], + componentB=label_to_mol[node2], + componentA_to_componentB=dict(json.loads(mapping))) + for node1, node2, mapping in graph.edges(data='mapping') + ] + + return cls(edges=edges, nodes=label_to_mol.values()) + + def to_graphml(self) -> str: + """Return the GraphML string representing this Network + + This is the primary serialization mechanism for this class. + + Returns + ------- + str : + string representing this network in GraphML format + """ + return "\n".join(nx.generate_graphml(self._serializable_graph())) + + @classmethod + def from_graphml(cls, graphml_str: str) -> LigandNetwork: + """Create from a GraphML string. + + Parameters + ---------- + graphml_str : str + GraphML string representation of a :class:`.Network` + + Returns + ------- + LigandNetwork + new network from the GraphML + """ + return cls._from_serializable_graph(nx.parse_graphml(graphml_str)) + + def enlarge_graph(self, *, edges=None, nodes=None) -> LigandNetwork: + """ + Create a new network with the given edges and nodes added + + Parameters + ---------- + edges : Iterable[:class:`.LigandAtomMapping`] + edges to append to this network + nodes : Iterable[:class:`.SmallMoleculeComponent`] + nodes to append to this network + + Returns + ------- + LigandNetwork + a new network adding the given edges and nodes to this network + """ + if edges is None: + edges = set([]) + + if nodes is None: + nodes = set([]) + + return LigandNetwork(self.edges | set(edges), self.nodes | set(nodes)) + + def _to_rfe_alchemical_network( + self, + components: dict[str, gufe.Component], + leg_labels: dict[str, list[str]], + protocol: gufe.Protocol, + *, + alchemical_label: str = "ligand", + autoname=True, + autoname_prefix="" + ) -> gufe.AlchemicalNetwork: + """ + Parameters + ---------- + components: dict[str, :class:`.Component`] + non-alchemical components (components that will be on both sides + of a transformation) + leg_labels: dict[str, list[str]] + mapping of the names for legs (the keys of this dict) to a list + of the component names. The component names must be the same as + used in the ``components`` dict. + protocol: :class:`.Protocol` + the protocol to apply + alchemical_label: str + the label for the component undergoing an alchemical + transformation (default ``'ligand'``) + """ + transformations = [] + for edge in self.edges: + for leg_name, labels in leg_labels.items(): + + # define a helper func to avoid repeated code + def sys_from_dict(component): + """ + Input component alchemically changing. Other info taken + from the outer scope. + """ + syscomps = {alchemical_label: component} + other_labels = set(labels) - {alchemical_label} + syscomps.update({label: components[label] + for label in other_labels}) + + if autoname: + name = f"{component.name}_{leg_name}" + else: + name = "" + + return gufe.ChemicalSystem(syscomps, name=name) + + sysA = sys_from_dict(edge.componentA) + sysB = sys_from_dict(edge.componentB) + if autoname: + prefix = f"{autoname_prefix}_" if autoname_prefix else "" + name = f"{prefix}{sysA.name}_{sysB.name}" + else: + name = "" + + transformation = gufe.Transformation(sysA, sysB, protocol, + mapping=edge, + name=name) + + transformations.append(transformation) + + return gufe.AlchemicalNetwork(transformations) + + def to_rbfe_alchemical_network( + self, + solvent: gufe.SolventComponent, + protein: gufe.ProteinComponent, + protocol: gufe.Protocol, + *, + autoname: bool = True, + autoname_prefix: str = "easy_rbfe", + **other_components + ) -> gufe.AlchemicalNetwork: + """Convert the ligand network to an AlchemicalNetwork + + Parameters + ---------- + protocol: Protocol + the method to apply to edges + autoname: bool + whether to automatically name objects by the ligand name and + state label + autoname_prefix: str + prefix for the autonaming; only used if autonaming is True + other_components: + additional non-alchemical components, keyword will be the string + label for the component + """ + components = { + 'protein': protein, + 'solvent': solvent, + **other_components + } + leg_labels = { + "solvent": ["ligand", "solvent"], + "complex": (["ligand", "solvent", "protein"] + + list(other_components)), + } + return self._to_rfe_alchemical_network( + components=components, + leg_labels=leg_labels, + protocol=protocol, + autoname=autoname, + autoname_prefix=autoname_prefix + ) + + # on hold until we figure out how to best hack in the PME/NoCutoff + # switch + # def to_rhfe_alchemical_network(self, *, solvent, protocol, + # autoname=True, + # autoname_prefix="easy_rhfe", + # **other_components): + # leg_labels = { + # "solvent": ["ligand", "solvent"] + list(other_components), + # "vacuum": ["ligand"] + list(other_components), + # } + # return self._to_rfe_alchemical_network( + # components={"solvent": solvent, **other_components}, + # leg_labels=leg_labels, + # protocol=protocol, + # autoname=autoname, + # autoname_prefix=autoname_prefix + # ) + + def is_connected(self) -> bool: + """Are all ligands in the network (indirectly) connected to each other + + A "False" value indicates that either some ligands have no edges or that + there are separate networks that do not link to each other. + """ + return nx.is_weakly_connected(self.graph) diff --git a/gufe/setup/network_planning/network_planner.py b/gufe/setup/network_planning/network_planner.py new file mode 100644 index 00000000..e69de29b diff --git a/gufe/transformations/transformation.py b/gufe/transformations/transformation.py index 58a898be..2352e873 100644 --- a/gufe/transformations/transformation.py +++ b/gufe/transformations/transformation.py @@ -10,7 +10,7 @@ from ..chemicalsystem import ChemicalSystem from ..protocols import Protocol, ProtocolDAG, ProtocolResult, ProtocolDAGResult -from ..mapping import ComponentMapping +from gufe.mapping import ComponentMapping class Transformation(GufeTokenizable): From c13f414c674f0867221c3641f65615eed9e593ab Mon Sep 17 00:00:00 2001 From: riesben Date: Wed, 11 Sep 2024 11:47:58 +0200 Subject: [PATCH 2/3] initial drafting --- gufe/__init__.py | 2 +- .../__init__.py | 0 .../atom_mapper.py | 0 .../atom_mapping.py | 0 .../atom_mapping_based/atom_mapping_scorer.py | 43 ++++++++++++++++++ .../ligand_atom_mapping.py | 0 .../ligandnetwork.py | 6 ++- .../network_planning/component_mapper.py | 30 +++++++++++++ .../component_mapping_scorer.py | 7 ++- .../setup/network_planning/network_planner.py | 44 +++++++++++++++++++ 10 files changed, 125 insertions(+), 7 deletions(-) rename gufe/setup/network_planning/{atom_based => atom_mapping_based}/__init__.py (100%) rename gufe/setup/network_planning/{atom_based => atom_mapping_based}/atom_mapper.py (100%) rename gufe/setup/network_planning/{atom_based => atom_mapping_based}/atom_mapping.py (100%) create mode 100644 gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py rename gufe/setup/network_planning/{atom_based => atom_mapping_based}/ligand_atom_mapping.py (100%) rename gufe/setup/network_planning/{atom_based => atom_mapping_based}/ligandnetwork.py (98%) diff --git a/gufe/__init__.py b/gufe/__init__.py index 891f059a..33459b6c 100644 --- a/gufe/__init__.py +++ b/gufe/__init__.py @@ -31,6 +31,6 @@ from .transformations import Transformation, NonTransformation from .alchemical_network import AlchemicalNetwork -from gufe.setup.network_planning.atom_based.ligandnetwork import LigandNetwork +from gufe.setup.network_planning.atom_mapping_based.ligandnetwork import LigandNetwork __version__ = version("gufe") diff --git a/gufe/setup/network_planning/atom_based/__init__.py b/gufe/setup/network_planning/atom_mapping_based/__init__.py similarity index 100% rename from gufe/setup/network_planning/atom_based/__init__.py rename to gufe/setup/network_planning/atom_mapping_based/__init__.py diff --git a/gufe/setup/network_planning/atom_based/atom_mapper.py b/gufe/setup/network_planning/atom_mapping_based/atom_mapper.py similarity index 100% rename from gufe/setup/network_planning/atom_based/atom_mapper.py rename to gufe/setup/network_planning/atom_mapping_based/atom_mapper.py diff --git a/gufe/setup/network_planning/atom_based/atom_mapping.py b/gufe/setup/network_planning/atom_mapping_based/atom_mapping.py similarity index 100% rename from gufe/setup/network_planning/atom_based/atom_mapping.py rename to gufe/setup/network_planning/atom_mapping_based/atom_mapping.py diff --git a/gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py b/gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py new file mode 100644 index 00000000..ef2795ee --- /dev/null +++ b/gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py @@ -0,0 +1,43 @@ +# This code is part of kartograf and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/gufe + +import abc +from ....tokenization import GufeTokenizable + +from .atom_mapping import AtomMapping + + +class AtomMappingScorer(GufeTokenizable): + """A generic class for scoring Atom mappings. + this class can be used for example to build graph algorithm based networks. + + Implementations of this class can require an arbitrary and non-standardised + number of input arguments to create. + + Implementations of this class provide the :meth:`.get_score` method + + """ + + def __call__(self, mapping: AtomMapping) -> float: + return self.get_score(mapping) + + @abc.abstractmethod + def get_score(self, mapping: AtomMapping) -> float: + """ calculate the score for an :class:`.AtomMapping` + the scoring function returns a value between 0 and 1. + a value close to 1.0 indicates a small change, a score close to zero indicates a large cost/change. + + Parameters + ---------- + mapping: AtomMapping + the mapping to be scored + args + kwargs + + Returns + ------- + float + a value between [0,1] where zero is a very bad score and one a very good one. + + """ + pass diff --git a/gufe/setup/network_planning/atom_based/ligand_atom_mapping.py b/gufe/setup/network_planning/atom_mapping_based/ligand_atom_mapping.py similarity index 100% rename from gufe/setup/network_planning/atom_based/ligand_atom_mapping.py rename to gufe/setup/network_planning/atom_mapping_based/ligand_atom_mapping.py diff --git a/gufe/setup/network_planning/atom_based/ligandnetwork.py b/gufe/setup/network_planning/atom_mapping_based/ligandnetwork.py similarity index 98% rename from gufe/setup/network_planning/atom_based/ligandnetwork.py rename to gufe/setup/network_planning/atom_mapping_based/ligandnetwork.py index 62751570..95fa2adc 100644 --- a/gufe/setup/network_planning/atom_based/ligandnetwork.py +++ b/gufe/setup/network_planning/atom_mapping_based/ligandnetwork.py @@ -9,11 +9,13 @@ import gufe from gufe import SmallMoleculeComponent -from gufe.mapping import LigandAtomMapping from gufe.tokenization import GufeTokenizable +from gufe.setup.network_planning.network_plan import NetworkPlan +from gufe.setup.network_planning.atom_mapping_based.ligand_atom_mapping import LigandAtomMapping -class LigandNetwork(GufeTokenizable): + +class LigandNetwork(NetworkPlan): """A directed graph connecting many ligands according to their atom mapping Parameters diff --git a/gufe/setup/network_planning/component_mapper.py b/gufe/setup/network_planning/component_mapper.py index e69de29b..98b887a7 100644 --- a/gufe/setup/network_planning/component_mapper.py +++ b/gufe/setup/network_planning/component_mapper.py @@ -0,0 +1,30 @@ +# This code is part of gufe and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/gufe +import abc +from collections.abc import Iterator +import gufe + +from gufe.tokenization import GufeTokenizable +from .component_mapping import ComponentMapping + + +class ComponentMapper(GufeTokenizable): + """A class for manufacturing mappings + + Implementations of this class can require an arbitrary and non-standardised + number of input arguments to create. + + Implementations of this class provide the :meth:`.suggest_mappings` method + """ + + @abc.abstractmethod + def suggest_mappings(self, + A: gufe.Component, + B: gufe.Component + ) -> Iterator[ComponentMapping]: + """Suggests possible mappings between two Components + + Suggests zero or more :class:`.AtomMapping` objects, which are possible + atom mappings between two :class:`.Component` objects. + """ + ... diff --git a/gufe/setup/network_planning/component_mapping_scorer.py b/gufe/setup/network_planning/component_mapping_scorer.py index a842f498..adc42741 100644 --- a/gufe/setup/network_planning/component_mapping_scorer.py +++ b/gufe/setup/network_planning/component_mapping_scorer.py @@ -4,8 +4,7 @@ import abc from gufe.tokenization import GufeTokenizable -from gufe.mapping import AtomMapping - +from .component_mapping import ComponentMapping class ComponentMappingScorer(GufeTokenizable): """A generic class for scoring Atom mappings. @@ -15,11 +14,11 @@ class ComponentMappingScorer(GufeTokenizable): Implementations of this class provide the :meth:`.get_score` method """ - def __call__(self, mapping: ComponentMappingScorer) -> float: + def __call__(self, mapping: ComponentMapping) -> float: return self.get_score(mapping) @abc.abstractmethod - def get_score(self, mapping: ComponentMappingScorer) -> float: + def get_score(self, mapping: ComponentMapping) -> float: """ calculate the score for an :class:`.AtomMapping` the scoring function returns a value between 0 and 1. a value close to 1.0 indicates a small change, a score close to zero indicates a large cost/change. diff --git a/gufe/setup/network_planning/network_planner.py b/gufe/setup/network_planning/network_planner.py index e69de29b..442ba568 100644 --- a/gufe/setup/network_planning/network_planner.py +++ b/gufe/setup/network_planning/network_planner.py @@ -0,0 +1,44 @@ +import abc +from typing import Iterable + +from ... import SmallMoleculeComponent + +from .network_plan import NetworkPlan +from .component_mapper import ComponentMapper +from .component_mapping_scorer import ComponentMappingScorer + +class NetworkPlanner(abc.ABC): + """A generic class for calculating :class:`.LigandNetworks`. + Implementations of this class can require an arbitrary and non-standardised + number of input arguments to create. + Implementations of this class provide the :meth:`.get_score` method + """ + + def __init__(self, mapper: ComponentMapper, scorer:ComponentMappingScorer, *args, *kwargs): + """ Generate a Ligand Network Planner. This class in general needs a mapper and a scorer. + Parameters + ---------- + mapper: AtomMapper + scorer: AtomMappingScorer + args + kwargs + """ + self.mapper = mapper + self.scorer = scorer + + + def __call__(self, ligands: Iterable[SmallMoleculeComponent])-> NetworkPlan: + return self.generate_ligand_network(*args, **kwargs) + + @abc.abstractmethod + def generate_ligand_network(self, ligands: Iterable[SmallMoleculeComponent])->LigandNetwork: + """Plan a Network which connects all ligands with minimal cost + Parameters + ---------- + ligands : Iterable[SmallMoleculeComponent] + the ligands to include in the Network + Returns + ------- + LigandNetwork + A Network, that connects all ligands with each other. + """ \ No newline at end of file From 2a243c8869216bc0646764ae3057e9594b4e194a Mon Sep 17 00:00:00 2001 From: riesben Date: Wed, 18 Sep 2024 10:20:15 +0200 Subject: [PATCH 3/3] minor improvements to the draft. there are few discussion points left. --- .../atom_mapping_based/atom_mapper.py | 3 +- .../atom_mapping_based/atom_mapping.py | 9 +- .../atom_mapping_based/atom_mapping_scorer.py | 4 +- .../atom_mapping_based/ligandnetwork.py | 41 +---- .../network_planning/component_mapper.py | 2 +- .../network_planning/component_mapping.py | 1 + .../component_mapping_scorer.py | 2 +- gufe/setup/network_planning/network_plan.py | 159 ++---------------- .../setup/network_planning/network_planner.py | 13 +- 9 files changed, 37 insertions(+), 197 deletions(-) diff --git a/gufe/setup/network_planning/atom_mapping_based/atom_mapper.py b/gufe/setup/network_planning/atom_mapping_based/atom_mapper.py index f33ef7c3..cbdb7f8c 100644 --- a/gufe/setup/network_planning/atom_mapping_based/atom_mapper.py +++ b/gufe/setup/network_planning/atom_mapping_based/atom_mapper.py @@ -27,4 +27,5 @@ def suggest_mappings(self, Suggests zero or more :class:`.AtomMapping` objects, which are possible atom mappings between two :class:`.Component` objects. """ - ... + raise NotImplementedError("This function was not implemented.") + diff --git a/gufe/setup/network_planning/atom_mapping_based/atom_mapping.py b/gufe/setup/network_planning/atom_mapping_based/atom_mapping.py index 03cb0ac5..d66a3603 100644 --- a/gufe/setup/network_planning/atom_mapping_based/atom_mapping.py +++ b/gufe/setup/network_planning/atom_mapping_based/atom_mapping.py @@ -37,22 +37,23 @@ def componentA_to_componentB(self) -> Mapping[int, int]: entity in the other component (e.g. the atom disappears), therefore resulting in a KeyError on query """ - ... + raise NotImplementedError("This function was not implemented.") @property @abc.abstractmethod def componentB_to_componentA(self) -> Mapping[int, int]: """Similar to A to B, but reversed.""" - ... + raise NotImplementedError("This function was not implemented.") @property @abc.abstractmethod def componentA_unique(self) -> Iterable[int]: """Indices of atoms in component A that aren't mappable to B""" - ... + raise NotImplementedError("This function was not implemented.") @property @abc.abstractmethod def componentB_unique(self) -> Iterable[int]: """Indices of atoms in component B that aren't mappable to A""" - ... + raise NotImplementedError("This function was not implemented.") + diff --git a/gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py b/gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py index ef2795ee..b533ede3 100644 --- a/gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py +++ b/gufe/setup/network_planning/atom_mapping_based/atom_mapping_scorer.py @@ -25,7 +25,7 @@ def __call__(self, mapping: AtomMapping) -> float: def get_score(self, mapping: AtomMapping) -> float: """ calculate the score for an :class:`.AtomMapping` the scoring function returns a value between 0 and 1. - a value close to 1.0 indicates a small change, a score close to zero indicates a large cost/change. + a value close to 1.0 indicates a small change - good score, a score close to zero indicates a large cost/change - bad score. Parameters ---------- @@ -40,4 +40,4 @@ def get_score(self, mapping: AtomMapping) -> float: a value between [0,1] where zero is a very bad score and one a very good one. """ - pass + raise NotImplementedError("This function was not implemented.") diff --git a/gufe/setup/network_planning/atom_mapping_based/ligandnetwork.py b/gufe/setup/network_planning/atom_mapping_based/ligandnetwork.py index 95fa2adc..4fc4fa2d 100644 --- a/gufe/setup/network_planning/atom_mapping_based/ligandnetwork.py +++ b/gufe/setup/network_planning/atom_mapping_based/ligandnetwork.py @@ -50,37 +50,6 @@ def _to_dict(self) -> dict: def _from_dict(cls, dct: dict): return cls.from_graphml(dct['graphml']) - @property - def graph(self) -> nx.MultiDiGraph: - """NetworkX graph for this network - - This graph will have :class:`.ChemicalSystem` objects as nodes and - :class:`.Transformation` objects as directed edges - """ - if self._graph is None: - graph = nx.MultiDiGraph() - # set iterator order depends on PYTHONHASHSEED, sorting ensures - # reproducibility - for node in sorted(self._nodes): - graph.add_node(node) - for edge in sorted(self._edges): - graph.add_edge(edge.componentA, edge.componentB, object=edge, - **edge.annotations) - - self._graph = nx.freeze(graph) - - return self._graph - - @property - def edges(self) -> FrozenSet[LigandAtomMapping]: - """A read-only view of the edges of the Network""" - return self._edges - - @property - def nodes(self) -> FrozenSet[SmallMoleculeComponent]: - """A read-only view of the nodes of the Network""" - return self._nodes - def _serializable_graph(self) -> nx.Graph: """ Create NetworkX graph with serializable attribute representations. @@ -309,12 +278,4 @@ def to_rbfe_alchemical_network( # protocol=protocol, # autoname=autoname, # autoname_prefix=autoname_prefix - # ) - - def is_connected(self) -> bool: - """Are all ligands in the network (indirectly) connected to each other - - A "False" value indicates that either some ligands have no edges or that - there are separate networks that do not link to each other. - """ - return nx.is_weakly_connected(self.graph) + # ) \ No newline at end of file diff --git a/gufe/setup/network_planning/component_mapper.py b/gufe/setup/network_planning/component_mapper.py index 98b887a7..86134bd0 100644 --- a/gufe/setup/network_planning/component_mapper.py +++ b/gufe/setup/network_planning/component_mapper.py @@ -27,4 +27,4 @@ def suggest_mappings(self, Suggests zero or more :class:`.AtomMapping` objects, which are possible atom mappings between two :class:`.Component` objects. """ - ... + raise NotImplementedError("This function was not implemented.") diff --git a/gufe/setup/network_planning/component_mapping.py b/gufe/setup/network_planning/component_mapping.py index 92a098f3..b5506cdf 100644 --- a/gufe/setup/network_planning/component_mapping.py +++ b/gufe/setup/network_planning/component_mapping.py @@ -17,6 +17,7 @@ class ComponentMapping(GufeTokenizable, abc.ABC): def __init__(self, componentA: gufe.Component, componentB: gufe.Component): self._componentA = componentA self._componentB = componentB + # self.componentA_to_componentB # TODO: is that something we want here, thinking beyond AtomMappings? def __contains__(self, item: gufe.Component): return item == self._componentA or item == self._componentB diff --git a/gufe/setup/network_planning/component_mapping_scorer.py b/gufe/setup/network_planning/component_mapping_scorer.py index adc42741..86d772c4 100644 --- a/gufe/setup/network_planning/component_mapping_scorer.py +++ b/gufe/setup/network_planning/component_mapping_scorer.py @@ -33,4 +33,4 @@ def get_score(self, mapping: ComponentMapping) -> float: float a value between [0,1] where zero is a very bad score and one a very good one. """ - pass \ No newline at end of file + raise NotImplementedError("This function was not implemented.") diff --git a/gufe/setup/network_planning/network_plan.py b/gufe/setup/network_planning/network_plan.py index d60e84fe..4708a441 100644 --- a/gufe/setup/network_planning/network_plan.py +++ b/gufe/setup/network_planning/network_plan.py @@ -8,25 +8,25 @@ from typing import FrozenSet, Iterable, Optional import gufe -from gufe import SmallMoleculeComponent -from gufe.mapping import LigandAtomMapping +from gufe import Component from gufe.tokenization import GufeTokenizable +from .component_mapping import ComponentMapping class NetworkPlan(GufeTokenizable): """A directed graph connecting many ligands according to their atom mapping Parameters ---------- - edges : Iterable[LigandAtomMapping] + edges : Iterable[ComponentMapping] edges for this network nodes : Iterable[SmallMoleculeComponent] nodes for this network """ def __init__( self, - edges: Iterable[LigandAtomMapping], - nodes: Optional[Iterable[SmallMoleculeComponent]] = None + edges: Iterable[ComponentMapping], + nodes: Optional[Iterable[Component]] = None ): if nodes is None: nodes = [] @@ -70,12 +70,12 @@ def graph(self) -> nx.MultiDiGraph: return self._graph @property - def edges(self) -> FrozenSet[LigandAtomMapping]: + def edges(self) -> FrozenSet[ComponentMapping]: """A read-only view of the edges of the Network""" return self._edges @property - def nodes(self) -> FrozenSet[SmallMoleculeComponent]: + def nodes(self) -> FrozenSet[Component]: """A read-only view of the nodes of the Network""" return self._nodes @@ -97,7 +97,6 @@ def _serializable_graph(self) -> nx.Graph: ( mol_to_label[edge.componentA], mol_to_label[edge.componentB], - json.dumps(list(edge.componentA_to_componentB.items())) ) for edge in self.edges ]) @@ -109,8 +108,8 @@ def _serializable_graph(self) -> nx.Graph: moldict=json.dumps(mol.to_dict(), sort_keys=True)) - for molA, molB, mapping in edge_data: - serializable_graph.add_edge(molA, molB, mapping=mapping) + for molA, molB in edge_data: + serializable_graph.add_edge(molA, molB) return serializable_graph @@ -120,14 +119,13 @@ def _from_serializable_graph(cls, graph: nx.Graph): This is the inverse of ``_serializable_graph``. """ - label_to_mol = {node: SmallMoleculeComponent.from_dict(json.loads(d)) + label_to_mol = {node: Component.from_dict(json.loads(d)) for node, d in graph.nodes(data='moldict')} edges = [ - LigandAtomMapping(componentA=label_to_mol[node1], - componentB=label_to_mol[node2], - componentA_to_componentB=dict(json.loads(mapping))) - for node1, node2, mapping in graph.edges(data='mapping') + ComponentMapping(componentA=label_to_mol[node1], + componentB=label_to_mol[node2]) + for node1, node2 in graph.edges(data='mapping') ] return cls(edges=edges, nodes=label_to_mol.values()) @@ -145,7 +143,7 @@ def to_graphml(self) -> str: return "\n".join(nx.generate_graphml(self._serializable_graph())) @classmethod - def from_graphml(cls, graphml_str: str) -> LigandNetwork: + def from_graphml(cls, graphml_str: str) -> ComponentMapping: """Create from a GraphML string. Parameters @@ -160,7 +158,7 @@ def from_graphml(cls, graphml_str: str) -> LigandNetwork: """ return cls._from_serializable_graph(nx.parse_graphml(graphml_str)) - def enlarge_graph(self, *, edges=None, nodes=None) -> LigandNetwork: + def enlarge_graph(self, *, edges=None, nodes=None) -> ComponentMapping: """ Create a new network with the given edges and nodes added @@ -182,132 +180,7 @@ def enlarge_graph(self, *, edges=None, nodes=None) -> LigandNetwork: if nodes is None: nodes = set([]) - return LigandNetwork(self.edges | set(edges), self.nodes | set(nodes)) - - def _to_rfe_alchemical_network( - self, - components: dict[str, gufe.Component], - leg_labels: dict[str, list[str]], - protocol: gufe.Protocol, - *, - alchemical_label: str = "ligand", - autoname=True, - autoname_prefix="" - ) -> gufe.AlchemicalNetwork: - """ - Parameters - ---------- - components: dict[str, :class:`.Component`] - non-alchemical components (components that will be on both sides - of a transformation) - leg_labels: dict[str, list[str]] - mapping of the names for legs (the keys of this dict) to a list - of the component names. The component names must be the same as - used in the ``components`` dict. - protocol: :class:`.Protocol` - the protocol to apply - alchemical_label: str - the label for the component undergoing an alchemical - transformation (default ``'ligand'``) - """ - transformations = [] - for edge in self.edges: - for leg_name, labels in leg_labels.items(): - - # define a helper func to avoid repeated code - def sys_from_dict(component): - """ - Input component alchemically changing. Other info taken - from the outer scope. - """ - syscomps = {alchemical_label: component} - other_labels = set(labels) - {alchemical_label} - syscomps.update({label: components[label] - for label in other_labels}) - - if autoname: - name = f"{component.name}_{leg_name}" - else: - name = "" - - return gufe.ChemicalSystem(syscomps, name=name) - - sysA = sys_from_dict(edge.componentA) - sysB = sys_from_dict(edge.componentB) - if autoname: - prefix = f"{autoname_prefix}_" if autoname_prefix else "" - name = f"{prefix}{sysA.name}_{sysB.name}" - else: - name = "" - - transformation = gufe.Transformation(sysA, sysB, protocol, - mapping=edge, - name=name) - - transformations.append(transformation) - - return gufe.AlchemicalNetwork(transformations) - - def to_rbfe_alchemical_network( - self, - solvent: gufe.SolventComponent, - protein: gufe.ProteinComponent, - protocol: gufe.Protocol, - *, - autoname: bool = True, - autoname_prefix: str = "easy_rbfe", - **other_components - ) -> gufe.AlchemicalNetwork: - """Convert the ligand network to an AlchemicalNetwork - - Parameters - ---------- - protocol: Protocol - the method to apply to edges - autoname: bool - whether to automatically name objects by the ligand name and - state label - autoname_prefix: str - prefix for the autonaming; only used if autonaming is True - other_components: - additional non-alchemical components, keyword will be the string - label for the component - """ - components = { - 'protein': protein, - 'solvent': solvent, - **other_components - } - leg_labels = { - "solvent": ["ligand", "solvent"], - "complex": (["ligand", "solvent", "protein"] - + list(other_components)), - } - return self._to_rfe_alchemical_network( - components=components, - leg_labels=leg_labels, - protocol=protocol, - autoname=autoname, - autoname_prefix=autoname_prefix - ) - - # on hold until we figure out how to best hack in the PME/NoCutoff - # switch - # def to_rhfe_alchemical_network(self, *, solvent, protocol, - # autoname=True, - # autoname_prefix="easy_rhfe", - # **other_components): - # leg_labels = { - # "solvent": ["ligand", "solvent"] + list(other_components), - # "vacuum": ["ligand"] + list(other_components), - # } - # return self._to_rfe_alchemical_network( - # components={"solvent": solvent, **other_components}, - # leg_labels=leg_labels, - # protocol=protocol, - # autoname=autoname, - # autoname_prefix=autoname_prefix - # ) + return NetworkPlan(self.edges | set(edges), self.nodes | set(nodes)) def is_connected(self) -> bool: """Are all ligands in the network (indirectly) connected to each other diff --git a/gufe/setup/network_planning/network_planner.py b/gufe/setup/network_planning/network_planner.py index 442ba568..aa25ba32 100644 --- a/gufe/setup/network_planning/network_planner.py +++ b/gufe/setup/network_planning/network_planner.py @@ -3,11 +3,13 @@ from ... import SmallMoleculeComponent +from gufe.tokenization import GufeTokenizable + from .network_plan import NetworkPlan from .component_mapper import ComponentMapper from .component_mapping_scorer import ComponentMappingScorer -class NetworkPlanner(abc.ABC): +class NetworkPlanner(GufeTokenizable): """A generic class for calculating :class:`.LigandNetworks`. Implementations of this class can require an arbitrary and non-standardised number of input arguments to create. @@ -27,11 +29,11 @@ def __init__(self, mapper: ComponentMapper, scorer:ComponentMappingScorer, *args self.scorer = scorer - def __call__(self, ligands: Iterable[SmallMoleculeComponent])-> NetworkPlan: + def __call__(self, *args, **kwargs)-> NetworkPlan: return self.generate_ligand_network(*args, **kwargs) @abc.abstractmethod - def generate_ligand_network(self, ligands: Iterable[SmallMoleculeComponent])->LigandNetwork: + def generate_ligand_network(self, ligands: Iterable[SmallMoleculeComponent])->NetworkPlan: """Plan a Network which connects all ligands with minimal cost Parameters ---------- @@ -39,6 +41,7 @@ def generate_ligand_network(self, ligands: Iterable[SmallMoleculeComponent])->Li the ligands to include in the Network Returns ------- - LigandNetwork + NetworkPlan A Network, that connects all ligands with each other. - """ \ No newline at end of file + """ + raise NotImplementedError("This function was not implemented.")