rafinat

Dirichlet-Multinomial posterior normalizations for single-cell RNA-seq count data.

Three count normalizations derived from a Dirichlet-Multinomial posterior (Liouville) view of the latent expression rates.

Each is a scikit-learn-style transformer: you instantiate it with its operating point, then .fit / .transform / .fit_transform a count matrix — so it drops into a preprocessing pipeline.

All take a dense (genes, cells) array of raw counts (genes = features/rows, cells = samples/columns) and return (genes, cells) (rafinat returns (genes + 1, cells) — the extra row is the cell-size scale coordinate; see the two rafinat embeddings for the shapes of method="lorentzian"). compost and hellnorm are pure numpy/scipy (CPU). rafinat uses the optional liudist package (Fisher-Rao geometry of Liouville laws; pulls in JAX, and benefits from a GPU JAX build) — installed via the [rafinat] extra.

import numpy as np, rafinat
X = np.random.poisson(0.5, size=(2000, 500)).astype(float)   # genes x cells

Z = rafinat.compost().fit_transform(X)      # compost p=0, trimmed
H = rafinat.hellnorm().fit_transform(X)     # Hellnorm, trimmed reference
R = rafinat.rafinat().fit_transform(X)      # rafinat; (genes + 1) x cells

fit / transform

fit estimates and freezes the two data-driven pieces — the concentration r and the per-cell reference level beta — and transform re-applies that frozen normalization. Because beta is per-cell, transform expects a matrix of the same (genes, cells) shape it was fitted on (re-fit for a different gene/cell set). fit_transform(X) is the one-shot form.

tr = rafinat.compost(p=0.5).fit(X)     # estimate & freeze tr.r_ and tr.beta_
Z  = tr.transform(X)                    # apply; == tr.fit_transform(X)

from sklearn.pipeline import Pipeline   # optional — also works without scikit-learn installed
pipe = Pipeline([("normalize", rafinat.compost())])
Z = pipe.fit_transform(X)

scikit-learn is an optional extra: if installed, the transformers inherit BaseEstimator / TransformerMixin (full Pipeline / clone / get_params support); otherwise a light built-in shim provides the same fit / transform / fit_transform / get_params API.

Choosing the operating point

The defaults that did best in the paper's benchmarks are baked in, but every knob is a constructor argument:

rafinat.compost(p=0.0)               # default: digamma log-corner, trimmed centering
rafinat.compost(p=0.5)               # posterior sqrt mean (order-1/2 power-mean cell size)
rafinat.compost(trim=0.0)            # ordinary (non-trimmed) per-cell centering
rafinat.compost(r="mle")             # pooled DM-MLE concentration instead of the isscr-matched r
rafinat.compost(zscore=True)         # + per-gene z-score (the optional '->Z' standardization)
rafinat.compost(p=0.5, beta="varest")   # per-cell beta = Laplace-MLE Gamma rate of the varest model
rafinat.compost(p=0.5, emulate_m=True)  # report the varest model's posterior expectation instead

rafinat.hellnorm(reference="uniform")     # classic log-map references: uniform / extrinsic / frechet
rafinat.hellnorm(trim=0.5)                # lighter top-trim

rafinat.rafinat(beta="atop10")       # simulation-leaning cell-size estimator
rafinat.rafinat(r_comp="mle")        # smoother directional concentration
rafinat.rafinat(zcomp=True)          # + per-gene z-score of the composition rows ('-> coordZ')

rafinat.rafinat(method="lorentzian")                     # coupling-aware embedding (see below)
rafinat.rafinat(method="lorentzian", coords="euclidean") # + classical-MDS to L2-ready coordinates

Each of these returns a transformer; call .fit_transform(X) (or .fit(X) then .transform(X)) on it. After fitting, the estimated values are exposed as fitted attributes (trailing underscore): compost.r_ / compost.beta_, and rafinat.r_comp_ / rafinat.ref_ / rafinat.C_ / rafinat.beta_ (both methods) — plus, for method="lorentzian", rafinat.ref_A_ (the frozen sigma origin) and rafinat.signature_ after transform (the per-row +1/-1 spacelike/timelike sign).

