Initial memory optimisations

Author: Tobiaspk

src/segger/cli/segment.py

@@ -154,7 +154,7 @@ def segment(
         group=group_prediction,
     )] = registry.get_default("prediction_graph_max_k"),
 
-    prediction_expansion_ratio: Annotated[float | None, registry.get_parameter(
+    prediction_graph_buffer_ratio: Annotated[float | None, registry.get_parameter(
         "prediction_graph_buffer_ratio",
         validator=validators.Number(gt=0),
         group=group_prediction,
@@ -323,7 +323,7 @@ def segment(
         transcripts_graph_max_dist=transcripts_max_dist,
         prediction_graph_mode=prediction_mode,
         prediction_graph_max_k=prediction_max_k,
-        prediction_graph_buffer_ratio=prediction_expansion_ratio,
+        prediction_graph_buffer_ratio=prediction_graph_buffer_ratio,
         tiling_margin_training=tiling_margin_training,
         tiling_margin_prediction=tiling_margin_prediction,
         tiling_nodes_per_tile=max_nodes_per_tile,

src/segger/data/tiling.py

@@ -199,7 +199,7 @@ def __init__(
     ):
         # Calculate QuadTree on points and set as tiles
         points = points_to_geoseries(positions, backend='cuspatial')
-        _, quadtree = get_quadtree_index(
+        _, quadtree, _ = get_quadtree_index(
             points,
             max_tile_size,
             with_bounds=True,

src/segger/geometry/quadtree.py

@@ -25,13 +25,14 @@ def get_quadtree_kwargs(
         A dictionary of keyword arguments including x_min, x_max, y_min,
         y_max, scale, and max_depth.
     """
-    # Calculate bounds
-    x_min = float(points.points.x.min())
-    x_max = float(points.points.x.max())
-    y_min = float(points.points.y.min())
-    y_max = float(points.points.y.max())
-
-    # Get hyperparams for quadtree
+    # Calculate bounds | Optimisation: Use interleaved view, without copying data
+    xy = cp.asarray(points.points.xy).reshape(-1, 2)  # zero-copy view                                                                                                                                                                                                                                                                                                                                                                                 
+    x_min = float(xy[:, 0].min())
+    x_max = float(xy[:, 0].max())                                                                                                                                                                                                                                                                                                                                                                                                                      
+    y_min = float(xy[:, 1].min())
+    y_max = float(xy[:, 1].max())       
+
+   # Get hyperparams for quadtree
     extent = max(x_max - x_min, y_max - y_min)
     max_depth = 1
     while extent // (1 << max_depth) > 0:
@@ -140,7 +141,7 @@ def get_quadtree_index(
     points: cuspatial.GeoSeries,
     max_size: int,
     with_bounds: bool = True,
-) -> tuple[cudf.Series, cudf.DataFrame]:
+) -> tuple[cudf.Series, cudf.DataFrame, dict]:
     """Build a cuSpatial quadtree from 2D point data.
 
     Parameters
@@ -190,7 +191,7 @@ def get_quadtree_index(
             y_max=y_max,
         )
 
-    return indices, quadtree
+    return indices, quadtree, kwargs
 
 
 def quadtree_to_geoseries(

src/segger/geometry/query.py

@@ -3,6 +3,7 @@
 import numpy as np
 import cuspatial
 import cudf
+import cupy as cp
 
 from .conversion import (
     polygons_to_geoseries,
@@ -47,14 +48,15 @@ def _points_in_polygons_contains(
         mapping each contained point to its containing polygon.
     """
     # Setup inputs for spatial join
+    cp.get_default_memory_pool().free_all_blocks()
+    cp.get_default_pinned_memory_pool().free_all_blocks()
     if max_size is None:
         max_size = 10000 if len(points) > 5e7 else 1000  # heuristic
-    point_indices, quadtree = get_quadtree_index(
+    point_indices, quadtree, kwargs = get_quadtree_index(
         points,
         max_size,
         with_bounds=False
     )
-    kwargs = get_quadtree_kwargs(points)
 
     # Perform spatial join in batches
     batch_idx = np.linspace(0, len(polygons), (batches or 1) + 1, dtype=int)