[main] Downsample transcripts for quadtree tiling construction

src/segger/cli/segment.py

@@ -191,6 +191,11 @@ def segment(
         group=group_tiling,
     )] = registry.get_default("tiling_nodes_per_tile"),
 
+    quadtree_downsample_n_transcripts: Annotated[int | None, registry.get_parameter(
+        "quadtree_downsample_n_transcripts",
+        group=group_tiling,
+    )] = registry.get_default("quadtree_downsample_n_transcripts"),
+
     max_edges_per_batch: Annotated[int, registry.get_parameter(
         "edges_per_batch",
         validator=validators.Number(gt=0),
@@ -353,6 +358,7 @@ def segment(
         tiling_margin_training=tiling_margin_training,
         tiling_margin_prediction=tiling_margin_prediction,
         tiling_nodes_per_tile=max_nodes_per_tile,
+        quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
         edges_per_batch=max_edges_per_batch,
         gene_corr_reference_path=gene_corr_reference_path,
         gene_missing_strategy=gene_missing_strategy,

src/segger/data/data_module.py

@@ -121,6 +121,10 @@ class ISTDataModule(LightningDataModule):
         Margin width (in µm) added to tiles during prediction.
     tiling_nodes_per_tile : int, default=50000
         Maximum number of nodes per tile for adaptive tiling.
+    quadtree_downsample_n_transcripts : int or None, default=50_000_000
+        Downsample points for quadtree construction (adaptive tiling and
+        prediction graph) to avoid numeric overflows (see segger issue
+        #77). Set to None to disable downsampling
     tiling_side_length : float, default=250.0
         Side length of square tiles (benchmarking only).
     training_fraction : float, default=0.75
@@ -153,6 +157,7 @@ class ISTDataModule(LightningDataModule):
     tiling_margin_training: float = 20.
     tiling_margin_prediction: float = 20.
     tiling_nodes_per_tile: int = 50_000
+    quadtree_downsample_n_transcripts: int | None = 50_000_000
     tiling_side_length: float = 250.  # TODO: Remove (benchmarking only)
     training_fraction: float = 0.75
     edges_per_batch: int = 1_000_000
@@ -237,6 +242,7 @@ def load(self):
             prediction_graph_mode=self.prediction_graph_mode,
             prediction_graph_max_k=self.prediction_graph_max_k,
             prediction_graph_buffer_ratio=self.prediction_graph_buffer_ratio,
+            quadtree_downsample_n_transcripts=self.quadtree_downsample_n_transcripts,
         )
 
         # Tile graph dataset
@@ -249,6 +255,7 @@ def load(self):
             self.tiling = QuadTreeTiling(
                 positions=node_positions,
                 max_tile_size=self.tiling_nodes_per_tile,
+                quadtree_downsample_n_transcripts=self.quadtree_downsample_n_transcripts,
             )
         #TODO: Remove (benchmarking only)
         elif self.tiling_mode == "square":

src/segger/data/tiling.py

@@ -209,18 +209,23 @@ class QuadTreeTiling(Tiling):
         the quadtree.
     max_tile_size : int
         The maximum number of points allowed in any single quadtree tile.
+    quadtree_downsample_n_transcripts : int, optional
+        Downsample transcripts for quadtree construction to avoid numeric overflows
+        (see segger issue #77). Set to None to disable downsampling
     """
     def __init__(
         self,
         positions: torch.Tensor,
         max_tile_size: int,
+        quadtree_downsample_n_transcripts: int | None = None,
     ):
         # Calculate QuadTree on points and set as tiles
         points = points_to_geoseries(positions, backend='cuspatial')
         _, quadtree, _ = get_quadtree_index(
             points,
             max_tile_size,
             with_bounds=True,
+            quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
         )
         self._tiles = quadtree_to_geoseries(quadtree, backend='geopandas')
 

src/segger/data/utils/heterodata.py

@@ -25,6 +25,7 @@ def setup_heterodata(
     prediction_graph_mode: Literal["nucleus", "cell", "uniform"],
     prediction_graph_max_k: int,
     prediction_graph_buffer_ratio: float,
+    quadtree_downsample_n_transcripts: int | None = 50_000_000,
     cells_embedding_key: str = 'X_pca',
     cells_clusters_column: str = 'phenograph_cluster',
     cells_encoding_column: str = 'cell_encoding',
@@ -158,6 +159,7 @@ def setup_heterodata(
         max_k=prediction_graph_max_k,
         buffer_ratio=prediction_graph_buffer_ratio,
         mode=prediction_graph_mode,
+        quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
     )
     logger.info(f"  tx-neighbors-bd edges: {data['tx', 'neighbors', 'bd'].edge_index.shape[1]:,}")
 

src/segger/data/utils/neighbors.py

@@ -203,6 +203,7 @@ def setup_prediction_graph(
     max_k: int,
     buffer_ratio: float,
     mode: Literal['nucleus', 'cell', 'uniform'] = 'cell',
+    quadtree_downsample_n_transcripts: int | None = None,
 ) -> torch.Tensor:
     """TODO: Add description.
     """
@@ -232,6 +233,7 @@ def setup_prediction_graph(
         polygons=polygons,
         predicate='contains',
         batches=100,
+        quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
     )
 
     return torch.tensor(

src/segger/geometry/quadtree.py

@@ -144,6 +144,7 @@ def get_quadtree_index(
     max_size: int,
     with_bounds: bool = True,
     max_retries: int = 5,
+    quadtree_downsample_n_transcripts: int | None = None,
 ) -> tuple[cudf.Series, cudf.DataFrame, dict]:
     """Build a cuSpatial quadtree from 2D point data.
 
@@ -158,6 +159,10 @@ def get_quadtree_index(
         DataFrame. Default is True.
     max_retries : int, optional
         Retries on invalid tree, each growing ``max_size`` by 5%. Default 6.
+    quadtree_downsample_n_transcripts : int, optional
+        Downsample transcripts for quadtree construction to avoid numeric overflows
+        (see segger issue #77). While this does not speed up the construction,
+        it becomes more memory efficient. Set to None to disable downsampling.
 
     Returns
     -------
@@ -166,7 +171,7 @@ def get_quadtree_index(
     quadtree : cudf.DataFrame
         DataFrame of quadtree tiles with spatial bounds and metadata.
     """
-    # Get hyperparams for quadtree
+    # Get hyperparams for quadtree (always on the full point set, see issue #77)
     kwargs = get_quadtree_kwargs(points)
     x_min = kwargs['x_min']
     x_max = kwargs['x_max']
@@ -175,6 +180,20 @@ def get_quadtree_index(
     scale = kwargs['scale']
     max_depth = kwargs['max_depth']
 
+    # Subsample points for tree construction to bound memory / avoid overflow
+    # on very large point counts (see segger issue #77)
+    n_points = len(points)
+    retry_step = 10000
+    if quadtree_downsample_n_transcripts is not None and n_points > quadtree_downsample_n_transcripts:
+        fraction = quadtree_downsample_n_transcripts / n_points
+        sample_idx = cp.random.randint(0, n_points, size=quadtree_downsample_n_transcripts)
+        points = points.iloc[sample_idx]
+
+        # adjust max-size and growth parameter
+        max_size = max(1, round(max_size * fraction))
+        retry_step = max(1, round(retry_step * fraction))
+        logger.debug(f"Subsampling {quadtree_downsample_n_transcripts / (1e6):.1f}M / {n_points / (1e6):.1f}M ({fraction:.1%}) points for quadtree ")
+
     logger.debug(f"Building quadtree on {len(points)} points with max_size={max_size}, max_depth={max_depth}")
 
     # Calculate quadtree on region (retry on invalid tree, see issue #40)
@@ -193,8 +212,8 @@ def get_quadtree_index(
         # check if valid (see segger issue #40)
         if is_quadtree_valid(quadtree, len(points)):
             break
-        logger.warning(f"Invalid quadtree generated with max_size={max_size}. Retry with max_size={max_size + 10000}.")
-        max_size += 10000
+        logger.warning(f"Invalid quadtree generated with max_size={max_size}. Retry with max_size={max_size + retry_step}.")
+        max_size += retry_step
     else:
         raise RuntimeError(
             f"cuSpatial returned an invalid quadtree after {max_retries + 1} "

src/segger/geometry/query.py

@@ -23,6 +23,7 @@ def _points_in_polygons_contains(
     polygons: cuspatial.GeoSeries,
     max_size: int | None = None,
     batches: int | None = None,
+    quadtree_downsample_n_transcripts: int | None = None,
 ) -> cudf.DataFrame:
     """Finds which points are strictly contained within polygons.
 
@@ -56,7 +57,8 @@ def _points_in_polygons_contains(
     point_indices, quadtree, kwargs = get_quadtree_index(
         points,
         max_size,
-        with_bounds=False
+        with_bounds=False,
+        quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
     )
 
     # Perform spatial join in batches
@@ -105,6 +107,7 @@ def _points_in_polygons_intersects(
     max_unassigned_points: int = 100_000,
     boundary_buffer: float = 1e-9,
     batches: int | None = None,
+    quadtree_downsample_n_transcripts: int | None = None,
 ) -> cudf.DataFrame:
     """Finds points that intersect polygons, including boundaries.
 
@@ -135,7 +138,12 @@ def _points_in_polygons_intersects(
         mapping each intersecting point to its polygon.
     """
     # GPU pass to find all points strictly contained by the polygons
-    contains = _points_in_polygons_contains(points, polygons, batches=batches)
+    contains = _points_in_polygons_contains(
+        points,
+        polygons,
+        batches=batches,
+        quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
+    )
 
     # Isolate points not found, which are potential boundary cases
     idx_all = cudf.RangeIndex(len(points))
@@ -151,7 +159,11 @@ def _points_in_polygons_intersects(
             ply_ixn.buffer(boundary_buffer),
             backend='cuspatial',
         )
-        in_buffer = _points_in_polygons_contains(pts_ixn, ply_buf)
+        in_buffer = _points_in_polygons_contains(
+            pts_ixn,
+            ply_buf,
+            quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
+        )
         in_buffer = in_buffer['index_query'].drop_duplicates()
         pts_ixn = pts_ixn.iloc[in_buffer]
 
@@ -181,7 +193,8 @@ def points_in_polygons(
     predicate: Literal['contains', 'intersects'] = 'intersects',
     max_unasigned_points: int = 100_000,
     boundary_buffer: float = 1e-9,
-    batches: int | None = None
+    batches: int | None = None,
+    quadtree_downsample_n_transcripts: int | None = None,
 ) -> cudf.DataFrame:
     """Finds which points fall inside which polygons using a given predicate.
 
@@ -231,14 +244,20 @@ def points_in_polygons(
 
     # Perform spatial join
     if predicate == 'contains':
-        return _points_in_polygons_contains(points, polygons, batches=batches)
+        return _points_in_polygons_contains(
+            points,
+            polygons,
+            batches=batches,
+            quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
+        )
     else:  # predicate == 'intersects'
         return _points_in_polygons_intersects(
             points,
             polygons,
             max_unasigned_points,
             boundary_buffer,
             batches,
+            quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts,
         )
 
 def polygons_in_polygons(