polish_nb

polish_nb(
    Y,
    W,
    alpha,
    psi,
    lambda_a=0.0,
    offset=None,
    absorb=None,
    absorb_batch=None,
    start_cols=None,
    psi_method='profile',
    psi_range=(0.001, 1000.0),
    warm=False,
    maxit=50,
    tol=1e-08,
    engine='batch',
    batch_size=None,
    block_size=None,
    backend='cpu',
    n_jobs=1,
    verbose=False,
)

Converge every gene of a shared NB fit to its own optimum (R polishNB).

fit_nb fits every gene in one IRLS loop with a shared cell weight vector and an aggregate convergence test; for a bright, cell-type-restricted gene that leaves it well short of its own optimum with an inflated dispersion. This takes such a fit and converges each gene separately by damped Newton on that gene’s penalised NB log-likelihood

.. math:: i f{NB}(y_i; _i, ) - _j _j _j^2, = W+ .

With psi_method="profile" the dispersion is then re-estimated by profile ML at the converged mean and the mean re-polished, twice; a gene whose dispersion optimum sits on the search bound keeps its input dispersion (psi_bound). A gene whose start is degenerate (a fitted log-mean below -10 at a cell with a positive count) or whose Newton diverges restarts from a sane point. A gene that cannot be polished from either start keeps its input coefficients and dispersion with polished=False.

The integer-count check reads only the first 20 genes: a cheap guard against an assay that is non-integer throughout, not a scan of every value.

Parameters

Name Type Description Default
Y genes x cells integer counts (dense, scipy sparse or tensor; densified one gene block at a time). required
W cells x p design, or an :func:~spanormpy._block_design.nb_block_design (CPU only; its groups’ blocks are absorbed by construction). required
alpha genes x p starting coefficients, typically fit_nb() output. required
psi starting per-gene dispersions (length nrow(Y), or one value). required
lambda_a ridge penalty: one value or one per column of W (for a block design, also one per column of [X \| Z]). Applied as given, with no scaling. 0.0
offset None, a log-scale offset with one value per cell, or a genes x cells matrix (finite), added to every linear predictor with its coefficient fixed at 1. None
absorb columns the per-gene solver absorbs by a Schur complement: None, a boolean vector (each marked column its own 1x1 block; 0/1 indicators partitioning the cells), or a per-column grouping (NaN/None for a dense column). None
absorb_batch the absorption for the shared-factor batched solver that runs on a device: None or a boolean vector (1x1 blocks only). None passes a boolean absorb through and sends a grouping to the dense batched solver. None
start_cols boolean vector over the columns of W marking indicator columns (such as cell-type intercepts) that the sane start fills with the gene’s log mean over that column’s cells, or None for the overall log mean on the first column. None
psi_method str "profile" or "fixed" (keep the input psi, converge only the mean). 'profile'
psi_range search interval for the profile dispersion; a gene whose optimum falls on either end keeps its input dispersion. (0.001, 1000.0)
warm bool alpha/psi are an already converged fit at a nearby penalty: a few Newton steps at the held dispersion, no dispersion search and no restart check. False
maxit int Newton iteration cap and relative log-likelihood tolerance, per gene. 50
tol int Newton iteration cap and relative log-likelihood tolerance, per gene. 50
engine str "batch" (default) runs a batch of genes through each Newton step; "gene" is the per-gene reference. The batched profile dispersion is a bisection and the per-gene one Brent’s method, so they agree on psi to optimize()’s tolerance. 'batch'
batch_size int | None genes per batched Newton, or None to size it from spanormpy.settings.polish_mem_budget (bytes per worker, default 1e9). None
block_size int | None genes per block (the unit of densification and dispatch), or None for at least one block per worker and at most 2,000 genes. None
backend str "cpu" (default), "auto" or "gpu". A device needs engine="batch" and float64, and forces serial dispatch. 'cpu'
n_jobs int threads over gene blocks. 1
verbose bool report progress on stderr. False

Returns

Name Type Description
dict alpha (genes x p), psi and loglik (the penalised log-likelihood at the returned fit, NaN for a gene that was not polished), and polish, a DataFrame with one row per gene: iterations, psi_fitnb, restarted, capped, singular, psi_bound and polished.

Examples

import numpy as np
import spanormpy as sn
rng = np.random.default_rng(1)
Y = rng.poisson(5, size=(20, 50))  # counts, genes x cells
W = np.column_stack([np.ones(50), np.linspace(-1, 1, 50)])  # intercept + one covariate
fit = sn.fit_nb(Y, W, verbose=False, backend="cpu")
pol = sn.polish_nb(Y, W, fit["alpha"], fit["psi"])
pol["polish"].head()  # per-gene convergence diagnostics
iterations psi_fitnb restarted capped singular psi_bound polished
0 2 0.003726 False False False True True
1 4 0.000617 False False False False True
2 2 0.001847 False False False True True
3 2 0.000098 False False False True True
4 4 0.004553 False False False False True