diff --git a/src/segger/cli/segment.py b/src/segger/cli/segment.py index 4848930..8832590 100644 --- a/src/segger/cli/segment.py +++ b/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, diff --git a/src/segger/data/data_module.py b/src/segger/data/data_module.py index ffb6d3d..d3fd6cc 100644 --- a/src/segger/data/data_module.py +++ b/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": diff --git a/src/segger/data/tiling.py b/src/segger/data/tiling.py index 1e2f920..bac0e85 100644 --- a/src/segger/data/tiling.py +++ b/src/segger/data/tiling.py @@ -209,11 +209,15 @@ 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') @@ -221,6 +225,7 @@ def __init__( points, max_tile_size, with_bounds=True, + quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts, ) self._tiles = quadtree_to_geoseries(quadtree, backend='geopandas') diff --git a/src/segger/data/utils/heterodata.py b/src/segger/data/utils/heterodata.py index e20a9bf..1745220 100644 --- a/src/segger/data/utils/heterodata.py +++ b/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]:,}") diff --git a/src/segger/data/utils/neighbors.py b/src/segger/data/utils/neighbors.py index f21b204..279c313 100644 --- a/src/segger/data/utils/neighbors.py +++ b/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( diff --git a/src/segger/geometry/quadtree.py b/src/segger/geometry/quadtree.py index a55c9e1..8850032 100644 --- a/src/segger/geometry/quadtree.py +++ b/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} " diff --git a/src/segger/geometry/query.py b/src/segger/geometry/query.py index dd4a0c3..8781da7 100644 --- a/src/segger/geometry/query.py +++ b/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,7 +244,12 @@ 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, @@ -239,6 +257,7 @@ def points_in_polygons( max_unasigned_points, boundary_buffer, batches, + quadtree_downsample_n_transcripts=quadtree_downsample_n_transcripts, ) def polygons_in_polygons(