|
11 | 11 | `BucketSorter.add`), kept on disk in a scratch group (`scratch_group`); |
12 | 12 | 3. read one bucket at a time, sort it, and append it to the output |
13 | 13 | (`BucketSorter.runs`). |
| 14 | +
|
| 15 | +Input that already arrives in row order -- as Cell Ranger and STARsolo write |
| 16 | +their MTX files -- skips the buckets (`grouped_runs`). Either way the runs |
| 17 | +go through `DuplicateSummer` and into a `CsrAppender`. |
14 | 18 | """ |
15 | 19 |
|
16 | 20 | from __future__ import annotations |
@@ -51,6 +55,23 @@ def count_keys(counts: np.ndarray, keys: np.ndarray) -> None: |
51 | 55 | counts[unique] += seen |
52 | 56 |
|
53 | 57 |
|
| 58 | +def sum_keys(sums: np.ndarray, keys: np.ndarray, weights: np.ndarray) -> None: |
| 59 | + """Add `weights[i]` to `sums[keys[i]]`, over the span the block covers. |
| 60 | +
|
| 61 | + As `count_keys`, and for the same reason: cost set by the block, not by |
| 62 | + the length of `sums`. |
| 63 | + """ |
| 64 | + if keys.size == 0: |
| 65 | + return |
| 66 | + lo, hi = int(keys.min()), int(keys.max()) |
| 67 | + span = hi - lo + 1 |
| 68 | + if span <= 4 * keys.size: |
| 69 | + sums[lo : hi + 1] += np.bincount(keys - lo, weights=weights, minlength=span) |
| 70 | + else: |
| 71 | + unique, inverse = np.unique(keys, return_inverse=True) |
| 72 | + sums[unique] += np.bincount(inverse, weights=weights, minlength=unique.size) |
| 73 | + |
| 74 | + |
54 | 75 | def bucket_bounds( |
55 | 76 | indptr: np.ndarray, *, chunk: int, bucket_entries: Optional[int] = None |
56 | 77 | ) -> np.ndarray: |
@@ -176,3 +197,132 @@ def runs(self) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]: |
176 | 197 | values = np.asarray(vals[...]) |
177 | 198 | order = np.lexsort((minor, major)) |
178 | 199 | yield major[order], minor[order], values[order] |
| 200 | + |
| 201 | + |
| 202 | +Triple = Tuple[np.ndarray, np.ndarray, np.ndarray] |
| 203 | + |
| 204 | +#: The largest nnz an int32 `indptr` can address. |
| 205 | +INT32_MAX = int(np.iinfo(np.int32).max) |
| 206 | + |
| 207 | + |
| 208 | +def grouped_runs(blocks: Iterator[Triple], *, sort: bool) -> Iterator[Triple]: |
| 209 | + """Runs in CSR order, from blocks already grouped by row. |
| 210 | +
|
| 211 | + The fast path, for input whose rows arrive in order -- as Cell Ranger and |
| 212 | + STARsolo write them. A row may straddle two blocks, so the last row of |
| 213 | + each block is held back and joined to the next. With `sort`, entries are |
| 214 | + put in column order within each row; without, they are taken to be in it |
| 215 | + already. |
| 216 | + """ |
| 217 | + carry: Optional[Triple] = None |
| 218 | + for rows, cols, values in blocks: |
| 219 | + if carry is not None: |
| 220 | + rows = np.concatenate((carry[0], rows)) |
| 221 | + cols = np.concatenate((carry[1], cols)) |
| 222 | + values = np.concatenate((carry[2], values)) |
| 223 | + carry = None |
| 224 | + if rows.size == 0: |
| 225 | + continue |
| 226 | + cut = int(np.searchsorted(rows, rows[-1], side="left")) |
| 227 | + carry = (rows[cut:], cols[cut:], values[cut:]) |
| 228 | + if cut: |
| 229 | + yield _ordered(rows[:cut], cols[:cut], values[:cut], sort) |
| 230 | + if carry is not None and carry[0].size: |
| 231 | + yield _ordered(*carry, sort) |
| 232 | + |
| 233 | + |
| 234 | +def _ordered(rows: np.ndarray, cols: np.ndarray, values: np.ndarray, sort: bool) -> Triple: |
| 235 | + if not sort: |
| 236 | + return rows, cols, values |
| 237 | + order = np.lexsort((cols, rows)) |
| 238 | + return rows[order], cols[order], values[order] |
| 239 | + |
| 240 | + |
| 241 | +class DuplicateSummer: |
| 242 | + """Merge repeated (row, column) entries by adding their values. |
| 243 | +
|
| 244 | + As scipy does when it builds CSR from coordinates, and so as scanpy's |
| 245 | + readers give them. Each run must hold whole rows, sorted, which both |
| 246 | + `grouped_runs` and `BucketSorter.runs` guarantee. |
| 247 | + """ |
| 248 | + |
| 249 | + def __init__(self) -> None: |
| 250 | + self.summed = 0 |
| 251 | + |
| 252 | + def __call__(self, runs: Iterator[Triple]) -> Iterator[Triple]: |
| 253 | + for rows, cols, values in runs: |
| 254 | + same = (rows[1:] == rows[:-1]) & (cols[1:] == cols[:-1]) |
| 255 | + if same.any(): |
| 256 | + starts = np.flatnonzero(np.concatenate(([True], ~same))) |
| 257 | + self.summed += int(rows.size - starts.size) |
| 258 | + values = np.add.reduceat(values, starts) |
| 259 | + rows, cols = rows[starts], cols[starts] |
| 260 | + yield rows, cols, values |
| 261 | + |
| 262 | + |
| 263 | +class CsrAppender: |
| 264 | + """A CSR matrix written row by row: entries must arrive in row order. |
| 265 | +
|
| 266 | + `data` and `indices` grow as entries are appended; `indptr` is written |
| 267 | + by `finish` from the rows seen, as int32 when nnz allows, as anndata and |
| 268 | + scipy would write it. |
| 269 | + """ |
| 270 | + |
| 271 | + def __init__( |
| 272 | + self, |
| 273 | + parent: Any, |
| 274 | + name: str, |
| 275 | + *, |
| 276 | + shape: Tuple[int, int], |
| 277 | + data_dtype: Any, |
| 278 | + index_dtype: Any = np.int32, |
| 279 | + compression: Optional[str] = None, |
| 280 | + ) -> None: |
| 281 | + from adata.core.subset import _growable_like |
| 282 | + from adata.elements import spec |
| 283 | + from adata.elements.write import set_shape_attr |
| 284 | + |
| 285 | + self.group = parent.create_group(name) |
| 286 | + spec.set_encoding(self.group, spec.CSR_MATRIX) |
| 287 | + set_shape_attr(self.group, shape) |
| 288 | + self.shape = shape |
| 289 | + self.data_dtype = np.dtype(data_dtype) |
| 290 | + self.index_dtype = np.dtype(index_dtype) |
| 291 | + self.compression = compression |
| 292 | + self.data = _growable_like( |
| 293 | + self.group, "data", self.data_dtype, None, compression=compression |
| 294 | + ) |
| 295 | + self.indices = _growable_like( |
| 296 | + self.group, "indices", self.index_dtype, None, compression=compression |
| 297 | + ) |
| 298 | + self.counts = np.zeros(shape[0], dtype=np.int64) |
| 299 | + self.nnz = 0 |
| 300 | + |
| 301 | + def append(self, rows: np.ndarray, cols: np.ndarray, values: np.ndarray) -> None: |
| 302 | + from adata.core import subset |
| 303 | + |
| 304 | + if rows.size == 0: |
| 305 | + return |
| 306 | + if self.data_dtype.kind in "iu": |
| 307 | + info = np.iinfo(self.data_dtype) |
| 308 | + if values.min() < info.min or values.max() > info.max: |
| 309 | + raise ValueError( |
| 310 | + f"A value does not fit {self.data_dtype}: summing duplicate " |
| 311 | + f"entries reached {int(values.max()):,}." |
| 312 | + ) |
| 313 | + count_keys(self.counts, rows) |
| 314 | + subset._append(self.data, values.astype(self.data_dtype, copy=False)) |
| 315 | + subset._append(self.indices, cols.astype(self.index_dtype, copy=False)) |
| 316 | + self.nnz += int(rows.size) |
| 317 | + |
| 318 | + def finish(self) -> int: |
| 319 | + from adata.core.subset import _sparse_dataset |
| 320 | + |
| 321 | + dtype = np.int32 if self.nnz <= INT32_MAX else np.int64 |
| 322 | + indptr = np.concatenate(([0], np.cumsum(self.counts))).astype(dtype) |
| 323 | + dataset = _sparse_dataset( |
| 324 | + self.group, "indptr", np.dtype(dtype), indptr.size, None, |
| 325 | + compression=self.compression, |
| 326 | + ) |
| 327 | + dataset[:] = indptr |
| 328 | + return self.nnz |
0 commit comments