Fix computereeb reset pos performance

cereeberus/cereeberus/compute/computereeb.py

@@ -171,7 +171,7 @@ def _dedup(lst):
         verts_at_level = []
         for rep, comp in components_at_level.items():
             nextNodeName = R.get_next_vert_name()
-            R.add_node(nextNodeName, now_min)
+            R.add_node(nextNodeName, now_min, reset_pos=False)
             vert_to_component[nextNodeName] = comp
             verts_at_level.append(nextNodeName)
 
@@ -182,7 +182,7 @@ def _dedup(lst):
                 if any(
                     is_face(prev_simp, simp) for simp in comp for prev_simp in prev_comp
                 ):
-                    R.add_edge(e, nextNodeName)
+                    R.add_edge(e, nextNodeName, reset_pos=False)
 
         # Step 4: Remove vertices and horizontal simplices – they live only at this exact height.
         for vert in vert_names:
@@ -213,7 +213,7 @@ def _dedup(lst):
         edges_at_prev_level = []
         for comp in components_above.values():
             e_name = "e_" + str(half_edge_index)
-            R.add_node(e_name, (now_min + now_max) / 2)
+            R.add_node(e_name, (now_min + now_max) / 2, reset_pos=False)
             vert_to_component[e_name] = comp
             half_edge_index += 1
             edges_at_prev_level.append(e_name)
@@ -224,6 +224,8 @@ def _dedup(lst):
                 if any(
                     is_face(simp, prev_simp) for simp in comp for prev_simp in prev_comp
                 ):
-                    R.add_edge(v, e_name)
+                    R.add_edge(v, e_name, reset_pos=False)
 
+    # All the nodes and edges above were added with reset_pos=False, so now we need to call set_pos_from_f() to set the positions
+    R.set_pos_from_f()
     return R

cereeberus/cereeberus/reeb/reebgraph.py

@@ -432,8 +432,14 @@ def remove_node(self, vertex, reset_pos=True):
         super().remove_node(vertex)
         del self.f[vertex]
 
-        if reset_pos and hasattr(self, "pos_f"):
+        # drop the old position for this vertex unconditionally. Skipping it leaves a pos_f entry for a vertex no longer in the graph.
+
+        if hasattr(self, "pos_f") and vertex in self.pos_f:
             del self.pos_f[vertex]
+        if hasattr(self, "pos") and vertex in self.pos:
+            del self.pos[vertex]
+
+        if reset_pos and hasattr(self, "pos_f"):
             self.set_pos_from_f()
 
     def remove_nodes_from(self, nodes, reset_pos=True):
@@ -475,19 +481,18 @@ def add_edge(self, u, v, reset_pos=True):
         else:
             # the function values are the same, so the edge collapses the two vertices
             # wlog we're going to get rid of v, and add all its edges to u
-
             # get the edges of v
             edges_in = self.in_edges(v)
             edges_out = self.out_edges(v)
 
             # add the edges to u
             for e in edges_in:
-                self.add_edge(e[0], u)
+                self.add_edge(e[0], u, reset_pos=False)
             for e in edges_out:
-                self.add_edge(u, e[1])
+                self.add_edge(u, e[1], reset_pos=False)
 
             # Remove v
-            self.remove_node(v)
+            self.remove_node(v, reset_pos=False)
 
         if reset_pos:
             self.set_pos_from_f()

tests/test_lowerstar_class.py

@@ -1,62 +1,64 @@
 import unittest
 from cereeberus import ReebGraph, LowerStar, computeReeb
 from cereeberus.data import Torus
+import time
 
 
 class TestLowerStarClass(unittest.TestCase):
-    
+
     def test_lower_star_assign_filtration(self):
         # Test that assigning filtration values to vertices correctly updates adjacent simplices.
         K = LowerStar()
         K.insert([0, 1, 2])
         K.insert([1, 3])
         K.insert([2, 3])
-        
+
         K.assign_filtration(0, 0.0)
         self.assertEqual(K.filtration([0]), 0.0)
         self.assertEqual(K.filtration([0, 1]), 0.0)
         self.assertEqual(K.filtration([0, 2]), 0.0)
         self.assertEqual(K.filtration([0, 1, 2]), 0.0)
