crispyx 0.2.1__tar.gz → 0.2.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: crispyx
-Version: 0.2.1
+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
@@ -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.1 → crispyx-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "crispyx"
-version = "0.2.1"
+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"}
@@ -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.1 → 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.1 → crispyx-0.2.2}/src/crispyx/_checkpoint.py

@@ -13,7 +13,6 @@ import json
 import logging
 import os
 import shutil
-import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
 
@@ -21,6 +20,7 @@ import h5py
 import numpy as np
 import pandas as pd
 
+from . import _messages
 from ._memory import drop_file_cache
 
 if TYPE_CHECKING:
@@ -304,6 +304,79 @@ def _sha256(array: np.ndarray) -> str:
     return hashlib.sha256(np.ascontiguousarray(array).tobytes()).hexdigest()
 
 
+#: 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
+
+
 class ResumableRun:
     """A result file filled in place, and a progress checkpoint, that outlive
     an interruption.
@@ -353,10 +426,9 @@ 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 _has_datasets(self.work_path, datasets)

{crispyx-0.2.1 → 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.1 → 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