Usage

Point assignment

import cuspa as cs

polygons = cs.io.from_geoarrow(geoarrow_table)
cell_ids = cs.tl.assign_points(points_xy, polygons)

cell_ids is an int32 device array. A value of -1 means the point was not assigned to a polygon.

Transcript aggregation

aggregate_to_cells consumes an iterable of (points_xy, gene_ids) CuPy batches. It does not read files or choose a batch size. Pass a one-item iterable for in-memory data:

csr = cs.tl.aggregate_to_cells(
    [(points_xy, gene_ids)],
    polygons,
    num_genes=num_genes,
)

Yield one CUDA batch at a time for larger inputs:

import cupy as cp


def gpu_batches(points_host, gene_ids_host, batch_size):
    for start in range(0, len(gene_ids_host), batch_size):
        stop = start + batch_size
        yield (
            cp.asarray(points_host[start:stop], dtype=cp.float32),
            cp.asarray(gene_ids_host[start:stop], dtype=cp.int32),
        )


csr = cs.tl.aggregate_to_cells(
    gpu_batches(points_host, gene_ids_host, batch_size=1_000_000),
    polygons,
    num_genes,
)

The returned CSRMatrix contains three device arrays:

• indptr

• indices

• data

Its shape field is a Python tuple.

Use csr.to_cupy_csr() or csr.to_torch_sparse() to create a framework-native sparse matrix without copying through host memory.

Batches are merged with exact integer arithmetic. Counts are int32, so each (cell, gene) count must be at most 2**31 - 1; aggregate_to_cells raises OverflowError otherwise. indptr and indices are int32 until the result has more than 2**31 - 1 non-zeros, and int64 after that. cupyx.scipy.sparse only supports int32 indices, so to_cupy_csr() raises OverflowError for such a matrix; to_torch_sparse() and the raw arrays still work.

Choosing a batch size

GPU memory must hold the largest input batch, its aggregation workspace, the polygon index, the accumulated CSR, and the merged CSR while a batch is added. Batching limits input memory; the final CSR must still fit on the GPU.

SpatialData adapter

import cuspa as cs

polygons = cs.io.from_spatialdata(sdata, "cell_boundaries")
transcripts = cs.io.transcripts_from_spatialdata(sdata, "transcripts")

csr = cs.tl.aggregate_to_cells(
    [(transcripts.xy, transcripts.gene_ids)],
    polygons,
    num_genes=len(transcripts.gene_names),
)

For large transcript tables, read x, y, and gene columns in record batches instead of materializing a point GeoDataFrame.

Buffer lifetime and streams

Polygons validates its buffers at construction and caches derived AABBs and a uniform-grid index. Treat the buffers as immutable. After changing them in place, call polygons.clear_cache() before the next query.

The optional stream= argument is a raw CUDA stream handle used by the native kernels. Operations may synchronize it while sizing indexes or results. For aggregate_to_cells, also make the corresponding CuPy stream current because buffer copies and merge buffers use CuPy’s current stream. Keep inputs alive and synchronize or use the same stream before consuming results elsewhere.

Limits

Polygon, ring, and spatial-index offsets use signed 32-bit indices and are checked against the 2^31-1 limit. Point and polygon indices in overlap_pairs results are int32, but the number of pairs uses 64-bit offsets and is limited only by device memory (8 bytes per pair). Each aggregation batch is limited to 2^31-1 transcripts. The aggregated CSR switches to int64 indptr and indices past 2^31-1 non-zero (cell, gene) pairs, so its size is limited only by device memory. Each count is int32 and limited to 2^31-1; larger counts raise OverflowError.