Fix mergetree bug

cereeberus/cereeberus/reeb/merge.py

@@ -19,7 +19,7 @@ class MergeTree(ReebGraph):
 
     """
 
-    def __init__(self, T=None, root=None, f={}, labels={}, seed=None, verbose=False):
+    def __init__(self, T=None, root=None, f={}, labels=None, seed=None, verbose=False):
         """
         Initialize a merge tree object.
 
@@ -42,7 +42,7 @@ def __init__(self, T=None, root=None, f={}, labels={}, seed=None, verbose=False)
             # Fix up the drawing locations
             self.fix_pos_f()
 
-        self.labels = labels
+        self.labels = labels if labels is not None else {}
 
     def __str__(self):
         return f"MergeTree with {len(self.nodes)} nodes and {len(self.edges)} edges."
@@ -64,6 +64,41 @@ def get_leaves(self):
 
         return [v for v in self.nodes if self.down_degree(v) == 0]
 
+    def smoothing_and_maps(self, eps=1, verbose=False):
+        """
+        Builds the ``eps``-smoothed merge tree and the associated vertex/edge maps.
+
+        ``ReebGraph.smoothing_and_maps`` already handles the ``v_inf`` root
+        correctly (it treats any infinite-valued vertex as a formal marker
+        and reattaches it after smoothing the finite part), but it returns a
+        plain ``ReebGraph``. This override just re-wraps that result as a
+        ``MergeTree`` so callers get back the expected type.
+
+        Parameters:
+            eps (float): The amount of smoothing to apply.
+            verbose (bool): Optional. If True, prints additional information.
+
+        Returns:
+            tuple: MergeTree, vertex_map, edge_map.
+        """
+        R_eps_generic, map_V, map_E = super().smoothing_and_maps(
+            eps=eps, verbose=verbose
+        )
+
+        R_eps = MergeTree()  # constructor already adds v_inf at f=inf
+        for v in R_eps_generic.nodes():
+            if v == "v_inf":
+                continue
+            R_eps.add_node(v, R_eps_generic.f[v], reset_pos=False)
+        for u, v, _ in R_eps_generic.edges(keys=True):
+            R_eps.add_edge(u, v, reset_pos=False)
+
+        # remap each label's target vertex through map_V, as smoothing renames vertices, a plain copy of self.lables would point to vertex names that no longer exist in the smoothed tree.
+        R_eps.labels = {key: map_V[val] for key, val in self.labels.items()}
+
+        R_eps.set_pos_from_f()
+        return R_eps, map_V, map_E
+
     def add_node(self, vertex, f_vertex, reset_pos=True):
         """
         Adds a node to the tree. Note that this will break the single connected component property of the tree so we assume you will do this in the process of adding more connecting edges.

cereeberus/cereeberus/reeb/reebgraph.py

@@ -9,7 +9,6 @@
 # from build.lib.cereeberus.reeb import graph
 
 
-
 class ReebGraph(nx.MultiDiGraph):
     """
     A Reeb graph stored as a networkx ``MultiDiGraph``. The function values are stored as a dictionary. The directedness of the edges follows the convention that the edge goes from the lower function value to the higher function value node.
@@ -79,23 +78,23 @@ def copy(self):
         """
         # Create a new ReebGraph with copies of the nodes and edges
         H = ReebGraph()
-        
+
         # Copy the function values dictionary
         H.f = self.f.copy()
-        
+
         # Copy all nodes and edges from the parent MultiDiGraph
         for v in self.nodes():
             H.add_node(v, self.f[v], reset_pos=False)
-        
+
         for u, v, key in self.edges(keys=True):
             super(ReebGraph, H).add_edge(u, v, key)
-        
+
         # Copy position information if it exists
-        if hasattr(self, 'pos_f') and self.pos_f:
+        if hasattr(self, "pos_f") and self.pos_f:
             H.pos_f = self.pos_f.copy()
-        if hasattr(self, 'pos') and self.pos:
+        if hasattr(self, "pos") and self.pos:
             H.pos = self.pos.copy()
-        
+
         return H
 
     def branch_decomp(self):
@@ -107,6 +106,7 @@ def branch_decomp(self):
                 ``decompose`` method has already been called on this graph.
         """
         from .branchdecomp import BranchDecomp
+
         bd = BranchDecomp()
         bd.decompose(self)
         return bd
@@ -185,18 +185,18 @@ def number_connected_components(self):
 
     def get_upward_path(self, start_vertex):
         """Return an upward path from the starting vertex by greedy dynamic choice.
-        
+
         Input:
             start_vertex: a vertex in the graph to start from
-            
+
         Output:
             path: a list of vertices representing the upward path
-            
+
         """
