rafinat 0.3.0

Bayesian normalizations for RNA-seq

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.

transformer	what it is	default operating point
compost	the DM/Liouville model	p=0 digamma log-corner, top-2/3-trimmed per-cell centering
hellnorm	the Hellinger / Fisher-Rao spherical (sqrt) normalization	top-2/3-trimmed per-cell sqrt reference
rafinat	the complete Fisher-Rao Liouville embedding	method="arclength": composition arc length (sharp r_snr) + radial cell-size coordinate, beta="total"

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 benchmark-winning defaults 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.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_, rafinat.r_comp_ / rafinat.ref_ and — for method="arclength" — rafinat.C_ / rafinat.beta_, or — for method="lorentzian" — rafinat.ref_A_ (the frozen sigma origin) plus rafinat.signature_ after transform (the per-row +1/-1 spacelike/timelike sign).

The two rafinat embeddings

rafinat's Fisher-Rao distance on the Dirichlet part is the great-arc (product) upper bound, which drops the rank-1 coupling carrying the total concentration A = Σ(x + r) ("n" / depth). The method knob chooses how that is handled:

• method="arclength" (default) — the great-arc composition arc length on depth-normalized counts, stacked with the explicit radial cell-size row sqrt(C)·ln β. All coordinates are spacelike, so ordinary L2 reproduces the distance. transform → (genes + 1, cells).
• method="lorentzian" — the coupling-aware embedding on raw counts, which appends the timelike total coordinate σ = ∫₀ᴬ √ψ'(s) ds the great-arc drops (a strictly tighter upper bound). Here σ carries depth, so it replaces the explicit β row. Because σ is timelike, the chart is pseudo-Euclidean — pick how it is handed off with coords:
  • coords="signed" (default) — the raw chart (genes + 1, cells); the last 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.

Depth-normalizing first would make every A equal and collapse σ to a constant, so method="lorentzian" embeds the raw counts (that is why it drops the depth-normalized-plus-β stack the arc-length method uses).

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). Dominates the classic uniform/extrinsic/Frechet references.
• rafinat — the Fisher-Rao Liouville distance factors as d² = d_composition² + C·(d ln β)²; method="arclength" realises it as the stack [ composition arc length (depth-normalized, sharp r_snr) ; sqrt(C)·ln(β_c) ], with the radial weight C = genes · r_mle decoupled from the (sharp) directional concentration so the cell-size axis stays alive at a parameter-free weight. method="lorentzian" instead uses the exact Gamma–Dirichlet split g_D = g_Γ − dσ⊗dσ: the great-arc coordinates minus the timelike total arc length σ = ∫₀ᴬ √ψ'(s) ds (liudist.lorentzian_embedding), a strictly tighter upper bound on d_composition that keeps the total/depth σ the great-arc discards (so σ stands in for the β row). coords="euclidean" then classical-MDS's that pseudo-Euclidean distance back to plain L2 coordinates.