-        
+
         K.assign_filtration(1, 3.0)
         self.assertEqual(K.filtration([1]), 3.0)
         self.assertEqual(K.filtration([1, 3]), 3.0)
         self.assertEqual(K.filtration([0, 1]), 3.0)  # Updated due to vertex 1
-        self.assertEqual(K.filtration([0, 1, 2]), 3.0)  # Updated due to vertex 1
-
+        # Updated due to vertex 1
+        self.assertEqual(K.filtration([0, 1, 2]), 3.0)
+
     def test_lower_star_sc_max_min_filtration(self):
         # Test that max and min filtration values are computed correctly.
         K = LowerStar()
         K.insert([0, 1])
         K.insert([1, 2])
         K.insert([2, 3])
-        
+
         K.assign_filtration(0, 1.0)
         K.assign_filtration(1, 2.0)
         K.assign_filtration(2, 3.0)
         K.assign_filtration(3, 4.0)
-        
+
         self.assertEqual(K.min_filtration(), 1.0)
         self.assertEqual(K.max_filtration(), 4.0)
-        
+
     def test_computeReeb(self):
         # Test the computation of the Reeb graph from a Lower Star Simplicial Complex.
         K = LowerStar()
         K.insert([0, 1, 2])
         K.insert([1, 3])
         K.insert([2, 3])
-        
+
         K.assign_filtration(0, 0.0)
         K.assign_filtration(1, 3.0)
         K.assign_filtration(2, 5.0)
         K.assign_filtration(3, 7.0)
-        
+
         R = computeReeb(K)
-        
+
         self.assertIsInstance(R, ReebGraph)
         self.assertGreater(len(R.nodes), 0)
         self.assertGreater(len(R.edges), 0)
-        
+
     def test_computeReeb_horizontal_edge(self):
         # Regression test: when two vertices share the same filtration value and
         # are connected by an edge (a "horizontal" edge), the level-set connected
@@ -92,11 +94,81 @@ def test_computeReeb_horizontal_edge(self):
     def test_torus_example_class(self):
         # Test the torus example from the documentation.
         T = Torus()
-        T.generate_grid(grid_size = 4)
-        T.assign_random_values(0,100, seed=1986)
-        
+        T.generate_grid(grid_size=4)
+        T.assign_random_values(0, 100, seed=1986)
+
         R = computeReeb(T)
-        
+
         self.assertIsInstance(R, ReebGraph)
         self.assertGreater(len(R.nodes), 0)