-        # Check that the vertex is in the graph 
+        # Check that the vertex is in the graph
         if start_vertex not in self.nodes:
-            raise ValueError(f"The vertex {start_vertex} is not in the Reeb graph.")   
-        
+            raise ValueError(f"The vertex {start_vertex} is not in the Reeb graph.")
+
         path = [start_vertex]
         while self.up_degree(path[-1]) > 0:
             s = next(self.successors(path[-1]))
@@ -518,7 +518,7 @@ def remove_edges_from(self, edges, reset_pos=True):
 
         if reset_pos:
             self.set_pos_from_f()
-    
+
     def remove_path_from(self, path, reset_pos=True):
         """Remove a path from the Reeb graph. A path is a list of vertices, and this method will remove one edge along each step of the path.
 
@@ -937,6 +937,65 @@ def smoothing_and_maps(self, eps=1, verbose=False):
                 {e: [e] for e in self.edges(keys=True)},
             )
 
+        # A vertex at infinity (like root of a mergetree) needs to be allowed to be smoothed. Perturbing it by a finite eps isn't meaningful. So we only smooth the finite part of the graph, then reattach each infinite vertex exactly as it was, reconneted to whichever finite vertices it end up adjacent to it after smoothing.
+
+        # For edges, they always point from lower to higher f value, a +inf vertex only has a predecessor and a -inf vertex only has a successor.
+
+        inf_nodes = [v for v in self.nodes if np.isinf(self.f[v])]
+        if inf_nodes:
+            # smooth the finite part, then reattach inf vertices afterward
+            finite_nodes = [v for v in self.nodes if v not in inf_nodes]
+
+            if not finite_nodes:
+                # the whole graph is infinite valued (eg MergeTree with only a root). There's nothing to smooth, so just return the graph
+                return(self, {v: v for v in self.nodes}, {e: [e] for e in self.edges(keys=True)})
+
+            f_finite = {v: self.f[v] for v in finite_nodes}
+            G_finite = ReebGraph(self.subgraph(finite_nodes), f_finite)
+            R_eps, map_V, map_E = G_finite.smoothing_and_maps(eps=eps, verbose=verbose)
+
+            for inf_v in inf_nodes:
+                # map_V[u] can land mid-component; use the component's actual top/bottom
+                components = list(nx.weakly_connected_components(R_eps))
+                comp_top = {}
+                comp_bottom = {}
+                for comp in components:
+                    tops = [v for v in comp if R_eps.up_degree(v) == 0]
+                    bottoms = [v for v in comp if R_eps.down_degree(v) == 0]
+                    for v in comp:
+                        comp_top[v] = tops
+                        comp_bottom[v] = bottoms
+
+                R_eps.add_node(inf_v, self.f[inf_v], reset_pos=False)
+                map_V[inf_v] = inf_v
+
+                preds = list(self.predecessors(inf_v))
+                attach_up = {u: comp_top[map_V[u]] for u in preds}
+                for tv in set(v for verts in attach_up.values() for v in verts):
+                    R_eps.add_edge(tv, inf_v, reset_pos=False)
+                for u in preds:
+                    new_edges_u = [
+                        (tv, inf_v, R_eps.number_of_edges(tv, inf_v) - 1)
+                        for tv in attach_up[u]
+                    ]
+                    for key in range(self.number_of_edges(u, inf_v)):
+                        map_E[(u, inf_v, key)] = new_edges_u
+
+                succs = list(self.successors(inf_v))
+                attach_down = {w: comp_bottom[map_V[w]] for w in succs}
+                for bv in set(v for verts in attach_down.values() for v in verts):
+                    R_eps.add_edge(inf_v, bv, reset_pos=False)
+                for w in succs:
+                    new_edges_w = [
+                        (inf_v, bv, R_eps.number_of_edges(inf_v, bv) - 1)
+                        for bv in attach_down[w]
+                    ]
+                    for key in range(self.number_of_edges(inf_v, w)):
+                        map_E[(inf_v, w, key)] = new_edges_w
+
+            R_eps.set_pos_from_f()
+            return R_eps, map_V, map_E
+
         # Get the list of critical values to place new nodes
         crit_vals = list(set(self.f.values()))
         new_crit_vals = [cv + eps for cv in crit_vals]
@@ -1023,13 +1082,27 @@ def smoothing_and_maps(self, eps=1, verbose=False):
                 ]
                 if len(lower_vert) > 1:
                     print(f"{i,c} has multiple lower vertices")
+                elif len(lower_vert) == 0:
+                    raise ValueError(
+                        f"No component found below critical value {cv!r} "
+                        f"for component {c!r}; this usually means a "
+                        "critical value has no finite neighbor to shift "
+                        "toward (e.g. an unhandled infinite value)."
+                    )
                 else:
                     lower_vert = lower_vert[0]
 
                 overlap_up = np.array([len(c.intersection(c_0)) for c_0 in C])
                 upper_vert = [comp_to_new_vert[u] for u in np.where(overlap_up > 0)[0]]
                 if len(upper_vert) > 1:
                     print(f"{i,c} has multiple upper vertices")
+                elif len(upper_vert) == 0:
+                    raise ValueError(
+                        f"No component found above critical value {cv!r} "
+                        f"for component {c!r}; this usually means a "
+                        "critical value has no finite neighbor to shift "
+                        "toward (e.g. an unhandled infinite value)."
+                    )
                 else:
                     upper_vert = upper_vert[0]
 

tests/test_merge_tree.py

@@ -58,9 +58,51 @@ def test_merge_labels(self):
         # Check that the LCA matrix is symmetric and of the same size as the leaf set 
         self.assertEqual(M.shape[0], len(Leaves))
 
-        # TODO: Add tests for specifc labeling functions, here I only have for the leaf version
-
+        self.assertEqual(M.shape[1], len(Leaves))
 
+    def test_merge_labels_by_labels_type(self):
+        MT = ex_mt.randomMergeTree(10)
 
+        # Label all leaves, then check the label-keyed LCA matrix matches
+        # the leaf-keyed one in shape and values.
+        MT.label_all_leaves()
+        M_labels = MT.LCA_matrix(type="labels")
+        M_leaves = MT.LCA_matrix()
+
+        self.assertEqual(M_labels.shape, M_leaves.shape)
+        np.testing.assert_array_equal(M_labels, M_leaves)
+
+        # add_label_edge should subdivide an edge and register a new label
+        leaves = MT.get_leaves()
+        u = leaves[0]
+        v = list(MT.successors(u))[0]
+        f_mid = (MT.f[u] + MT.f[v]) / 2
+        n_labels_before = len(MT.labels)
+        MT.add_label_edge(u, v, "mid_vertex", f_mid, label="mid")
+        self.assertIn("mid", MT.labels)
+        self.assertEqual(len(MT.labels), n_labels_before + 1)
+
+    def test_smoothing_preserves_type_and_root(self):
+        MT = ex_mt.randomMergeTree(9)
+        MT_eps = MT.smoothing(1)
+
+        self.assertIsInstance(MT_eps, MergeTree)
+        self.assertTrue('v_inf' in MT_eps.nodes)
+        self.assertEqual(MT_eps.f['v_inf'], np.inf)
+        self.assertEqual(MT_eps.up_degree('v_inf'), 0)
+        self.assertEqual(set(MT_eps.nodes), set(MT_eps.f.keys()))
+        self.assertEqual(set(MT_eps.nodes), set(MT_eps.pos_f.keys()))
+
+    def test_smoothing_two_separate_branches_to_root(self):
+        MT = MergeTree()
+        MT.add_node('p', 0)
+        MT.add_node('q', 1)
+        MT.add_edge('p', 'v_inf')
+        MT.add_edge('q', 'v_inf')
+
+        MT_eps, _, map_E = MT.smoothing_and_maps(0.5)
+        self.assertIsInstance(MT_eps, MergeTree)
+        self.assertEqual(MT_eps.down_degree('v_inf'), 2)
+        self.assertNotEqual(map_E[('p', 'v_inf', 0)], map_E[('q', 'v_inf', 0)])
 if __name__ == '__main__':
     unittest.main()

tests/test_reeb_class.py

@@ -266,6 +266,22 @@ def test_smoothing_internal_namespaced_vertices(self):
         ]
         self.assertGreater(len(internal_nodes), 0)
 
+    def test_smoothing_with_infinite_node(self):
+        # A vertex at +/- infinity (e.g. a MergeTree's v_inf root) should
+        # smooth without crashing, staying on top/bottom of the result.
+        R = ReebGraph()
+        R.add_node('a', 0)
+        R.add_node('b', 1)
+        R.add_node('c', 2)
+        R.add_node('top', float('inf'))
+        R.add_edge('a', 'b')
+        R.add_edge('b', 'c')
+        R.add_edge('c', 'top')
+
+        R_eps = R.smoothing(1)
+        self.check_reeb(R_eps)
+        self.assertEqual(R_eps.up_degree('top'), 0)
+
     def test_matrices(self):
         # This test makes sure you can get the adjacency matrix and boundary matrix of a Reeb graph.
         R = ex_rg.juggling_man()