From f0efdbe41d9067848a674c649ce3bd30d9ea8fff Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 14 Nov 2024 10:58:42 -0800 Subject: [PATCH 1/8] adding subgraph method --- gufe/network.py | 19 +++++++++++++++++-- gufe/tests/test_alchemicalnetwork.py | 15 +++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/gufe/network.py b/gufe/network.py index e02be31a..adc2a0ee 100644 --- a/gufe/network.py +++ b/gufe/network.py @@ -1,7 +1,7 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/gufe -from typing import Iterable, Optional +from typing import Generator, Iterable, Optional, Self import networkx as nx from .tokenization import GufeTokenizable @@ -102,7 +102,7 @@ def _to_dict(self) -> dict: "name": self.name} @classmethod - def _from_dict(cls, d: dict): + def _from_dict(cls, d: dict) -> Self: return cls(nodes=frozenset(d['nodes']), edges=frozenset(d['edges']), name=d.get('name')) @@ -115,6 +115,21 @@ def to_graphml(self) -> str: """Currently not implemented""" raise NotImplementedError + @classmethod + def _from_nx_graph(cls, nx_graph) -> Self: + chemical_systems = [n for n in nx_graph.nodes()] + tranformations = [e[2]['object'] for e in nx_graph.edges(data=True)] # list of transformations + return cls(nodes=chemical_systems, edges=tranformations) + + def connected_subgraphs(self) -> Generator[Self, None, None]: + """Return all connected subgraphs of the alchemical network in order of size + """ + node_groups = nx.weakly_connected_components(self.graph) + for node_group in node_groups: + nx_subgraph = self.graph.subgraph(node_group) + alc_subgraph = self._from_nx_graph(nx_subgraph) + yield(alc_subgraph) + @classmethod def from_graphml(cls, str): """Currently not implemented""" diff --git a/gufe/tests/test_alchemicalnetwork.py b/gufe/tests/test_alchemicalnetwork.py index ef148bac..438fa412 100644 --- a/gufe/tests/test_alchemicalnetwork.py +++ b/gufe/tests/test_alchemicalnetwork.py @@ -47,3 +47,18 @@ def test_connectivity(self, benzene_variants_star_map): else: edges = alnet.graph.edges(node) assert len(edges) == 0 + + def test_weakly_connected_subgraphs(self, benzene_variants_star_map): + # remove two edges to create a network w/ two floating nodes + edge_list = [e for e in benzene_variants_star_map.edges] + alnet = benzene_variants_star_map.copy_with_replacements(edges=edge_list[:-2]) + + subgraphs = [subgraph for subgraph in alnet.connected_subgraphs()] + + # check that one graph is all protein components + # and the other is all solvent components + + assert len(subgraphs) == 4 + + assert 0 + From 348ef73c728839027d299f901eca9876b9663502 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 14 Nov 2024 11:39:34 -0800 Subject: [PATCH 2/8] adding check for correct network subgraph generation --- gufe/network.py | 2 +- gufe/tests/test_alchemicalnetwork.py | 16 ++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/gufe/network.py b/gufe/network.py index adc2a0ee..65cf9baa 100644 --- a/gufe/network.py +++ b/gufe/network.py @@ -122,7 +122,7 @@ def _from_nx_graph(cls, nx_graph) -> Self: return cls(nodes=chemical_systems, edges=tranformations) def connected_subgraphs(self) -> Generator[Self, None, None]: - """Return all connected subgraphs of the alchemical network in order of size + """Return all connected subgraphs of the alchemical network in order of size (largest first). """ node_groups = nx.weakly_connected_components(self.graph) for node_group in node_groups: diff --git a/gufe/tests/test_alchemicalnetwork.py b/gufe/tests/test_alchemicalnetwork.py index 438fa412..7418ff07 100644 --- a/gufe/tests/test_alchemicalnetwork.py +++ b/gufe/tests/test_alchemicalnetwork.py @@ -51,14 +51,18 @@ def test_connectivity(self, benzene_variants_star_map): def test_weakly_connected_subgraphs(self, benzene_variants_star_map): # remove two edges to create a network w/ two floating nodes edge_list = [e for e in benzene_variants_star_map.edges] - alnet = benzene_variants_star_map.copy_with_replacements(edges=edge_list[:-2]) + alnet = benzene_variants_star_map.copy_with_replacements(edges=edge_list[:-1]) subgraphs = [subgraph for subgraph in alnet.connected_subgraphs()] - # check that one graph is all protein components - # and the other is all solvent components + assert set([len(subgraph.nodes) for subgraph in subgraphs]) == {6,7,1} - assert len(subgraphs) == 4 + # which graph has the removed node is not deterministic, so we just + # check that one graph is all-solvent and the other is all-protein + for subgraph in subgraphs: + components = [frozenset(n.components.keys()) for n in subgraph.nodes] + if {'solvent','protein','ligand'} in components: + assert set(components) == {frozenset({'solvent','protein','ligand'})} - assert 0 - + else: + assert set(components) == {frozenset({'solvent','ligand'})} From ea09e235374ab13684f323e66e66ff47dd26ad57 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 14 Nov 2024 11:48:26 -0800 Subject: [PATCH 3/8] adding try/except for self typing import --- gufe/network.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/gufe/network.py b/gufe/network.py index 65cf9baa..d3f66854 100644 --- a/gufe/network.py +++ b/gufe/network.py @@ -1,7 +1,7 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/gufe -from typing import Generator, Iterable, Optional, Self +from typing import Generator, Iterable, Optional import networkx as nx from .tokenization import GufeTokenizable @@ -9,6 +9,11 @@ from .chemicalsystem import ChemicalSystem from .transformations import Transformation +# Self typing is supported in python 3.11+ +try: + from typing import Self +except ImportError: + from __future__ import annotations class AlchemicalNetwork(GufeTokenizable): _edges: frozenset[Transformation] From e568ccf0d05f69a06dd69b09cfb517b2f7f3eb44 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 14 Nov 2024 11:49:28 -0800 Subject: [PATCH 4/8] putting functions in a nicer order --- gufe/network.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/gufe/network.py b/gufe/network.py index d3f66854..20f8f740 100644 --- a/gufe/network.py +++ b/gufe/network.py @@ -120,22 +120,22 @@ def to_graphml(self) -> str: """Currently not implemented""" raise NotImplementedError + @classmethod + def from_graphml(cls, str): + """Currently not implemented""" + raise NotImplementedError + @classmethod def _from_nx_graph(cls, nx_graph) -> Self: + """Create an alchemical network from a networkx representation.""" chemical_systems = [n for n in nx_graph.nodes()] - tranformations = [e[2]['object'] for e in nx_graph.edges(data=True)] # list of transformations - return cls(nodes=chemical_systems, edges=tranformations) + transformations = [e[2]['object'] for e in nx_graph.edges(data=True)] + return cls(nodes=chemical_systems, edges=transformations) def connected_subgraphs(self) -> Generator[Self, None, None]: - """Return all connected subgraphs of the alchemical network in order of size (largest first). - """ + """Return a generator of all connected subgraphs of the alchemical network.""" node_groups = nx.weakly_connected_components(self.graph) for node_group in node_groups: nx_subgraph = self.graph.subgraph(node_group) alc_subgraph = self._from_nx_graph(nx_subgraph) yield(alc_subgraph) - - @classmethod - def from_graphml(cls, str): - """Currently not implemented""" - raise NotImplementedError From aff2dfbd1a9c9cb9deffc554eb1cb7d2788ad86f Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 14 Nov 2024 11:54:00 -0800 Subject: [PATCH 5/8] fixing typing import --- gufe/network.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/gufe/network.py b/gufe/network.py index 20f8f740..e2b21ec8 100644 --- a/gufe/network.py +++ b/gufe/network.py @@ -2,6 +2,7 @@ # For details, see https://github.com/OpenFreeEnergy/gufe from typing import Generator, Iterable, Optional +from typing_extensions import Self # Self is included in typing as of python 3.11 import networkx as nx from .tokenization import GufeTokenizable @@ -9,11 +10,7 @@ from .chemicalsystem import ChemicalSystem from .transformations import Transformation -# Self typing is supported in python 3.11+ -try: - from typing import Self -except ImportError: - from __future__ import annotations + class AlchemicalNetwork(GufeTokenizable): _edges: frozenset[Transformation] From f0a834605ab63cb3647a6f24a590f7f239d25c35 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 14 Nov 2024 12:26:53 -0800 Subject: [PATCH 6/8] updating test name to match function name --- gufe/tests/test_alchemicalnetwork.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gufe/tests/test_alchemicalnetwork.py b/gufe/tests/test_alchemicalnetwork.py index 7418ff07..06f64e68 100644 --- a/gufe/tests/test_alchemicalnetwork.py +++ b/gufe/tests/test_alchemicalnetwork.py @@ -48,7 +48,7 @@ def test_connectivity(self, benzene_variants_star_map): edges = alnet.graph.edges(node) assert len(edges) == 0 - def test_weakly_connected_subgraphs(self, benzene_variants_star_map): + def test_connected_subgraphs(self, benzene_variants_star_map): # remove two edges to create a network w/ two floating nodes edge_list = [e for e in benzene_variants_star_map.edges] alnet = benzene_variants_star_map.copy_with_replacements(edges=edge_list[:-1]) From 0c982ee2f6344e09f3e8241a52acbc741cb0f522 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Mon, 18 Nov 2024 12:59:46 -0800 Subject: [PATCH 7/8] adding connected network fixture and additional test --- gufe/tests/conftest.py | 19 +++++++++++++------ gufe/tests/test_alchemicalnetwork.py | 12 +++++++++--- 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/gufe/tests/conftest.py b/gufe/tests/conftest.py index 62487740..3080f469 100644 --- a/gufe/tests/conftest.py +++ b/gufe/tests/conftest.py @@ -256,9 +256,8 @@ def complex_equilibrium(solvated_complex): protocol=DummyProtocol(settings=DummyProtocol.default_settings()) ) - @pytest.fixture -def benzene_variants_star_map( +def benzene_variants_star_map_transformations( benzene, toluene, phenol, @@ -320,7 +319,15 @@ def benzene_variants_star_map( mapping=None, ) - return gufe.AlchemicalNetwork( - list(solvated_ligand_transformations.values()) - + list(solvated_complex_transformations.values()) - ) + return list(solvated_ligand_transformations.values()), list(solvated_complex_transformations.values()) + + +@pytest.fixture +def benzene_variants_star_map(benzene_variants_star_map_transformations): + solvated_ligand_transformations, solvated_complex_transformations = benzene_variants_star_map_transformations + return gufe.AlchemicalNetwork(solvated_ligand_transformations+solvated_complex_transformations) + +@pytest.fixture +def benzene_variants_ligand_star_map(benzene_variants_star_map_transformations): + solvated_ligand_transformations, _ = benzene_variants_star_map_transformations + return gufe.AlchemicalNetwork(solvated_ligand_transformations) diff --git a/gufe/tests/test_alchemicalnetwork.py b/gufe/tests/test_alchemicalnetwork.py index 06f64e68..8231c51d 100644 --- a/gufe/tests/test_alchemicalnetwork.py +++ b/gufe/tests/test_alchemicalnetwork.py @@ -48,8 +48,9 @@ def test_connectivity(self, benzene_variants_star_map): edges = alnet.graph.edges(node) assert len(edges) == 0 - def test_connected_subgraphs(self, benzene_variants_star_map): - # remove two edges to create a network w/ two floating nodes + def test_connected_subgraphs_multiple_subgraphs(self, benzene_variants_star_map): + """Identify two separate networks and one floating nodes as subgraphs.""" + # remove an edge to create a network w/ two subnetworks and one floating node edge_list = [e for e in benzene_variants_star_map.edges] alnet = benzene_variants_star_map.copy_with_replacements(edges=edge_list[:-1]) @@ -63,6 +64,11 @@ def test_connected_subgraphs(self, benzene_variants_star_map): components = [frozenset(n.components.keys()) for n in subgraph.nodes] if {'solvent','protein','ligand'} in components: assert set(components) == {frozenset({'solvent','protein','ligand'})} - else: assert set(components) == {frozenset({'solvent','ligand'})} + + def test_connected_subgraphs_one_subgraph(self, benzene_variants_ligand_star_map): + """Return the same network if it only contains one connected component.""" + alnet = benzene_variants_ligand_star_map + subgraphs = [subgraph for subgraph in alnet.connected_subgraphs()] + assert subgraphs == [alnet] From ba85863932a693374c316e16f8ee00f9b8b019fe Mon Sep 17 00:00:00 2001 From: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> Date: Tue, 19 Nov 2024 10:05:37 -0800 Subject: [PATCH 8/8] Apply suggestions from code review Co-authored-by: Irfan Alibay --- gufe/network.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gufe/network.py b/gufe/network.py index e2b21ec8..59929d68 100644 --- a/gufe/network.py +++ b/gufe/network.py @@ -118,7 +118,7 @@ def to_graphml(self) -> str: raise NotImplementedError @classmethod - def from_graphml(cls, str): + def from_graphml(cls, str) -> Self: """Currently not implemented""" raise NotImplementedError