-        self.assertGreater(len(R.edges), 0)
+        self.assertGreater(len(R.edges), 0)
+
+    def test_computeReeb_pos_f_consistency(self):
+        # Regression test for the reset_pos fix: computeReeb builds
+        # the whole graph with reset_pos=False and only calls
+        # set_pos_from_f() once at the end. Make sure that single call
+        # actually leaves every bookkeeping structure fully in sync -
+        # in particular, no vertex should be missing from pos_f, and none
+        # should be a stale leftover with no corresponding node.
+        K = LowerStar()
+        K.insert([0, 1, 2])
+        K.insert([1, 3])
+        K.insert([2, 3])
+
+        K.assign_filtration(0, 0.0)
+        K.assign_filtration(1, 3.0)
+        K.assign_filtration(2, 5.0)
+        K.assign_filtration(3, 7.0)
+
+        R = computeReeb(K)
+
+        # nodes, f, and pos_f should all track exactly the same vertex set
+        self.assertEqual(set(R.nodes), set(R.f.keys()))
+        self.assertEqual(set(R.nodes), set(R.pos_f.keys()))
+
+        # y-coordinate of every position should be exactly f(v)
+        for v in R.nodes:
+            self.assertEqual(R.pos_f[v][1], R.f[v])
+
+        # every edge should point toward the higher function value
+        for edge in R.edges:
+            u, v = edge[:2]
+            self.assertGreater(R.f[v], R.f[u])
+
+    def test_computeReeb_torus_pos_f_consistency(self):
+        # Same consistency check as above, but on a bigger graph with
+        # multiple critical points and merges, to exercise more of the
+        # deferred reset_pos path (including tied-value collapses).
+        T = Torus()
+        T.generate_grid(grid_size=6)
+        T.assign_random_values(0, 100, seed=1986)
+
+        R = computeReeb(T)
+
+        self.assertEqual(set(R.nodes), set(R.f.keys()))
+        self.assertEqual(set(R.nodes), set(R.pos_f.keys()))
+        for v in R.nodes:
+            self.assertEqual(R.pos_f[v][1], R.f[v])
+
+    def test_computeReeb_performance(self):
+        # Regression test for the O(N) layout-recomputation bug: computeReeb
+        # used to call set_pos_from_f() (a full scipy.optimize.minimize
+        # layout pass) on every single add_node/add_edge instead of once at
+        # the end. On this input the buggy version takes ~9s; the fixed
+        # version takes well under 1s. Use a generous bound (5s) so this
+        # isn't flaky on a slow CI runner, while still failing hard if the
+        # O(N) behavior is ever reintroduced.
+        T = Torus()
+        T.generate_grid(grid_size=6)
+        T.assign_random_values(0, 100, seed=1986)
+
+        t0 = time.time()
+        R = computeReeb(T)
+        dt = time.time() - t0
+
+        self.assertGreater(R.number_of_nodes(), 0)
+        self.assertLess(
+            dt, 5.0,
+            f"computeReeb took {dt:.2f}s on a {R.number_of_nodes()}-node graph "
+            "-- did the reset_pos=False deferral get reverted?"
+        )

tests/test_reeb_class.py

@@ -301,7 +301,51 @@ def test_set_pos_from_f_preserves_y_function_values(self):
             self.assertEqual(R.pos_f[v][1], R.f[v])
 
 
+    def test_remove_node_deferred_pos_cleanup(self):
+        # Regression test: remove_node(reset_pos=False) must still drop the
+        # removed vertex's pos_f entry immediately. Previously this cleanup
+        # was bundled inside `if reset_pos:`, so a deferred removal left a
+        # dangling pos_f entry for a vertex no longer in the graph.
+        R = ex_rg.simple_loops()
+        R.set_pos_from_f()  # establish an initial pos_f for every node
+
+        v = next(iter(R.nodes))
+        R.remove_node(v, reset_pos=False)
+
+        self.assertNotIn(v, R.nodes)
+        self.assertNotIn(
+            v, R.pos_f,
+            "pos_f still has a stale entry for a removed vertex"
+        )
+        # No blanket check_reeb() here: an isolated deferred removal can
+        # leave pos_f short of an entry for OTHER now-orphaned nodes until
+        # the next set_pos_from_f() call; that's expected, not a bug.
 
+    def test_add_edge_collapse_respects_reset_pos_false(self):
+        R = ReebGraph()
+        R.add_node('a', 0.0, reset_pos=False)
+        R.add_node('b', 1.0, reset_pos=False)
+        R.add_node('c', 1.0, reset_pos=False)  # same f as 'b' -> collapse
+        R.add_node('d', 2.0, reset_pos=False)
+
+        R.add_edge('a', 'b', reset_pos=False)
+        R.add_edge('b', 'd', reset_pos=False)
+        # Triggers the tied-value collapse branch inside add_edge.
+        R.add_edge('a', 'c', reset_pos=False)
+
+        # pos_f exists (ReebGraph.__init__ always calls set_pos_from_f once,
+        # even on an empty graph) but should still be EMPTY here -- nothing
+        # after construction should have triggered a recompute, including
+        # the collapse's internal recursive add_edge/remove_node calls.
+        self.assertEqual(
+            R.pos_f, {},
+            "pos_f was populated even though every call used reset_pos=False "
+            "-- the collapse branch must be forwarding reset_pos incorrectly"
+        )
+
+        # Now do the single deferred layout call, as computeReeb does.
+        R.set_pos_from_f()
+        self.check_reeb(R)