From 1a33b8cae7a7e01cb3d6eba7f1ac8ff51dc67c81 Mon Sep 17 00:00:00 2001
From: richard gowers <richardjgowers@gmail.com>
Date: Tue, 7 May 2024 15:46:20 +0100
Subject: [PATCH] adds LigandNetwork.remove_edges

---
 gufe/ligandnetwork.py                     | 32 ++++++++++++++++++-
 gufe/tests/test_ligand_network.py         | 38 +++++++++++++++++++++++
 news/added_LigandNetwork_delete_edges.rst | 23 ++++++++++++++
 3 files changed, 92 insertions(+), 1 deletion(-)
 create mode 100644 news/added_LigandNetwork_delete_edges.rst

diff --git a/gufe/ligandnetwork.py b/gufe/ligandnetwork.py
index 177aa61e..5e10a35c 100644
--- a/gufe/ligandnetwork.py
+++ b/gufe/ligandnetwork.py
@@ -5,7 +5,7 @@
 from itertools import chain
 import json
 import networkx as nx
-from typing import FrozenSet, Iterable, Optional
+from typing import FrozenSet, Iterable, Optional, Union
 import gufe
 
 from gufe import SmallMoleculeComponent
@@ -184,6 +184,36 @@ def enlarge_graph(self, *, edges=None, nodes=None) -> LigandNetwork:
 
         return LigandNetwork(self.edges | set(edges), self.nodes | set(nodes))
 
+    def remove_edges(self, edges: Union[LigandAtomMapping, list[LigandAtomMapping]]) -> LigandNetwork:
+        """Create a new copy of this network with some edges removed
+
+        Note that this will not remove any nodes, potentially resulting in
+        disconnected networks
+
+        Parameters
+        ----------
+        edges : list[LigandAtomMapping] or LigandAtomMapping
+          the edges to remove, these *must* be present in the network
+
+        Returns
+        -------
+        network : LigandNetwork
+        """
+        if isinstance(edges, LigandAtomMapping):
+            edges = [edges]
+
+        to_remove = set(edges)
+        current = set(self.edges)
+
+        # check that all edges to remove are present
+        if extras := to_remove - current:
+            raise ValueError("Some edges weren't already present: "
+                             f"{extras}")
+
+        new_edges = current - to_remove
+
+        return LigandNetwork(new_edges, self.nodes)
+
     def _to_rfe_alchemical_network(
         self,
         components: dict[str, gufe.Component],
diff --git a/gufe/tests/test_ligand_network.py b/gufe/tests/test_ligand_network.py
index a81a0888..e394aa66 100644
--- a/gufe/tests/test_ligand_network.py
+++ b/gufe/tests/test_ligand_network.py
@@ -440,3 +440,41 @@ def test_empty_ligand_network(mols):
 
     assert len(n.edges) == 0
     assert len(n.nodes) == 1
+
+
+class TestLigandNetworkRemoveEdges:
+    def test_remove_edges_single(self, std_edges):
+        e1, e2, e3 = std_edges
+
+        n = LigandNetwork(edges=[e1, e2, e3])
+
+        n2 = n.remove_edges(e1)
+
+        assert len(n2.edges) == 2
+        assert len(n2.nodes) == 3
+        assert e1 not in n2.edges
+        assert e2 in n2.edges
+        assert e3 in n2.edges
+        assert n.nodes == n2.nodes
+
+    def test_remove_edges_pair(self, std_edges):
+        e1, e2, e3 = std_edges
+
+        n = LigandNetwork(edges=[e1, e2, e3])
+
+        n2 = n.remove_edges([e1, e2])
+
+        assert len(n2.edges) == 1
+        assert len(n2.nodes) == 3
+        assert e1 not in n2.edges
+        assert e2 not in n2.edges
+        assert e3 in n2.edges
+        assert n.nodes == n2.nodes
+
+    def test_remove_edges_fail(self, std_edges):
+        e1, e2, e3 = std_edges
+
+        n = LigandNetwork([e1, e2])
+
+        with pytest.raises(ValueError):
+            n.remove_edges(e3)
diff --git a/news/added_LigandNetwork_delete_edges.rst b/news/added_LigandNetwork_delete_edges.rst
new file mode 100644
index 00000000..e0c35c69
--- /dev/null
+++ b/news/added_LigandNetwork_delete_edges.rst
@@ -0,0 +1,23 @@
+**Added:**
+
+* added LigandNetwork.remove_edges
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>