destriper 0.1.0__py3-none-any.whl

destriper/__init__.py

@@ -0,0 +1,45 @@
+"""destriper — estimate and correct row/column multiplicative stripe
+artifacts in high-resolution spatial transcriptomics count data.
+
+Numerical core::
+
+    import destriper as ds
+    result = ds.fit(counts, row_indices, column_indices, nucl_labels)
+    corrected_totals = ds.destripe_tot_counts(tot_counts, rows, cols, labels, result)
+    corrected_matrix, achieved = ds.rescale(count_matrix, corrected_totals)
+
+Optional AnnData layer (``pip install destriper[anndata]``)::
+
+    result = ds.fit_adata(adata, nucl_key="nucleus_id")
+    ds.destripe_adata(adata, result, target_layer="destriped")
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from destriper._core import fit
+from destriper._destripe import destripe_tot_counts, rescale
+from destriper._result import FitResult
+
+if TYPE_CHECKING:
+    from destriper._adata import destripe_adata, fit_adata
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "fit",
+    "destripe_tot_counts",
+    "rescale",
+    "FitResult",
+    "fit_adata",
+    "destripe_adata",
+]
+
+
+def __getattr__(name):
+    # Lazy AnnData layer: import (and require) anndata only on first use.
+    if name in ("fit_adata", "destripe_adata"):
+        from destriper import _adata
+
+        return getattr(_adata, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

destriper/_adata.py

@@ -0,0 +1,141 @@
+"""Optional AnnData convenience layer.
+
+Requires ``anndata`` (``pip install destriper[anndata]``). Imported lazily
+by ``destriper.__getattr__`` so the numerical core has no hard dependency
+on anndata.
+"""
+import logging
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+
+from destriper._core import fit as _fit
+from destriper._destripe import (
+    destripe_factors,
+    destripe_tot_counts as _destripe_tot_counts,
+    rescale as _rescale,
+)
+from destriper._result import FitResult
+
+logger = logging.getLogger("destriper")
+
+
+def _require_anndata():
+    try:
+        import anndata  # noqa: F401
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            "The AnnData layer requires `anndata`. "
+            "Install it with `pip install destriper[anndata]`."
+        ) from e
+
+
+def _totals(adata, count_key):
+    if count_key is not None:
+        return np.asarray(adata.obs[count_key].to_numpy(), dtype=float)
+    return np.asarray(adata.X.sum(axis=1)).ravel().astype(float)
+
+
+def fit_adata(
+    adata,
+    *,
+    nucl_key: str = "nucleus_id",
+    count_key: str | None = "total_counts",
+    row_key: str = "array_row",
+    col_key: str = "array_col",
+    cv_group_key: str | None = None,
+    max_iter_theta: int = 5,
+    max_iter: int = 100_000,
+    cv="spatial",
+) -> FitResult:
+    """Fit the stripe GLM from an AnnData.
+
+    Uses only nuclear bins (rows with a non-null ``obs[nucl_key]``); unlabelled
+    bins are ignored. ``count_key`` supplies per-bin totals from ``.obs``; if
+    ``None`` they are computed from ``adata.X``. See :func:`destriper.fit`
+    for the other parameters. ``cv_group_key`` names an ``.obs`` column of CV
+    group ids (overrides ``cv``).
+    """
+    _require_anndata()
+    obs = adata.obs
+    labels = obs[nucl_key].to_numpy()
+    is_nucl = ~pd.isna(labels)
+    n_total = len(labels)
+    n_nucl = int(is_nucl.sum())
+    if n_nucl == 0:
+        raise ValueError(f"No nuclear bins: all values in obs[{nucl_key!r}] are null.")
+    if n_nucl < n_total:
+        logger.info(
+            "Fitting on %d nuclear bins (of %d total); %d unlabelled bins ignored.",
+            n_nucl,
+            n_total,
+            n_total - n_nucl,
+        )
+
+    counts = _totals(adata, count_key)[is_nucl]
+    rows = np.asarray(obs[row_key].to_numpy())[is_nucl]
+    cols = np.asarray(obs[col_key].to_numpy())[is_nucl]
+    nucl_labels = labels[is_nucl]
+
+    cv_arg = cv
+    if cv_group_key is not None:
+        cv_arg = np.asarray(obs[cv_group_key].to_numpy())[is_nucl]
+
+    return _fit(
+        counts,
+        rows,
+        cols,
+        nucl_labels,
+        max_iter_theta=max_iter_theta,
+        max_iter=max_iter,
+        cv=cv_arg,
+    )
+
+
+def destripe_adata(
+    adata,
+    result: FitResult,
+    *,
+    source_layer: str | None = None,
+    target_layer: str = "destriped",
+    nucl_key: str = "nucleus_id",
+    row_key: str = "array_row",
+    col_key: str = "array_col",
+) -> None:
+    """Apply a fitted :class:`FitResult` to an AnnData in place.
+
+    Writes:
+
+    - ``adata.layers[target_layer]`` — the destriped count matrix
+    - ``adata.obs["ds_destripe_factor"]`` — the per-bin stripe factor
+    - ``adata.obs["ds_corrected_counts"]`` — the target corrected totals from
+      :func:`destriper.destripe_tot_counts`. Note this is the *target*, not
+      the rescaled-matrix row-sums (``layers[target_layer].sum(1)``); the two
+      differ on originally-empty bins (see :func:`destriper.rescale`).
+    - ``adata.uns["destriper"]["result"]`` — the serialized fit
+      (:meth:`FitResult.to_uns`): ``row_factors``, ``col_factors`` and
+      ``nucl_concentration`` as one-column DataFrames plus the scalar fit
+      parameters. Rebuild it with :meth:`FitResult.from_uns`.
+    """
+    _require_anndata()
+    obs = adata.obs
+    rows = np.asarray(obs[row_key].to_numpy())
+    cols = np.asarray(obs[col_key].to_numpy())
+    labels = obs[nucl_key].to_numpy()
+
+    X = adata.X if source_layer is None else adata.layers[source_layer]
+    if not sp.issparse(X):
+        X = sp.csr_matrix(X)
+
+    tot = np.asarray(X.sum(axis=1)).ravel().astype(float)
+    corrected = _destripe_tot_counts(tot, rows, cols, labels, result)
+    X_corr, _achieved = _rescale(X, corrected)
+
+    adata.layers[target_layer] = X_corr
+    adata.obs["ds_destripe_factor"] = destripe_factors(
+        result.row_factors, result.col_factors, rows, cols
+    )
+    adata.obs["ds_corrected_counts"] = corrected
+    uns = adata.uns.setdefault("destriper", {})
+    uns["result"] = result.to_uns()

destriper/_core.py

@@ -0,0 +1,167 @@
+import logging
+import warnings
+from contextlib import contextmanager
+
+import numpy as np
+import pandas as pd
+
+from destriper._df import build_fit_df, strip_id_prefix, warn_large_regions
+from destriper._result import FitResult
+
+logger = logging.getLogger("destriper")
+
+_V6_MODEL_ARGS = dict(
+    init_method="ones",
+    init_method_args={"c_mean": True},
+    family="nbinom",
+    family_params={"theta": 1.0},
+    n_alphas=10,
+    alphas=None,
+    P2="hw_only",
+    P1="hw_only",
+    freeze_c=False,
+    sklearnCV=True,
+    sklearnCV_one_SE_rule=False,
+    warm_start_alpha=True,
+    fit_theta_iter=True,
+    fit_theta_iter_loc="out",
+    selection="random",
+    solver="coordinate_descent",
+    parallel=True,
+    reestimate_c_on_val=True,
+)
+
+
+@contextmanager
+def _quiet_fit_warnings():
+    """Silence a few known-benign warnings from the vendored CV code.
+
+    The cross-validation ``c`` re-estimation divides by per-group sums that can
+    be zero on some folds (result discarded via ``np.where``), and a pandas
+    groupby-apply triggers a deprecation warning. Neither affects results; we
+    scope the filters by message so genuine warnings still surface.
+    """
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", message=".*divide by zero encountered.*", category=RuntimeWarning
+        )
+        warnings.filterwarnings(
+            "ignore", message=".*invalid value encountered.*", category=RuntimeWarning
+        )
+        warnings.filterwarnings(
+            "ignore",
+            message=".*DataFrameGroupBy.apply operated on the grouping columns.*",
+            category=FutureWarning,
+        )
+        yield
+
+
+def _groups_to_splits(groups):
+    """Leave-one-group-out CV splits (positional) from a per-bin group array."""
+    groups = np.asarray(groups)
+    all_pos = np.arange(len(groups))
+    splits = [
+        (np.setdiff1d(all_pos, test, assume_unique=False), test)
+        for g in pd.unique(groups)
+        if len((test := all_pos[groups == g])) > 0
+    ]
+    if len(splits) < 2:
+        raise ValueError("cv group array must define at least 2 non-empty groups.")
+    return splits
+
+
+def _resolve_cv(cv, n_bins):
+    if isinstance(cv, str):
+        if cv not in ("spatial", "default"):
+            raise ValueError(
+                f"cv string must be 'spatial' or 'default', got {cv!r}."
+            )
+        return cv
+    if isinstance(cv, (int, np.integer, bool)) and not isinstance(cv, bool):
+        n = int(cv)
+        if n < 2:
+            raise ValueError("integer cv (number of KFold folds) must be >= 2.")
+        return n
+    # per-bin group-id array -> leave-one-group-out
+    groups = np.asarray(cv)
+    if groups.ndim != 1 or len(groups) != n_bins:
+        raise ValueError(
+            "cv group array must be 1-D with one entry per bin "
+            f"(expected length {n_bins}, got {groups.shape})."
+        )
+    return _groups_to_splits(groups)
+
+
+def fit(
+    counts,
+    row_indices,
+    column_indices,
+    nucl_labels,
+    *,
+    max_iter_theta: int = 5,
+    max_iter: int = 100_000,
+    cv="spatial",
+) -> FitResult:
+    """Fit the row/column stripe GLM on per-bin total counts of nucleus bins.
+
+    Fits ``k_ij ~ NB(mean = c_p * h_i * w_j)`` (log link) with the tuned v6
+    recipe. Only nucleus (labelled) bins are used; the returned factors are
+    later applied to all bins via :func:`destripe_tot_counts`.
+
+    Parameters
+    ----------
+    counts : 1-D int array
+        Per-bin total counts (nucleus bins only).
+    row_indices, column_indices : 1-D int arrays
+        Array-grid row / column index per bin.
+    nucl_labels : 1-D array
+        Nucleus id per bin; must be non-null.
+    max_iter_theta : int, default 5
+        Cap on the dispersion (theta) iteration.
+    max_iter : int, default 100000
+        Cap on the GLM solver iterations per fit.
+    cv : "spatial" | "default" | int | 1-D array, default "spatial"
+        Cross-validation folds for regularization selection: ``"spatial"``
+        (4 grid quadrants, leave-one-quadrant-out), ``"default"`` (5-fold
+        KFold), an int (KFold with that many folds), or a per-bin group-id
+        array (leave-one-group-out).
+
+    Returns
+    -------
+    FitResult
+    """
+    df = build_fit_df(counts, row_indices, column_indices, nucl_labels)
+    warn_large_regions(df, threshold=100)
+    cv_split_method = _resolve_cv(cv, n_bins=len(df))
+
+    # Lazy import: pulls in glum/tabmat only when actually fitting.
+    from destriper._fit.glum_wrapper import GlumWrapper
+
+    model = GlumWrapper(
+        data_df=df,
+        cv_split_method=cv_split_method,
+        fit_theta_max_iter=int(max_iter_theta),
+        max_iter=int(max_iter),
+        **_V6_MODEL_ARGS,
+    )
+    with _quiet_fit_warnings():
+        model.fit()
+
+    sol = model.sol_
+    status = model.status_dict
+
+    concentration = sol.c.copy()
+    concentration.index = pd.Index(
+        [strip_id_prefix(x) for x in concentration.index], name="nucleus"
+    )
+
+    return FitResult(
+        row_factors=sol.h,
+        col_factors=sol.w,
+        nucl_concentration=concentration,
+        dispersion=float(model.theta),
+        converged=bool(status["converged"]),
+        theta_converged=bool(status["theta_iter_converged_"]),
+        theta_iterations=int(status["theta_iter_"]),
+        alpha=float(model.alpha),
+    )

destriper/_destripe.py

@@ -0,0 +1,123 @@
+"""Total-count destriping arithmetic (numpy / scipy.sparse, no AnnData).
+
+Reimplements the correction logic from the source ``spatialAdata`` methods
+(``get_destripe_factors``, ``_n_counts_adjusted_dividing``,
+``_n_counts_adjusted_qm``) and ``destripe_counts`` as plain free functions.
+"""
+import logging
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+
+from destriper._quantile_matching import get_qm_fun
+
+logger = logging.getLogger("destriper")
+
+
+def destripe_factors(row_factors, col_factors, row_indices, column_indices):
+    """Per-bin multiplicative stripe factor ``f = h[row] * w[col]``.
+
+    Rows/cols absent from the fitted factors get factor ``1.0`` (left
+    unchanged). Mirrors the source ``get_destripe_factors``.
+    """
+    row_indices = np.asarray(row_indices)
+    column_indices = np.asarray(column_indices)
+    row_filled = row_factors.reindex(np.unique(row_indices), fill_value=1.0)
+    col_filled = col_factors.reindex(np.unique(column_indices), fill_value=1.0)
+    return row_filled.loc[row_indices].values * col_filled.loc[column_indices].values
+
+
+def _adjust_dividing(k, f):
+    """Cytoplasm bins: ``adj = k / f`` with the source's guards."""
+    k = np.asarray(k, dtype=float)
+    f = np.asarray(f, dtype=float)
+    with np.errstate(divide="ignore", invalid="ignore"):
+        adj = k / f
+    inf_sel = np.isinf(adj)  # f == 0 and k > 0 -> keep original count
+    adj[inf_sel] = k[inf_sel]
+    zero_zero = (k == 0) & (f == 0)  # 0 / 0 -> 0
+    adj[zero_zero] = 0.0
+    return adj
+
+
+def destripe_tot_counts(tot_counts, row_indices, column_indices, nucl_labels, result):
+    """Compute per-bin destriped (target) total counts for *all* bins.
+
+    Nucleus bins — those whose label is in ``result.nucl_concentration`` — are
+    corrected by negative-binomial quantile matching against ``mu = c_p`` using
+    the fitted dispersion. Every other bin (label unseen in the fit, or null =
+    cytoplasm) is corrected by plain division ``k / f``. Labels not seen in the
+    fit are logged and treated as non-nuclei.
+    """
+    k = np.asarray(tot_counts, dtype=float)
+    n = len(k)
+    row_indices = np.asarray(row_indices)
+    column_indices = np.asarray(column_indices)
+    labels = np.asarray(nucl_labels, dtype=object)
+    if not (len(row_indices) == len(column_indices) == len(labels) == n):
+        raise ValueError(
+            "tot_counts, row_indices, column_indices, nucl_labels must have equal length."
+        )
+
+    f = destripe_factors(
+        result.row_factors, result.col_factors, row_indices, column_indices
+    )
+
+    # Canonical string labels (None for null); route by membership in fitted nuclei.
+    c = result.nucl_concentration
+    is_null = pd.isna(labels)
+    canon = np.array(
+        [None if is_null[i] else str(labels[i]) for i in range(n)], dtype=object
+    )
+    mu = pd.Series(canon).map(c)  # c_p for nucleus bins, NaN otherwise
+    is_nucleus = mu.notna().values
+
+    unseen = sorted(
+        {lab for lab, nn in zip(canon, ~is_null) if nn} - set(c.index)
+    )
+    if unseen:
+        shown = ", ".join(repr(u) for u in unseen[:20])
+        more = "" if len(unseen) <= 20 else f", ... (+{len(unseen) - 20} more)"
+        logger.info(
+            "Considering %d label(s) not seen in the fit as non-nuclei (cytoplasm): %s%s",
+            len(unseen),
+            shown,
+            more,
+        )
+
+    adj = np.empty(n, dtype=float)
+    cyto = ~is_nucleus
+    adj[cyto] = _adjust_dividing(k[cyto], f[cyto])
+    if is_nucleus.any():
+        qm = get_qm_fun("nbinom", dist_params={"r": float(result.dispersion)})
+        adj[is_nucleus] = qm(
+            k[is_nucleus], mu.values[is_nucleus].astype(float), f[is_nucleus]
+        )
+    return adj
+
+
+def rescale(count_matrix, corrected_tot_counts):
+    """Row-scale a sparse count matrix so each bin sums to its corrected total.
+
+    Returns ``(corrected_matrix, achieved_tot_counts)``. Originally-empty rows
+    stay empty (no scalar can make a zero row sum to a positive target), so
+    ``achieved_tot_counts`` differs from ``corrected_tot_counts`` exactly on
+    those rows. Mirrors the source ``destripe_counts``.
+    """
+    if not sp.issparse(count_matrix):
+        raise TypeError("count_matrix must be a scipy.sparse matrix.")
+    X = count_matrix.tocsr().astype(float)
+    corrected = np.asarray(corrected_tot_counts, dtype=float)
+    if X.shape[0] != len(corrected):
+        raise ValueError(
+            f"count_matrix has {X.shape[0]} rows but corrected_tot_counts has "
+            f"length {len(corrected)}."
+        )
+    original = np.asarray(X.sum(axis=1)).ravel()
+    with np.errstate(divide="ignore", invalid="ignore"):
+        scale = corrected / original
+    scale[original == 0] = 1.0  # leave empty rows unscaled (they stay zero)
+    X_out = sp.diags(scale).dot(X)
+    achieved = np.asarray(X_out.sum(axis=1)).ravel()
+    return X_out, achieved

