Fix push_both, update version number

numba_celltree/__init__.py

@@ -1,5 +1,5 @@
 from numba_celltree.celltree import CellTree2d
 
-__version__ = "0.2.0"
+__version__ = "0.2.1"
 
 __all__ = ("CellTree2d",)

numba_celltree/constants.py

@@ -142,8 +142,7 @@ class CellTreeData(NamedTuple):
 # that can be included is given by 2 ** (depth of stack - 1).
 # (int(math.ceil(math.log(MAX_N_FACE, 2))) + 1)
 # This is only true for relatively balanced trees. MAX_N_FACE = int(2e9)
-# results in required stack of 32. 128 suffices for some more degenerate input
-# (triangulation result of complex polygon earcuts).
+# results in required stack of 32.
 INITIAL_TREE_DEPTH = 32
 # Floating point slack
 TOLERANCE_ON_EDGE = 1e-9

numba_celltree/creation.py

@@ -230,8 +230,8 @@ def pessimistic_n_nodes(n_polys: int):
 @nb.njit(inline="always")
 def push_both(root_stack, dim_stack, root, dim, size):
     root_stack, size_root = push(root_stack, root, size)
-    _, _ = push(dim_stack, dim, size)
-    return size_root
+    dim_stack, _ = push(dim_stack, dim, size)
+    return root_stack, dim_stack, size_root
 
 
 @nb.njit(inline="always")
@@ -361,7 +361,9 @@ def build(
                 if dim_flag >= 0:
                     dim_flag = (not dim) - 2
                     nodes[root_index]["dim"] = not root.dim
-                    size = push_both(root_stack, dim_stack, root_index, dim_flag, size)
+                    root_stack, dim_stack, size = push_both(
+                        root_stack, dim_stack, root_index, dim_flag, size
+                    )
                 else:  # Already split once, convert to leaf.
                     nodes[root_index]["Lmax"] = -1
                     nodes[root_index]["Rmin"] = -1
@@ -387,8 +389,12 @@ def build(
         node_index = push_node(nodes, left_child, node_index)
         node_index = push_node(nodes, right_child, node_index)
 
-        size = push_both(root_stack, dim_stack, child_ind + 1, right_child.dim, size)
-        size = push_both(root_stack, dim_stack, child_ind, left_child.dim, size)
+        root_stack, dim_stack, size = push_both(
+            root_stack, dim_stack, child_ind + 1, right_child.dim, size
+        )
+        root_stack, dim_stack, size = push_both(
+            root_stack, dim_stack, child_ind, left_child.dim, size
+        )
 
     return node_index
 

tests/test_utils.py

@@ -37,6 +37,9 @@ def test_push_size_increases():
     stack, size = ut.push(stack, 1, size)
     assert stack.size == 2
     assert size_before < stack.size
+    stack, size = ut.push(stack, 10, size)
+    assert stack.size == 4
+    assert np.array_equal(stack[:3], [0, 1, 10])
 
 
 def test_copy():