spacr.power_model¶

Bayesian sparse-regression inference for pooled-screen power analysis.

This is the inference half of a Python port of spaCRPower (https://github.com/maomlab/spaCRPower), by Matthew O’Meara (maom@umich.edu, ORCID 0000-0002-3128-5331), released under the MIT licence. The MIT licence requires the copyright notice be retained, so:

spaCRPower — Copyright (c) Matthew O’Meara. Released under the MIT licence. Ported to Python for spaCR; the statistical model (Poisson regression on log10 read-fraction with a horseshoe prior on the per-gene coefficients, scored by AUROC / average precision against ground-truth hit status) is his. Errors in the translation are ours.

The simulator half lives in spacr.power_simulate; this module only consumes the tidy (gene, well) table it produces and never imports it at module scope.

The question this module answers¶

How many wells, how many cells per well, and how good a classifier do I need before a pooled CRISPR screen can actually find its hits? You answer it by simulating screens you know the truth for, fitting the model you would really use, and measuring how well the fit recovers the hits you planted. scan_parameters() sweeps that loop over a design grid.

The model¶

One row per well. Npositive positive cells out of Ntotal imaged, with a per-gene covariate log10expression — the log10 fraction of the well’s sequencing reads assigned to that gene:

Npositive_w ~ Poisson(exp(alpha + sum_g beta_g * x_wg + log(Ntotal_w)))

beta_g is the evidence that gene g is a hit: higher means the wells carrying more of gene g had more positive cells than their cell count alone explains. There are far more genes than wells (p >> n), so beta is given a regularized horseshoe prior (Piironen & Vehtari, Electron. J. Statist. 11(2), 2017) — heavy tails so genuine hits escape shrinkage, a sharp spike at zero so the hundreds of non-hits collapse onto it, and a Student-t slab so the tails stay proper.

Backends¶

The R original fits this with brms + cmdstanr (full NUTS). There is no Python drop-in: cmdstanpy needs a C++ toolchain at install time, which is not acceptable for a pip-installed package. So the backend is pluggable:

"torch" (default, always available): Mean-field ADVI — automatic-differentiation variational inference — in a non-centred parameterisation, using the torch spaCR already depends on. Zero new dependencies, seconds not minutes, and it will run on the GPU if asked.
"numpyro" / "pymc" (optional, lazily imported): Exact NUTS, if the user has installed either. Never added to install_requires.
"auto": numpyro, else pymc, else torch — and it says which, both in the log and in PowerFit.backend, so a result can never be mistaken for one produced by a different method. Pin backend= explicitly if you need two machines to agree bit for bit.

What ADVI costs you versus NUTS¶

ADVI is an optimisation, not a sampler, and mean-field means the approximating family is a product of independent Gaussians in the unconstrained space. Concretely:

Intervals are too narrow. Mean-field VI systematically underestimates posterior variance, because it cannot represent posterior correlations and pays no penalty for ignoring them. Treat q5/q95 from the torch backend as a spread indicator, not as a calibrated 90 % interval, and never quote them as one in a paper.
No convergence diagnostics worth the name. There is no R-hat, no effective sample size. All you get is “did the ELBO stop moving”, which is reported as PowerFit.converged and is a much weaker claim than NUTS convergence.
It can land in a local optimum, and the horseshoe’s funnel geometry is exactly the sort of thing that causes it.
The posterior mean of beta — the ranking statistic — is the part that survives. AUROC and average precision depend only on the order of the beta estimates, and that ordering is what mean-field VI gets right even when it gets the spread wrong. This is why ADVI is a defensible default for power analysis specifically, and why it is not a defensible default for reporting a per-gene credible interval.

Use backend="numpyro" or backend="pymc" when the interval itself is the deliverable.

Things this module refuses to get wrong¶

A gene with no across-well variance is not “not a hit”. If a gene’s log10expression column is constant, its coefficient is perfectly confounded with the intercept and carries no information whatsoever. The horseshoe will happily shrink it to ~0, which reads exactly like “tested, and not a hit”. That is a believable wrong answer, so those genes get NaN and are counted in n_unidentified_dropped instead.

A degenerate evaluation is NaN, never 0.5. If every gene is a hit, or none is, AUROC is undefined. Returning 0.5 there would be indistinguishable from “the method has no signal”, which is a completely different finding. Degenerate cases return NaN plus a reason string.

A failed fit is reported as failed. scan_parameters() records status="failed" or status="not_converged" with the exception text, and NaN metrics. A sweep that silently backfills 0.5 for the points that blew up looks, on a plot, exactly like a design that is at chance — which is the single most expensive mistake this module could make.

The sign convention is pinned down and tested. See evaluate_model_fit().

Attributes¶

BACKENDS

Concrete backends, in the order backend="auto" prefers them reversed.

Exceptions¶

PowerFitError

The power model could not be fit, so it has no estimate to give.

Classes¶

`ModelData`	Well-level design for the Poisson model, one row per well.
`PowerFit`	Posterior draws of the per-gene coefficients, plus how they were got.

Functions¶

`available_backends`(→ Dict[str, bool])	Report which backends this interpreter can actually run.
`evaluate_model_fit`(→ pandas.DataFrame)	Score the fit against ground-truth hit status: AUROC and average precision.
`fit_and_evaluate`(→ Tuple[PowerFit, pandas.DataFrame, ...)	Run the whole inference chain on one simulated screen.
`fit_model`(→ PowerFit)	Fit the horseshoe Poisson model and return posterior draws of `beta`.
`gather_model_estimate`(→ pandas.DataFrame)	Summarise the per-gene coefficients, one row per gene.
`prepare_model_data`(→ ModelData)	Collapse a tidy `(gene, well)` screen table to the well-level design.
`resolve_backend`(→ str)	Turn a requested backend into the one that will actually be used.
`scan_parameters`(→ pandas.DataFrame)	Sweep a design grid: simulate, fit, score, at every combination.

Module Contents¶

exception spacr.power_model.PowerFitError[source]¶

Bases: spacr.errors.SpacrError

The power model could not be fit, so it has no estimate to give.

Raised rather than returning a fit object with plausible-looking numbers in it. scan_parameters() catches this per grid point and records the point as failed; everywhere else it propagates.

Initialize self. See help(type(self)) for accurate signature.

class spacr.power_model.ModelData[source]¶

Well-level design for the Poisson model, one row per well.

Parameters:

wells – retained well identifiers in design-row order.
genes – gene identifiers in the column order of log10expression and fitted gene effects.
Npositive – positive-cell counts per retained well, summed over genes.
Ntotal – imaged-cell counts per retained well; every value is positive because the model uses log(Ntotal) as its offset.
log10expression – (n_wells, n_genes) matrix of log10(reads_gene_well / reads_well + 1e-4).
dropped_wells – wells removed because their imaged-cell count was not positive.
zero_read_wells – retained wells with no sequencing reads; their expression fractions use the pseudocount rather than nan.
unidentified_genes – genes whose expression column is constant across retained wells and whose fitted effect must remain unidentified.

to_frame() → pandas.DataFrame[source]¶

Return the design as a wide DataFrame, for eyeballing.

Columns are well, Npositive, Ntotal and one log10expression column per gene, named after the gene.

Returns:: pandas.DataFrame with n_wells rows.

property n_genes: int[source]¶: Number of genes (columns of the design, length of beta).

property n_wells: int[source]¶: Number of wells retained (rows of the design).

class spacr.power_model.PowerFit[source]¶

Posterior draws of the per-gene coefficients, plus how they were got.

Variables:

backend – the backend that actually ran. Recorded so a result can never be mistaken for one produced by a different method.
requested_backend – what the caller asked for ("auto", or a name).
method – "advi" or "nuts".
draws – (n_draws, n_genes) posterior draws of beta. Columns whose gene is unidentified are NaN throughout.
intercept_draws – (n_draws,) posterior draws of the intercept, on the centred covariate scale.
genes – gene labels, aligned with the columns of draws.
converged – whether the fit met its own convergence criterion. For ADVI that means the ELBO stopped improving; for NUTS it means no divergences and R-hat within tolerance.
diagnostics – backend-specific detail (ELBO trace, R-hat, timings, prior scales). Always includes beta_scale, naming the units beta is in.
data – the ModelData that was fit.
seed – the seed the fit was run with.

property n_draws: int[source]¶: Number of posterior draws held.

spacr.power_model.available_backends() → Dict[str, bool][source]¶

Report which backends this interpreter can actually run.

Returns:: mapping from backend name to availability. torch is always True in a working spaCR install, because spaCR depends on it.

>>> available_backends()["torch"]
True

spacr.power_model.evaluate_model_fit(data: pandas.DataFrame, model_estimate: pandas.DataFrame) → pandas.DataFrame[source]¶

Score the fit against ground-truth hit status: AUROC and average precision.

Port of evaluate_model_fit.

The sign convention, spelled out, because it is the easiest thing in this port to get backwards and the easiest to not notice. The R computes mean_inv = -mean, recodes hit as a factor with levels c("no", "yes"), and hands both to yardstick. yardstick’s default event_level is "first", so the event it scores is "no" — the non-hits — and the score it scores them with is -mean. Scoring -m for the event “not a hit” is identical to scoring +m for the event “is a hit”: both ROC curves are the same curve. So here we do the equivalent and much more legible thing, roc_auc_score(y_true=hit, y_score=mean), and spacr.power_model.evaluate_model_fit()’s tests pin the orientation on a case whose answer is known — strong planted hits must score near 1, not near 0. A flipped convention returns 1 - AUROC, which for a mediocre design looks entirely plausible.

Average precision uses scikit-learn’s step-wise definition (sum_n (R_n - R_{n-1}) P_n), the same estimator yardstick uses, so the two are comparable.

Parameters:

data – any frame carrying ground truth, with columns gene and hit (0/1 or boolean). The tidy screen table works directly; hit must be constant within a gene.
model_estimate – output of gather_model_estimate(), or any frame with gene and mean.

Returns:

one-row pandas.DataFrame with model_ap, model_auroc, ap_baseline (the prevalence, which is what average precision would be at chance), n_genes_scored, n_hits, n_non_hits, n_unidentified_dropped, n_missing_truth and reason. Degenerate cases give NaN metrics and a non-empty reason; they never give 0.5.

Raises:

PowerFitError – on missing columns, or a gene whose hit status is inconsistent between rows.

>>> import pandas as pd
>>> truth = pd.DataFrame({"gene": ["A", "B", "C", "D"],
...                       "hit": [1, 0, 0, 0]})
>>> est = pd.DataFrame({"gene": ["A", "B", "C", "D"],
...                     "mean": [2.0, 0.1, 0.0, -0.1]})
>>> float(evaluate_model_fit(truth, est)["model_auroc"].iloc[0])
1.0

spacr.power_model.fit_and_evaluate(data: pandas.DataFrame, *, fill_missing: bool = False, **fit_kwargs: Any) → Tuple[PowerFit, pandas.DataFrame, pandas.DataFrame][source]¶

Run the whole inference chain on one simulated screen.

prepare_model_data -> fit_model -> gather_model_estimate -> evaluate_model_fit, which is the sequence every vignette in the R package runs and the sequence scan_parameters() runs at each grid point.

Parameters:

data – tidy screen table (see prepare_model_data()).
fill_missing – passed to prepare_model_data().
fit_kwargs – passed to fit_model().

Returns:

(fit, model_estimate, model_evaluation).

spacr.power_model.fit_model(model_data: ModelData, *, backend: str = 'auto', seed: int = 0, standardize: bool = False, expected_hits: float | None = None, scale_global: float | None = None, df_local: float = 10.0, df_global: float = 1.0, df_slab: float = 4.0, scale_slab: float = 2.0, n_draws: int = 1000, n_steps: int = 3000, n_mc_samples: int = 8, learning_rate: float = 0.05, convergence_tol: float = 0.001, device: str = 'cpu', n_warmup: int = 1000, n_samples: int = 1000, n_chains: int = 4, max_tree_depth: int = 12, target_accept: float = 0.9) → PowerFit[source]¶

Fit the horseshoe Poisson model and return posterior draws of beta.

Port of compile_model + fit_model from spaCRPower’s R/fit_model.R, which run brms::brm(Npositive ~ 1 + log10expression + offset(log(Ntotal)), family = poisson, prior = horseshoe(df = 10)). df_local=10 here is that df = 10.

Parameters:

model_data – output of prepare_model_data().
backend – "auto", "torch", "numpyro" or "pymc". See resolve_backend(); the choice is recorded on the result.
seed – RNG seed. The same seed and backend give the same estimates.
standardize – scale each covariate column to unit SD before fitting, so the horseshoe shrinks every gene on a common scale. Changes the units of beta (recorded in diagnostics["beta_scale"]). Default False, matching the R.
expected_hits – prior guess at the number of true hits, used for the horseshoe’s global scale. Default: 5 % of the library.
scale_global – set the horseshoe global scale tau0 outright, ignoring expected_hits.
df_local – degrees of freedom of the local shrinkage half-Student-t. brms’s horseshoe(df=...).
df_global – degrees of freedom of the global half-Student-t; 1 is a half-Cauchy, the standard choice.
df_slab – degrees of freedom of the regularising slab.
scale_slab – scale of the regularising slab; caps how large a coefficient the prior will entertain.
n_draws – posterior draws to keep (torch backend; the NUTS backends keep n_samples * n_chains).
n_steps – ADVI optimisation steps (torch backend).
n_mc_samples – Monte-Carlo samples per ELBO gradient (torch).
learning_rate – Adam learning rate (torch).
convergence_tol – relative ELBO drift below which the ADVI fit is called converged (torch).
device – torch device; "cpu" by default because the problem is small and CPU float64 is bit-reproducible across machines.
n_warmup – NUTS warmup iterations (numpyro/pymc).
n_samples – NUTS post-warmup draws per chain (numpyro/pymc).
n_chains – NUTS chains (numpyro/pymc).
max_tree_depth – NUTS max tree depth (numpyro); the R uses 12.
target_accept – NUTS target acceptance (pymc).

Returns:

PowerFit.

Raises:

PowerFitError – on an unusable design, an unavailable named backend, or a diverged optimisation.

>>> import numpy as np
>>> md = ModelData(
...     wells=np.array(["w1", "w2"]),
...     genes=np.array(["A", "B"]),
...     Npositive=np.array([1, 9]),
...     Ntotal=np.array([100, 100]),
...     log10expression=np.array([[-0.3, -0.3], [-1.0, -0.05]]),
... )
>>> fit = fit_model(md, backend="torch", n_steps=50, n_draws=16, seed=0)
>>> fit.backend, fit.draws.shape
('torch', (16, 2))

spacr.power_model.gather_model_estimate(fit: PowerFit) → pandas.DataFrame[source]¶

Summarise the per-gene coefficients, one row per gene.

Port of gather_model_estimate, which runs posterior::summarize_draws() and keeps the b_log10expression* variables. The variable column reproduces that R naming so the two implementations’ outputs can be joined and compared directly.

Higher mean means stronger evidence the gene is a hit. Wells carrying more of a hit gene have more positive cells than their imaged cell count alone accounts for, which is a positive coefficient on log read fraction.

Parameters:: fit – result of fit_model().
Returns:: pandas.DataFrame with columns gene, variable, mean, sd, q5, q95, prob_positive (posterior probability that beta > 0) and identified. Unidentified genes carry NaN in every numeric column and identified=False.
Raises:: PowerFitError – if fit is not a PowerFit.

>>> import numpy as np
>>> fit = PowerFit(
...     backend="torch", requested_backend="torch", method="advi",
...     draws=np.array([[1.0, 0.0], [3.0, 0.0]]),
...     intercept_draws=np.array([-4.0, -4.0]),
...     genes=np.array(["A", "B"]), converged=True)
>>> gather_model_estimate(fit)[["gene", "mean"]].to_dict("records")
[{'gene': 'A', 'mean': 2.0}, {'gene': 'B', 'mean': 0.0}]

spacr.power_model.prepare_model_data(data: pandas.DataFrame, *, fill_missing: bool = False) → ModelData[source]¶

Collapse a tidy (gene, well) screen table to the well-level design.

Port of prepare_model_data from spaCRPower’s R/fit_model.R. Per well:

Npositive       = sum over genes of `positive`
Ntotal          = cells imaged in the well
log10expression = log10(n_reads_per_gene_per_well / total_reads + 1e-4)

Ntotal is read from an imaging_n_cells_per_well column if the simulator provides one (it must then be constant within a well, or the table is malformed and we say so); otherwise it is summed from imaging_n_cells_per_gene_per_well. The two agree by construction, since the per-gene counts are a multinomial split of the per-well total.

Two deliberate departures from the R:

The R writes Npositive = sum(well_data$positive[1]) — sum of a single element, i.e. the positive count of whichever gene happens to sort first in the well, not the well’s total. That is a bug in the R (the surrounding documentation says “number of positive”), and reproducing it would throw away most of the response. We sum over all genes in the well.
Genes whose covariate column ends up constant across wells are flagged in ModelData.unidentified_genes. Their coefficients are not estimable and are reported as NaN rather than as a shrunk-to-zero “not a hit”.

Parameters:

data – tidy screen table with one row per (gene, well) pair. Required columns: gene, well, positive, n_reads_per_gene_per_well, and one of imaging_n_cells_per_well / imaging_n_cells_per_gene_per_well.
fill_missing – if the gene x well grid is incomplete, False (the default) raises, because a missing pair is ambiguous — zero reads, or not measured? True fills the missing pairs with zeros for every count column.

Returns:

ModelData.

Raises:

PowerFitError – on a missing required column, a duplicated (gene, well) pair, an incomplete grid with fill_missing=False, a within-well inconsistent imaging_n_cells_per_well, Npositive > Ntotal, negative counts, or an empty design.

Worked example — two wells, two genes, gene B dominating well w2:

>>> import pandas as pd
>>> tidy = pd.DataFrame({
...     "well": ["w1", "w1", "w2", "w2"],
...     "gene": ["A", "B", "A", "B"],
...     "positive": [1, 0, 0, 9],
...     "imaging_n_cells_per_gene_per_well": [50, 50, 10, 90],
...     "n_reads_per_gene_per_well": [500, 500, 100, 900],
... })
>>> md = prepare_model_data(tidy)
>>> list(md.wells), list(md.genes)
(['w1', 'w2'], ['A', 'B'])
>>> list(md.Npositive), list(md.Ntotal)
([1, 9], [100, 100])
>>> md.log10expression.round(3).tolist()
[[-0.301, -0.301], [-1.0, -0.046]]

spacr.power_model.resolve_backend(backend: str = 'auto') → str[source]¶

Turn a requested backend into the one that will actually be used.

"auto" prefers exact NUTS where it is installed — numpyro, then pymc — and otherwise uses torch, which always is. A named backend is never silently substituted: asking for pymc without pymc installed raises, because a power analysis quietly computed by a different inference method than the one in your methods section is a reproducibility failure, not a convenience.

Parameters:: backend – "auto", "torch", "numpyro" or "pymc".
Returns:: the concrete backend name.
Raises:: PowerFitError – for an unknown name, or a known one that is not installed.

>>> resolve_backend("torch")
'torch'

spacr.power_model.scan_parameters(*, progress_file: str | None = None, resume: bool = True, on_error: str = 'record', backend: str = 'auto', seed: int = 0, n_replicates: int = 1, verbose: bool = False, simulate_fn: Callable[..., pandas.DataFrame] | None = None, fit_kwargs: Mapping[str, Any] | None = None, on_point: Callable[[Mapping[str, Any]], Any] | None = None, **parameters: Any) → pandas.DataFrame[source]¶

Sweep a design grid: simulate, fit, score, at every combination.

Port of scan_parameters from spaCRPower’s R/scan_parameters.R. Every keyword argument that is not one of the named options above is a simulator parameter; pass a scalar to hold it fixed and a sequence to sweep it. The grid is the Cartesian product, in the order the sweeps were given (matching R’s tidyr::expand_grid), with param_index numbered from 1 as in the R.

A point that fails is reported as failed. A fit that raises lands as status="failed" with the exception text in error; a fit that runs but does not meet its convergence criterion lands as status="not_converged". Both carry NaN metrics. This is the whole point: a sweep that backfilled 0.5 for its broken points would plot as a design sitting at chance, and “this design cannot find its hits” and “this fit crashed” are conclusions with opposite consequences.

Resuming. With progress_file, each completed row is appended to a TSV immediately, so a killed sweep loses at most the point it was on. Re-running with resume=True (the default) reads the file and skips points whose run_key is already there. The key is a digest of the parameter values, the replicate, the seed and the backend – not the row number – so extending or reordering the grid does not invalidate finished work.

Parameters:

progress_file – TSV to append completed rows to; created, with its parent directory, if absent.
resume – skip points already present in progress_file. With resume=False and a non-empty progress file, this raises rather than appending a second copy of every point.
on_error – "record" (default) logs the failure, writes the row and continues; "raise" re-raises, for debugging one point.
backend – inference backend, see resolve_backend().
seed – master seed. Each point’s seed is derived from (seed, param_index, replicate), so a point’s data does not depend on how many points ran before it.
n_replicates – independent simulated screens per grid point. One screen at one setting is a single draw from a noisy process; the sweep is much easier to read with three or five.
verbose – log each point’s parameters and score as it runs.
simulate_fn – screen simulator; defaults to spacr.power_simulate.simulate_screen().
fit_kwargs – extra keyword arguments for fit_model().
on_point –
called once per completed (point, replicate) with a mapping {"index", "total", "point_index", "n_points", "replicate", "resumed", "row"} — index and total count fits, one-based, so index / total is a progress fraction, and row is that fit’s result record.

Returning exactly False stops the sweep and returns the rows completed so far; anything else (including None) continues. This exists for one reason: a sweep is minutes long and a GUI has to be able to show where it is and to stop it without killing the thread. Between points is the only safe granularity — one fit is atomic — so a cancel is honoured after the fit in flight, not during it. The callback runs on the calling thread, inline; it must not block, and a callback that raises aborts the sweep rather than being swallowed, because a broken progress reporter is a bug in the caller and hiding it would leave the sweep running with nothing watching it.
parameters – simulator parameters, scalar or sequence.

Returns:

pandas.DataFrame, one row per (grid point, replicate), with every parameter column plus the columns in _SCAN_RESULT_COLUMNS. Rows loaded from a resumed progress file are included. DataFrame.attrs carries n_planned (fits the full grid would have run) and cancelled (whether on_point stopped it early), so a short frame can be told from a small grid.

Raises:

PowerFitError – if no parameters were given, on_error is not recognised, n_replicates < 1, the progress file exists with a different column layout (appending misaligned rows to a TSV is a silent data-corruption bug, so we stop instead), or the progress file is non-empty and resume=False.

Example – sweep the number of wells, three replicates each, resumable. Every keyword here is forwarded verbatim to the simulator, so the full set is whatever spacr.power_simulate.simulate_screen() accepts; the ones below are the design knobs a power analysis usually turns:

scores = scan_parameters(
    n_genes_in_library=200,
    gene_hit_rate=0.05,
    n_wells_per_screen=[24, 48, 96],   # the sweep
    class_pos_mu=0.99,
    class_neg_mu=0.01,
    ...,                               # remaining simulator arguments
    n_replicates=3,
    progress_file="scan/wells.tsv",
)
solvable = scores.query("status == 'ok' and model_auroc >= 0.9")

spacr.power_model.BACKENDS: Tuple[str, ...] = ('torch', 'numpyro', 'pymc')[source]¶

Concrete backends, in the order backend="auto" prefers them reversed.

resolve_backend prefers exact NUTS (numpyro, then pymc) and falls back to torch, which is always present because spaCR depends on it.

Nested helpers¶

_fit_numpyro_nuts.model()¶

Define the regularized-horseshoe Poisson model for NumPyro NUTS.

spacr/power_model.py:1011

_fit_torch_advi._log_joint(theta: 'torch.Tensor') → 'torch.Tensor'¶

Evaluate the unnormalized horseshoe-Poisson log posterior.

The parameter-independent Poisson factorial is omitted, so each leading draw receives the same optimizer-equivalent objective.

spacr/power_model.py:829

_fit_torch_advi._unpack(theta: 'torch.Tensor')¶

Split the final parameter axis while preserving leading draws.

spacr/power_model.py:820

prepare_model_data._matrix(column: str) → np.ndarray¶

Pivot column to captured well-by-gene order as float64.

spacr/power_model.py:414

scan_parameters._report(row: Dict[str, Any], point_index: int, replicate: int, resumed: bool) → bool¶

Hand one finished row to on_point; True means “keep going”.

Kept as a closure over the loop counters rather than inlined twice so the resumed and freshly-computed paths cannot drift apart — a progress bar that skips the rows a resume restored is a progress bar that runs backwards on the second attempt.

spacr/power_model.py:1949