destriper/_df.py

@@ -0,0 +1,79 @@
+import logging
+
+import numpy as np
+import pandas as pd
+
+logger = logging.getLogger("destriper")
+
+_ID_PREFIX = "id_"
+
+
+def _as_1d(name, arr, n=None):
+    arr = np.asarray(arr)
+    if arr.ndim != 1:
+        raise ValueError(f"{name} must be 1-D, got shape {arr.shape}.")
+    if n is not None and len(arr) != n:
+        raise ValueError(
+            f"{name} has length {len(arr)} but expected {n} (same as counts)."
+        )
+    return arr
+
+
+def build_fit_df(counts, row_indices, column_indices, nucl_labels):
+    """Assemble ``df[["i", "j", "p", "k"]]`` for :func:`fit_GLM_glum`.
+
+    - ``i`` = ``row_indices`` cast to int
+    - ``j`` = ``column_indices`` cast to int
+    - ``p`` = ``"id_" + str(label)`` as an ``object`` column (never NaN)
+    - ``k`` = ``counts`` cast to int
+
+    Every bin is assumed nuclear; a null label raises (pass only nuclear bins).
+    """
+    counts = _as_1d("counts", counts)
+    n = len(counts)
+    row_indices = _as_1d("row_indices", row_indices, n)
+    column_indices = _as_1d("column_indices", column_indices, n)
+    nucl_labels = _as_1d("nucl_labels", np.asarray(nucl_labels, dtype=object), n)
+
+    if pd.isna(list(nucl_labels)).any():
+        raise ValueError(
+            "nucl_labels must not contain NaN/None. fit() expects only nuclear "
+            "bins; filter unassigned (cytoplasm) bins out before fitting."
+        )
+
+    df = pd.DataFrame(
+        {
+            "i": row_indices.astype(int),
+            "j": column_indices.astype(int),
+            "p": np.array([_ID_PREFIX + str(x) for x in nucl_labels], dtype=object),
+            "k": counts.astype(int),
+        }
+    )[["i", "j", "p", "k"]]
+    return df
+
+
+def warn_large_regions(df, threshold=100):
+    """Warn about nucleus labels with more than ``threshold`` bins.
+
+    At 2 um bin resolution a nucleus spanning >100 bins usually signals a
+    segmentation merge error.
+    """
+    sizes = df.groupby("p", observed=True).size()
+    big = sizes[sizes > threshold]
+    if len(big) > 0:
+        labels = [strip_id_prefix(p) for p in big.index[:20].tolist()]
+        more = "" if len(big) <= 20 else f" (and {len(big) - 20} more)"
+        logger.warning(
+            "%d nucleus label(s) have more than %d bins%s: %s; "
+            "this often indicates a segmentation merge error.",
+            len(big),
+            threshold,
+            more,
+            ", ".join(f"{lab!r} ({int(sizes.loc[_ID_PREFIX + str(lab)])})" for lab in labels),
+        )
+
+
+def strip_id_prefix(p):
+    """Inverse of the ``"id_"`` prefixing applied in :func:`build_fit_df`."""
+    s = str(p)
+    return s[len(_ID_PREFIX):] if s.startswith(_ID_PREFIX) else s

destriper/_fit/_pandas_log.py

@@ -0,0 +1,7 @@
+import textwrap
+
+
+def log_df(df, logger):
+    df_str = df.to_string()
+    indented = textwrap.indent(df_str, "    ")
+    logger("DataFrame:\n%s", indented)