From 92ac5a6e603ad02b2cef6f2126ea6d88fd946379 Mon Sep 17 00:00:00 2001
From: Tobiaspk <tobiaspk1@gmail.com>
Date: Thu, 3 Sep 2026 11:59:47 -0400
Subject: [PATCH] Fix quadtree leaf boundaries to match cuSpatial's actual
 clamped scale

get_quadtree_kwargs() computed scale without checking cuSpatial's own
minimum-scale requirement, which cuSpatial silently clamps up to and
only reports via a UserWarning. get_quadrant_bounds() then computed
each leaf's polygon using the pre-clamp scale, so tile boundaries no
longer matched the tree cuSpatial actually built. On the CosMx
pancreas dataset this dropped ~50% of transcript points from every
tile on the initial join, cascading into an all -1 label tensor and
a bincount crash.

Compute scale directly from cuSpatial's min-scale formula so it's
never undershot, and pass the real (scale, max_depth) into
get_quadrant_bounds so leaf polygons match the actual tree.
---
 src/segger/geometry/quadtree.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/segger/geometry/quadtree.py b/src/segger/geometry/quadtree.py
index 38b4091..29397f9 100644
--- a/src/segger/geometry/quadtree.py
+++ b/src/segger/geometry/quadtree.py
@@ -40,7 +40,10 @@ def get_quadtree_kwargs(
     max_depth = 1
     while extent // (1 << max_depth) > 0:
         max_depth += 1
-    scale = extent // (1 << max_depth - 1)
+    max_depth = min(max_depth, 15)
+
+    # Match cuSpatial's own min-scale formula so it never silently clamps this up.
+    scale = extent / ((1 << max_depth) + 2)
 
     # Return as dictionary
     return dict(
@@ -49,7 +52,7 @@ def get_quadtree_kwargs(
         y_min=y_min,
         y_max=y_max,
         scale=scale,
-        max_depth=min(max_depth, 15),
+        max_depth=max_depth,
     )
 
 
@@ -100,6 +103,8 @@ def get_quadrant_bounds(
     x_max: float,
     y_min: float,
     y_max: float,
+    scale: float,
+    max_depth: int,
 ):
     """
     Add spatial bounds to each leaf in a cuSpatial quadtree.
@@ -115,6 +120,12 @@ def get_quadrant_bounds(
         Full extent of the quadtree in x-direction.
     y_min, y_max : float
         Full extent of the quadtree in y-direction.
+    scale : float
+        The `scale` actually used by cuSpatial to build `quadtree` (i.e. the
+        leaf cell width). Must match, or leaf boundaries will be computed for
+        a different cell size than the tree cuSpatial built.
+    max_depth : int
+        The `max_depth` actually used by cuSpatial to build `quadtree`.
 
     Returns
     -------
@@ -122,13 +133,10 @@ def get_quadrant_bounds(
         Input DataFrame with added bounding box columns: 'x_min', 'x_max',
         'y_min', and 'y_max'.
     """
-    width =  x_max - x_min
-    height = y_max - y_min
     levels = quadtree['level'].astype(float) + 1
     coords = cp.array(keys_to_coordinates(quadtree['key'].to_numpy()))
-    quadrant_max = np.ceil(np.log2(max(width, height)))
-    quadrant_dim = 2 ** (quadrant_max - levels)
-    
+    quadrant_dim = scale * 2 ** (max_depth - levels)
+
     quadtree['x_min'] = x_min + coords[0] * quadrant_dim
     quadtree['x_max'] = quadtree['x_min'] + quadrant_dim
     quadtree['y_min'] = y_min + coords[1] * quadrant_dim
@@ -215,6 +223,8 @@ def get_quadtree_index(
             x_max=x_max,
             y_min=y_min,
             y_max=y_max,
+            scale=scale,
+            max_depth=max_depth,
         )
 
     return indices, quadtree, kwargs