The two rafinat embeddings

Both methods realise the full Liouville distance d² = d_x² + C·(Δ ln β)² — a Dirichlet (composition) part d_x plus the radial cell-size row sqrt(C)·ln β (β = the tunable, robust cell size; C = genes · r_mle). They differ in how they approximate d_x, whose great-arc (product) form drops the rank-1 coupling carrying the total concentration A = Σ(x + r) — the posterior sharpness, distinct from the cell-size scale β:

• method="arclength" (default) — the great-arc composition arc length on depth-normalized counts (which zeroes the coupling by construction), stacked with the sqrt(C)·ln β row. All coordinates are spacelike, so ordinary L2 reproduces the distance. transform → (genes + 1, cells).
• method="lorentzian" — the coupling-aware d_x on raw counts: it appends the timelike total-concentration coordinate σ = ∫₀ᴬ √ψ'(s) ds the great-arc drops (a strictly tighter upper bound on d_x). σ is not the cell-size axis — it is the concentration coupling inside d_x, so it sits alongside the sqrt(C)·ln β cell-size row, not in place of it. The stack is [ arc ; σ (timelike) ; sqrt(C)·ln β ]. Because σ is timelike the chart is pseudo-Euclidean — pick how it is handed off with coords:
  • coords="signed" (default) — the raw chart (genes + 2, cells); the sigma row is timelike (see signature_). Cheap and batched (batch_cells), but the downstream method must honour the signature — plain L2 on it is not the Lorentzian distance.
  • coords="euclidean" — classical-MDS the Lorentzian distance into true Euclidean coordinates (k, cells) (k ≤ cells − 1) whose ordinary L2 reproduces it — a drop-in for an L2 PCA→kNN pipeline. This step is O(cells²) to build the matrix + O(cells³) to diagonalise and is transductive (no frozen basis), so prefer the one-shot fit_transform.

method="lorentzian" embeds the raw counts on purpose: depth-normalizing first would make every A equal and collapse σ to a constant (identical to arclength). So it keeps the concentration variation arclength normalizes away, while still carrying the same β cell-size row.

The optional per-gene standardization (compost(zscore=True), rafinat(zcomp=True)) is disabled by default, matching the benchmark's shipped defaults.

Install

pip install rafinat                # compost / hellnorm (numpy + scipy only)
pip install "rafinat[rafinat]"     # + the rafinat() method — adds the liudist backend (pulls in JAX)

compost and hellnorm need only numpy/scipy; the heavy liudist + JAX stack is pulled in only by the [rafinat] extra, i.e. only if you use the rafinat() method.

Method provenance

• compost — digamma(r + x) with r = 1/(4·alpha) + 1/2 (so digamma(r+x) ≈ log(x + 1/(4·alpha))), minus a per-cell location estimated on the low-expression bulk (the top high-expression genes — the biologically variable ones — are dropped from the cell-size estimate).
• hellnorm — sqrt(x / sum x) minus a top-trimmed per-cell mean (the sqrt-geometry analog of compost's trimmed centering). Ahead of the classic uniform/extrinsic/Frechet references in the benchmarks.
• rafinat — the Fisher-Rao Liouville distance factors as d² = d_x² + C·(d ln β)² (a Dirichlet part d_x + the radial cell-size sqrt(C)·ln β, C = genes · r_mle decoupled from the sharp directional concentration so the cell-size axis stays alive at a parameter-free weight). method="arclength" takes d_x as the great-arc arc length on depth-normalized counts; method="lorentzian" instead uses the exact Gamma–Dirichlet split g_D = g_Γ − dσ⊗dσ, keeping the timelike total-concentration arc length σ = ∫₀ᴬ √ψ'(s) ds (liudist.lorentzian_embedding) the great-arc discards — a strictly tighter estimate of d_x. σ is the concentration coupling inside d_x, so it sits alongside the same sqrt(C)·ln β cell-size row, giving the stack [ arc ; σ (−) ; sqrt(C)·ln β (+) ]; coords="euclidean" then classical-MDS's that pseudo-Euclidean distance back to plain L2 coordinates.