crispyx 0.2.0__tar.gz → 0.2.2__tar.gz

{crispyx-0.2.0/src/crispyx.egg-info → crispyx-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: crispyx
-Version: 0.2.0
+Version: 0.2.2
 Summary: Memory-efficient streaming analysis of large-scale CRISPR and Perturb-seq screens on disk-backed AnnData files
 Author: Jin-Hong Du
 License: Modified MIT License
@@ -106,7 +106,7 @@ Requires-Dist: numba>=0.59
 Requires-Dist: pandas>=1.5
 Requires-Dist: scipy>=1.10
 Requires-Dist: h5py>=3.0
-Requires-Dist: joblib>=1.0
+Requires-Dist: joblib>=1.4
 Requires-Dist: scikit-learn>=1.0
 Requires-Dist: scikit-misc>=0.1.4
 Requires-Dist: scanpy>=1.9.2
@@ -130,7 +130,6 @@ Requires-Dist: sphinx-rtd-theme>=1.2; extra == "docs"
 Requires-Dist: sphinx-copybutton>=0.5; extra == "docs"
 Requires-Dist: nbsphinx>=0.9; extra == "docs"
 Requires-Dist: ipykernel>=6.0; extra == "docs"
-Requires-Dist: tomli>=2.0; python_version < "3.11" and extra == "docs"
 Dynamic: license-file
 
 # crispyx

{crispyx-0.2.0 → crispyx-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "crispyx"
-version = "0.2.0"
+dynamic = ["version"]
 description = "Memory-efficient streaming analysis of large-scale CRISPR and Perturb-seq screens on disk-backed AnnData files"
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -40,7 +40,7 @@ dependencies = [
     "pandas>=1.5",
     "scipy>=1.10",
     "h5py>=3.0",
-    "joblib>=1.0",
+    "joblib>=1.4",
     "scikit-learn>=1.0",
     "scikit-misc>=0.1.4",
     "scanpy>=1.9.2",
@@ -69,7 +69,6 @@ docs = [
     "sphinx-copybutton>=0.5",
     "nbsphinx>=0.9",
     "ipykernel>=6.0",
-    "tomli>=2.0; python_version < '3.11'",
 ]
 
 [project.urls]
@@ -78,6 +77,9 @@ Documentation = "https://crispyx.readthedocs.io"
 Repository = "https://github.com/jinhongdu-lab/crispyx"
 "Bug Tracker" = "https://github.com/jinhongdu-lab/crispyx/issues"
 
+[tool.setuptools.dynamic]
+version = {attr = "crispyx._version.__version__"}
+
 [tool.setuptools.packages.find]
 where = ["src"]
 

{crispyx-0.2.0 → crispyx-0.2.2}/src/crispyx/__init__.py

@@ -15,12 +15,7 @@ import os
 # Override with NUMBA_THREADING_LAYER=omp if you manage the OMP conflict yourself.
 os.environ.setdefault("NUMBA_THREADING_LAYER", "workqueue")
 
-from importlib.metadata import PackageNotFoundError, version
-
-try:
-    __version__ = version("crispyx")
-except PackageNotFoundError:
-    __version__ = "0.0.0.dev"
+from ._version import __version__
 
 # ---------------------------------------------------------------------------
 # Public API re-exports

{crispyx-0.2.0 → crispyx-0.2.2}/src/crispyx/_checkpoint.py

@@ -13,14 +13,16 @@ import json
 import logging
 import os
 import shutil
-import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 import h5py
 import numpy as np
 import pandas as pd
 
+from . import _messages
+from ._memory import drop_file_cache
+
 if TYPE_CHECKING:
     from tqdm import tqdm
 
@@ -302,26 +304,104 @@ def _sha256(array: np.ndarray) -> str:
     return hashlib.sha256(np.ascontiguousarray(array).tobytes()).hexdigest()
 
 
-class ResumableRun:
-    """Result arrays and a progress checkpoint that outlive an interruption.
+#: Elements per block when a result matrix is read or rewritten a block of
+#: rows at a time (32 MB of float64), so no step holds a whole result.
+_BLOCK_ELEMENTS = 1 << 22
+
+
+def _row_blocks(n_rows: int, n_genes: int, chunk_rows: int | None = None):
+    """Slices covering ``n_rows`` rows of ``n_genes`` columns, about
+    :data:`_BLOCK_ELEMENTS` elements each. With ``chunk_rows`` (the rows of
+    an HDF5 chunk), each slice is whole bands of chunks, so no chunk is read
+    twice."""
+    step = max(1, _BLOCK_ELEMENTS // max(n_genes, 1))
+    if chunk_rows:
+        step = max(1, round(step / chunk_rows)) * chunk_rows
+    for start in range(0, n_rows, step):
+        yield slice(start, min(start + step, n_rows))
+
+
+def _chunk_rows(dataset) -> int | None:
+    """Rows in one HDF5 chunk of ``dataset`` (``None`` when contiguous or
+    not an HDF5 dataset)."""
+    chunks = getattr(dataset, "chunks", None)
+    return chunks[0] if chunks else None
+
+
+def _create_array(group: h5py.Group, name: str, **kwargs) -> h5py.Dataset:
+    """``group.create_dataset`` tagged as an anndata dense array."""
+    ds = group.create_dataset(name, **kwargs)
+    ds.attrs["encoding-type"] = "array"
+    ds.attrs["encoding-version"] = "0.2.0"
+    return ds
+
+
+def _create_unfilled_array(
+    group: h5py.Group, name: str, shape: tuple[int, ...], dtype, chunks: tuple[int, ...] | None = None,
+) -> h5py.Dataset:
+    """An anndata dense array whose storage is allocated now but never
+    filled: creating it writes nothing (on disk it is a sparse hole until
+    written), and the file's layout is final from the start, so filling it
+    later only writes data -- a kill mid-write leaves a readable file.
+    Unwritten elements read as zero. Contiguous, or in HDF5 ``chunks``
+    (every chunk allocated up front, so the layout is still final)."""
+    dcpl = h5py.h5p.create(h5py.h5p.DATASET_CREATE)
+    dcpl.set_alloc_time(h5py.h5d.ALLOC_TIME_EARLY)
+    dcpl.set_fill_time(h5py.h5d.FILL_TIME_NEVER)
+    if chunks is not None:
+        dcpl.set_chunk(tuple(int(c) for c in chunks))
+    return _create_array(group, name, shape=shape, dtype=dtype, dcpl=dcpl)
+
+
+#: Bytes in one HDF5 chunk of a result written in column blocks.
+_COLUMN_BLOCK_CHUNK_BYTES = 1 << 20
+
+
+def _column_block_chunks(n_groups: int, n_genes: int, chunk_size: int) -> tuple[int, int]:
+    """HDF5 chunk shape for a result written one gene chunk (a block of
+    ``chunk_size`` columns) at a time: ``(rows, chunk_size)``, ~1 MiB.
+
+    In a row-major contiguous dataset a gene chunk is one small write per
+    group row (10,922 x 4.6 KB per layer on a genome-wide screen), each
+    straddling partial pages that the filesystem must read before writing.
+    On a network filesystem that read is a round trip: on Lustre a
+    gene chunk took 219 s to write contiguously even with its pages cached,
+    and 918 s when they were evicted between chunks, against 22 s in these
+    chunks. On a local disk the two layouts write alike (the chunked one
+    slightly faster), so it is used everywhere; reading one row of a large
+    result reads its band of chunks (~2 ms locally), and small results are
+    read into memory whole anyway.
+    """
+    width = max(1, min(int(chunk_size), n_genes))
+    rows = max(1, min(n_groups, _COLUMN_BLOCK_CHUNK_BYTES // (width * 8)))
+    return rows, width
 
-    The arrays are memmaps in a hidden ``.{output name}.resume`` directory
-    beside the output (a temporary directory would vanish with the process,
-    taking the finished work with it). The checkpoint JSON records the call's
-    :func:`run_fingerprint` plus caller-defined progress, and :meth:`save`
-    flushes the arrays before writing it, so everything a checkpoint lists
-    as done is on disk.
 
-    A caller that writes its partial output itself (rather than into
-    arrays) keeps it in :attr:`directory` too and names it in ``requires``.
+class ResumableRun:
+    """A result file filled in place, and a progress checkpoint, that outlive
+    an interruption.
+
+    The work file is :attr:`work_path`, ``result.h5ad`` in a hidden
+    ``.{output name}.resume`` directory beside the output (a temporary
+    directory would vanish with the process, taking the finished work with
+    it). ``scaffold(path)`` creates it with every dataset the run fills,
+    pre-sized; the run writes its results into :attr:`file` as they are
+    computed, and :meth:`finish` moves the completed file to the output.
+    Nothing is held in memory or mapped, so the run's footprint does not
+    grow with the result.
+
+    The checkpoint JSON records the call's :func:`run_fingerprint` plus
+    caller-defined progress. :meth:`save` hands the file's pending data to the
+    operating system and evicts its pages before writing the checkpoint, so
+    everything a checkpoint lists as done survives the process being killed
+    and little of it stays charged to the process's memory.
 
     A run resumes only when ``resume`` is set, the checkpoint's fingerprint
-    matches this call (:func:`fingerprint_difference`), and every array
-    exists with the expected shape and dtype (plus every file named in
-    ``requires``). Otherwise the stale checkpoint and directory are
-    discarded and the run starts fresh, so a checkpoint can never be paired
-    with another run's data; a resume that throws away a checkpoint warns
-    and says why.
+    matches this call (:func:`fingerprint_difference`), and the work file
+    opens with every dataset in ``datasets`` at its expected shape and dtype.
+    Otherwise the stale checkpoint and directory are discarded and the run
+    starts fresh, so a checkpoint can never be paired with another run's
+    data; a resume that throws away a checkpoint warns and says why.
     """
 
     def __init__(
@@ -330,13 +410,15 @@ class ResumableRun:
         checkpoint_path: Path,
         *,
         fingerprint: dict,
-        arrays: dict[str, tuple[tuple[int, ...], "np.typing.DTypeLike", float | int | bool]],
         resume: bool,
-        requires: tuple[str, ...] = (),
+        scaffold: "Callable[[Path], None]",
+        datasets: "dict[str, tuple[tuple[int, ...], np.typing.DTypeLike]]",
     ) -> None:
         self.checkpoint_path = checkpoint_path
         self.fingerprint = fingerprint
         self.directory = output_path.with_name(f".{output_path.name}.resume")
+        self.work_path = self.directory / "result.h5ad"
+        self.file: h5py.File | None = None
         checkpoint = _read_checkpoint(checkpoint_path, required_keys=("fingerprint",)) if resume else None
         mismatch = (
             fingerprint_difference(checkpoint["fingerprint"], fingerprint)
@@ -344,56 +426,65 @@ class ResumableRun:
             else "it cannot be read"
         )
         if checkpoint is not None and mismatch is not None:
-            warnings.warn(
+            _messages.emit(
                 f"Not resuming from the checkpoint at {checkpoint_path}: {mismatch}. "
                 "Discarding it and starting over.",
-                stacklevel=3,
             )
         self.resumed = (
-            checkpoint is not None
-            and mismatch is None
-            and all((self.directory / name).exists() for name in requires)
-            and all(
-                (self.directory / f"{name}.dat").exists()
-                and (self.directory / f"{name}.dat").stat().st_size
-                == int(np.prod(shape, dtype=np.int64)) * np.dtype(dtype).itemsize
-                for name, (shape, dtype, _fill) in arrays.items()
-            )
+            checkpoint is not None and mismatch is None and _has_datasets(self.work_path, datasets)
         )
         self.progress = {k: v for k, v in checkpoint.items() if k != "fingerprint"} if self.resumed else {}
-        self.arrays: dict[str, np.memmap] = {}
         if not self.resumed:
             self.discard()
-        self.directory.mkdir(parents=True, exist_ok=True)
-        for name, (shape, dtype, fill) in arrays.items():
-            path = self.directory / f"{name}.dat"
-            if self.resumed:
-                self.arrays[name] = np.memmap(path, mode="r+", dtype=dtype, shape=shape)
-            else:
-                array = np.memmap(path, mode="w+", dtype=dtype, shape=shape)
-                if fill != 0:
-                    array.fill(fill)
-                self.arrays[name] = array
+            self.directory.mkdir(parents=True, exist_ok=True)
+            scaffold(self.work_path)
+        self.file = h5py.File(self.work_path, "r+")
 
     def save(self, **progress) -> None:
-        """Flush the arrays, then record ``progress`` in the checkpoint."""
-        for array in self.arrays.values():
-            array.flush()
+        """Put the work file's data on disk, then record ``progress`` in the
+        checkpoint."""
+        self.file.flush()
+        drop_file_cache(self.work_path)
         _write_checkpoint_atomic(self.checkpoint_path, {"fingerprint": self.fingerprint, **progress})
 
-    def discard(self) -> None:
-        """Remove the checkpoint and the arrays (after success, or when stale).
+    def close(self) -> None:
+        """Close the work file, for the steps that finish it by path."""
+        if self.file is not None:
+            self.file.close()
+            self.file = None
+        drop_file_cache(self.work_path)
 
-        Drops this object's references to the arrays; once the caller drops
-        its own, the mapped pages are released.
-        """
-        self.arrays = {}
+    def finish(self, output_path: Path) -> None:
+        """Move the completed work file to ``output_path`` and remove the
+        checkpoint and directory."""
+        self.close()
+        os.replace(self.work_path, output_path)
+        self.discard()
+
+    def discard(self) -> None:
+        """Remove the checkpoint and the work file (after success, or when stale)."""
+        if self.file is not None:
+            self.file.close()
+            self.file = None
         self.checkpoint_path.unlink(missing_ok=True)
-        # ignore_errors: on Windows a file still mapped by a live result view
-        # cannot be deleted; the next run on this output removes it.
         shutil.rmtree(self.directory, ignore_errors=True)
 
 
+def _has_datasets(path: Path, datasets: dict) -> bool:
+    """Whether the h5 file at ``path`` opens and holds every dataset in
+    ``datasets`` with its ``(shape, dtype)``."""
+    try:
+        with h5py.File(path, "r") as handle:
+            return all(
+                name in handle
+                and handle[name].shape == tuple(shape)
+                and handle[name].dtype == np.dtype(dtype)
+                for name, (shape, dtype) in datasets.items()
+            )
+    except (OSError, KeyError):  # absent, or a kill mid-write left it unreadable
+        return False
+
+
 class _DummyProgress:
     """Dummy progress bar that does nothing (for when verbose=False)."""
     

{crispyx-0.2.0 → crispyx-0.2.2}/src/crispyx/_disk.py

@@ -11,7 +11,6 @@ resource allocator. There is no override parameter analogous to
 from __future__ import annotations
 
 import shutil
-import warnings
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -194,25 +193,25 @@ def warn_if_disk_space_low(
     ``verbose``-gated confirmation alongside this unconditional warning
     without re-running the free-space lookup.
     """
+    from . import _messages  # _messages imports this module
+
     estimate = assess_bytes(required_bytes, output_path)
     label = f"{context}: " if context else ""
 
     if estimate.free_bytes is not None and not estimate.sufficient:
-        warnings.warn(
+        _messages.emit(
             f"{label}estimated disk usage ({format_bytes(estimate.required_bytes)}) "
             f"leaves less than {min_free_fraction:.0%} free space at {estimate.path} "
             f"({format_bytes(estimate.free_bytes)} currently free). The operation may fail "
             "with 'No space left on device'. Free up space or point output_path/TMPDIR at a "
             "volume with more room before rerunning.",
-            stacklevel=3,
         )
     elif estimate.required_gb > large_file_gb:
-        warnings.warn(
+        _messages.emit(
             f"{label}this operation will write approximately "
             f"{format_bytes(estimate.required_bytes)} to disk. "
             "Streaming keeps memory bounded, but the on-disk footprint (and the time to write "
             "it) is not -- make sure the output volume has room before running large batches.",
-            stacklevel=3,
         )
 
     return estimate

{crispyx-0.2.0 → crispyx-0.2.2}/src/crispyx/_irls.py

@@ -30,6 +30,7 @@ from __future__ import annotations
 from typing import Literal
 
 import numpy as np
+from scipy.linalg import cho_factor, cho_solve
 
 __all__ = [
     "EPS",
@@ -40,11 +41,15 @@ __all__ = [
     "eta_floor",
     "gram_batched",
     "irls_weights",
+    "PIVOT_FLOOR",
     "precondition_columns",
+    "sandwich_variances_2x2",
     "schur_complement",
     "schur_solve",
     "schur_variances",
+    "solve_spd_2x2",
     "working_response",
+    "wls_solve",
 ]
 
 #: Division guard.  Not a modelling floor -- see the module docstring.
@@ -161,6 +166,74 @@ def gram_batched(
     return out
 
 
+#: Pivot floor of the 2x2 factorisation: a singular system gives large but
+#: finite values rather than inf or nan.
+PIVOT_FLOOR = 1e-12
+
+
+def solve_spd_2x2(a, b, c, r0, r1):
+    """Solve ``[[a, b], [b, c]] x = [r0, r1]`` elementwise -- one system per
+    gene when the arguments are arrays -- through its ``LDL'``
+    factorisation, whose pivots ``a`` and ``c - b**2 / a`` are floored at
+    :data:`PIVOT_FLOOR`. Returns ``(x0, x1)``.
+
+    Only arithmetic and ufuncs, so a Numba kernel can compile it too.
+    """
+    d0 = np.maximum(a, PIVOT_FLOOR)
+    lower = b / d0
+    d1 = np.maximum(c - lower * b, PIVOT_FLOOR)
+    x1 = (r1 - lower * r0) / d1
+    return (r0 - b * x1) / d0, x1
+
+
+def sandwich_variances_2x2(m00, m01, m11, ridge):
+    """Sandwich variances ``e_k' H M H e_k`` of a two-column fit's
+    coefficients, for Fisher information ``M = [[m00, m01], [m01, m11]]``
+    and ``H = (M + ridge I)^-1``, with ``H e_k`` from :func:`solve_spd_2x2`.
+    Returns ``(intercept, second column)``.
+    """
+    a, c = m00 + ridge, m11 + ridge
+
+    def quadratic(v0, v1):
+        return v0 * v0 * m00 + 2 * v0 * v1 * m01 + v1 * v1 * m11
+
+    return (quadratic(*solve_spd_2x2(a, m01, c, 1.0, 0.0)),
+            quadratic(*solve_spd_2x2(a, m01, c, 0.0, 1.0)))
+
+
+def wls_solve(
+    x_weighted: np.ndarray, z_weighted: np.ndarray | None = None, ridge: float = 0.0
+) -> tuple[np.ndarray | None, np.ndarray]:
+    """Coefficients and covariance ``(X'WX + ridge I)^-1`` of a weighted
+    least-squares fit, from its weighted design ``W^1/2 X`` (and response
+    ``W^1/2 z``; without one the coefficients are ``None``).
+
+    A Cholesky factor of the normal matrix when it is positive definite;
+    otherwise an SVD of the weighted design, stacked on ``sqrt(ridge) I``,
+    gives the minimum-norm solution and the pseudo-inverse with singular
+    values below ``matrix_rank``'s tolerance dropped -- without forming a
+    normal matrix whose condition number is the design's squared.
+    """
+    n_features = x_weighted.shape[1]
+    gram = x_weighted.T @ x_weighted
+    gram[np.diag_indices(n_features)] += ridge
+    try:
+        factor = cho_factor(gram, check_finite=False)
+    except np.linalg.LinAlgError:
+        design, response = x_weighted, z_weighted
+        if ridge:
+            design = np.vstack([x_weighted, np.sqrt(ridge) * np.eye(n_features)])
+            if z_weighted is not None:
+                response = np.concatenate([z_weighted, np.zeros(n_features)])
+        u, s, vt = np.linalg.svd(design, full_matrices=False)
+        keep = s > s.max(initial=0.0) * max(design.shape) * np.finfo(np.float64).eps
+        scaled = vt[keep].T / s[keep]                       # V S^-1
+        beta = None if response is None else scaled @ (u[:, keep].T @ response)
+        return beta, scaled @ scaled.T
+    beta = None if z_weighted is None else cho_solve(factor, x_weighted.T @ z_weighted, check_finite=False)
+    return beta, cho_solve(factor, np.eye(n_features), check_finite=False)
+
+
 def irls_weights(mu: np.ndarray, alpha: np.ndarray | float | None) -> np.ndarray:
     """Log-link IRLS weights ``mu**2 / V(mu)``.
 
@@ -497,15 +570,17 @@ def schur_variances(
         (H^-1)_XX = S^-1
         (H^-1)_GG = D^-1 + D^-1 B' S^-1 B D^-1
 
-    so the covariate variances come free from the factorisation already
-    computed by :func:`schur_solve` and the group variances need only the
-    quadratic form ``B_k' S^-1 B_k`` per group.
+    so one solve of each ``S`` against ``[I | B]`` gives the covariate
+    variances and the quadratic forms ``B_k' S^-1 B_k`` the group variances
+    need.
     """
     inverse_diagonal = 1.0 / diagonal                        # (a, p)
-    schur_inverse = np.linalg.inv(schur)                     # (p, dx, dx)
-    feature_variance = np.diagonal(schur_inverse, axis1=1, axis2=2)
-
+    n_genes, dx = schur.shape[:2]
     blocks = np.ascontiguousarray(cross.transpose(2, 1, 0))  # (p, dx, a)
-    quadratic = (blocks * np.matmul(schur_inverse, blocks)).sum(axis=1)   # (p, a)
+    # One factorisation of each S, solved for [I | B] together.
+    rhs = np.concatenate([np.broadcast_to(np.eye(dx), (n_genes, dx, dx)), blocks], axis=2)
+    solved = np.linalg.solve(schur, rhs)
+    feature_variance = np.diagonal(solved[:, :, :dx], axis1=1, axis2=2)
+    quadratic = (blocks * solved[:, :, dx:]).sum(axis=1)     # (p, a)
     group_variance = inverse_diagonal.T + quadratic * inverse_diagonal.T**2
     return feature_variance, group_variance