Merge pull request #59 from YosefLab/codex/add-treedata._has_overlap-attribute

Track tree overlap state

Author: colganwi

src/treedata/_core/aligned_mapping.py

@@ -184,6 +184,7 @@ def __setitem__(self, key: str, value: nx.DiGraph):
             self._update_tree_labels()
 
         self._data[key] = value.copy()
+        self.parent._update_has_overlap()
 
     def __delitem__(self, key: str):
         """Delete item from the mapping."""
@@ -196,6 +197,7 @@ def __delitem__(self, key: str):
         self._update_tree_labels()
 
         del self._data[key]
+        self.parent._update_has_overlap()
 
     def __len__(self) -> int:
         """Get length of the mapping."""

src/treedata/_core/treedata.py

@@ -184,6 +184,7 @@ def _init_as_actual(
             self._obst = X.obst
             self._vart = X.vart
             self._alignment = X._alignment
+            self._has_overlap = X.has_overlap
 
         # init from scratch
         else:
@@ -199,8 +200,10 @@ def _init_as_actual(
                 self._allow_overlap = bool(allow_overlap)
             else:
                 raise ValueError("allow_overlap has to be a boolean")
+            self._has_overlap = False
             self._obst = AxisTrees(self, 0, vals=obst)
             self._vart = AxisTrees(self, 1, vals=vart)
+            self._update_has_overlap()
 
     def _init_as_view(self, tdata_ref: TreeData, oidx: Index1D | None, vidx: Index1D | None):
         super()._init_as_view(tdata_ref, oidx=oidx, vidx=vidx)
@@ -209,6 +212,7 @@ def _init_as_view(self, tdata_ref: TreeData, oidx: Index1D | None, vidx: Index1D
         self._tree_label = tdata_ref._tree_label
         self._alignment = tdata_ref._alignment
         self._allow_overlap = tdata_ref._allow_overlap
+        self._has_overlap = tdata_ref._has_overlap
 
         # view of obst and vart
         self._obst = tdata_ref.obst._view(self, oidx)
@@ -259,6 +263,17 @@ def allow_overlap(self) -> bool:
         """Whether overlapping trees are allowed."""
         return self._allow_overlap
 
+    @property
+    def has_overlap(self) -> bool:
+        """
+        Flag indicating whether stored trees contain overlapping nodes.
+
+        Returns
+        -------
+        bool - ``True`` when any stored trees share nodes, ``False`` otherwise.
+        """
+        return self._has_overlap
+
     @property
     def alignment(self) -> Literal["leaves", "nodes", "subset"]:
         """Mapping between trees and observations/variables."""
@@ -278,11 +293,13 @@ def is_view(self) -> bool:
     def obst(self, value):
         obst = AxisTrees(self, 0, vals=dict(value))
         self._obst = obst
+        self._update_has_overlap()
 
     @vart.setter
     def vart(self, value):
         vart = AxisTrees(self, 1, vals=dict(value))
         self._vart = vart
+        self._update_has_overlap()
 
     @allow_overlap.setter
     def allow_overlap(self, value):
@@ -295,6 +312,34 @@ def allow_overlap(self, value):
                         f"One or more trees in {attr} have overlapping nodes. Cannot set allow_overlap to False."
                     )
         self._allow_overlap = value
+        self._update_has_overlap()
+
+    def _update_has_overlap(self) -> None:
+        """
+        Update the cached overlap indicator.
+
+        Ensures the cached `_has_overlap` flag matches the current state of stored
+        trees.
+
+        Parameters
+        ----------
+        None
+            This method does not accept any parameters.
+
+        Returns
+        -------
+        None - This function updates the `_has_overlap` attribute in place.
+        """
+        if not self._allow_overlap:
+            self._has_overlap = False
+            return
+
+        has_overlap = False
+        if hasattr(self, "_obst"):
+            has_overlap = has_overlap or self._obst._check_tree_overlap()
+        if hasattr(self, "_vart"):
+            has_overlap = has_overlap or self._vart._check_tree_overlap()
+        self._has_overlap = has_overlap
 
     @alignment.setter
     def alignment(self, value):

tests/test_base.py

@@ -52,6 +52,7 @@ def test_attributes(X, tree, axis):
     assert getattr(tdata, f"{dim}t").parent is tdata
     assert list(getattr(tdata, f"{dim}t").dim_names) == ["0", "1", "2"]
     assert tdata.allow_overlap is False
+    assert tdata.has_overlap is False
     assert tdata.label is None
 
 
@@ -107,17 +108,30 @@ def test_tree_overlap(X, tree):
     tdata = td.TreeData(X, obst={"0": tree, "1": second_tree}, allow_overlap=True)
     check_graph_equality(tdata.obst["0"], tree)
     check_graph_equality(tdata.obst["1"], second_tree)
+    assert tdata.has_overlap is True
     # Test set allow_overlap to True
     tdata = td.TreeData(X, obst={"0": tree}, allow_overlap=False)
     assert tdata.allow_overlap is False
+    assert tdata.has_overlap is False
     tdata.allow_overlap = True
     tdata.obst["1"] = tree
     assert list(tdata.obst.keys()) == ["0", "1"]
     assert tdata.allow_overlap
+    assert tdata.has_overlap is True
     # Cannot set allow_overlap to False when overlap is present
     with pytest.raises(ValueError):
         tdata.allow_overlap = False
     assert tdata.allow_overlap
+    assert tdata.has_overlap is True
+
+
+def test_has_overlap_updates_on_delete(X, tree):
+    second_tree = nx.DiGraph()
+    second_tree.add_edges_from([("root", "0"), ("root", "1")])
+    tdata = td.TreeData(X, obst={"0": tree, "1": second_tree}, allow_overlap=True)
+    assert tdata.has_overlap is True
+    del tdata.obst["1"]
+    assert tdata.has_overlap is False
 
 
 def test_alignment(X, tree):