diff --git a/docs/changelog.md b/docs/changelog.md index 539fcdef..ad9e3e29 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) ## v1.5.0 +### 🚀 Features + +* `Statsmodels` fits linear mixed models when the design contains a random intercept such as `~condition + (1 | donor)`, by REML with Satterthwaite t-tests like `lmerTest`, so single cells can be tested while accounting for the samples they come from; the kernel is shared with Dialogue and the `(1 | variable)` syntax with Milo ([#1147](https://github.com/scverse/pertpy/pull/1147)) @Sizerta @Zethson + ### ⚠️ Changed results * Dialogue reproduces the R reference end to end: REML mixed models with Satterthwaite t-tests like `lmerTest`, R's PMD penalty, permutation test and cell-type pruning, Pearson pre-selection of signature genes, one-sided gene p-values, and `test_phenotype_association` now follows `DIALOGUE.pheno` (first level or `True` as the indicator, pooled `All` row, no Fisher combination) ([#1177](https://github.com/scverse/pertpy/pull/1177)) @Zethson diff --git a/docs/tutorials/notebooks b/docs/tutorials/notebooks index c6404f02..51649f61 160000 --- a/docs/tutorials/notebooks +++ b/docs/tutorials/notebooks @@ -1 +1 @@ -Subproject commit c6404f023149ece7ec90331f538b9c33c6a14222 +Subproject commit 51649f61d5476815d539a0469b07f719ed63b8f0 diff --git a/src/pertpy/_parallel.py b/src/pertpy/_parallel.py index 7df481b1..d7b4c460 100644 --- a/src/pertpy/_parallel.py +++ b/src/pertpy/_parallel.py @@ -5,8 +5,10 @@ import math from collections.abc import Iterable, Iterator +from contextlib import contextmanager from typing import Any +import numba import numpy as np from joblib import Parallel from tqdm.auto import tqdm @@ -29,6 +31,17 @@ def _parallelize_with_joblib(delayed_objects: Iterable[Any], *, total: int | Non return iter(Parallel(return_as="list", **kwargs)(delayed_objects)) +@contextmanager +def _numba_threads(n_tasks: int) -> Iterator[None]: + """Cap the numba thread count at `n_tasks` within the block, since every thread of a parallel region costs time even when idle.""" + n_threads = numba.get_num_threads() + numba.set_num_threads(max(1, min(n_threads, n_tasks))) + try: + yield + finally: + numba.set_num_threads(n_threads) + + def _block_slices(n_items: int, *, n_blocks: int = 1, max_block_size: int | None = None) -> list[slice]: """Split `range(n_items)` into contiguous blocks, of `max_block_size` items at most.""" if n_items <= 0: diff --git a/src/pertpy/tools/_dialogue.py b/src/pertpy/tools/_dialogue.py index 0b63aeaa..e3c5c3aa 100644 --- a/src/pertpy/tools/_dialogue.py +++ b/src/pertpy/tools/_dialogue.py @@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, cast import matplotlib.pyplot as plt -import numba import numpy as np import pandas as pd import scanpy as sc @@ -29,7 +28,9 @@ from statsmodels.stats.multitest import multipletests from pertpy._doc import _doc_params, doc_common_plot_args +from pertpy._parallel import _numba_threads from pertpy._types import CSBase, cast_dense, cast_frame, cast_matrix +from pertpy.tools._mixed_models import _random_intercept_fits if TYPE_CHECKING: from collections.abc import Sequence @@ -413,188 +414,6 @@ def _iterative_nnls( return coef -@njit(cache=True) -def _solve_small(a: np.ndarray, b: np.ndarray) -> np.ndarray: - """Solve ``a @ x = b`` by Gaussian elimination with partial pivoting, all NaN if ``a`` is singular to working precision.""" - m = a.shape[0] - a = a.copy() - x = b.copy() - tol = 1e-7 * np.abs(np.diag(a)) - for k in range(m): - piv = k + np.argmax(np.abs(a[k:, k])) - if not np.abs(a[piv, k]) > tol[k]: - return np.full(m, np.nan) - for j in range(m): - a[k, j], a[piv, j] = a[piv, j], a[k, j] - x[k], x[piv] = x[piv], x[k] - for i in range(k + 1, m): - f = a[i, k] / a[k, k] - a[i, k:] -= f * a[k, k:] - x[i] -= f * x[k] - for k in range(m - 1, -1, -1): - x[k] = (x[k] - np.sum(a[k, k + 1 :] * x[k + 1 :])) / a[k, k] - return x - - -@njit(cache=True) -def _random_intercept_gls( - within: np.ndarray, sums: np.ndarray, sizes: np.ndarray, gamma: float -) -> tuple[np.ndarray, float, np.ndarray, np.ndarray]: - """GLS fit of a random-intercept model, given the ratio ``gamma`` of the random-intercept to the residual variance. - - ``within`` holds the within-group cross-products and ``sums`` the per-group sums of the design columns followed by the response. - Returns the fixed effects, the generalized residual sum of squares, the generalized cross-products, and the per-group residual sums divided by ``1 + n_g * gamma``. - """ - p = within.shape[0] - 1 - shrink = 1.0 / (1.0 + sizes * gamma) - cross = within + (sums.T * (shrink / sizes)) @ sums - beta = _solve_small(cross[:p, :p], cross[:p, p].copy()) - rss = cross[p, p] - np.sum(cross[:p, p] * beta) - resid = sums[:, p].copy() - for a in range(p): - resid -= sums[:, a] * beta[a] - return beta, rss, cross, resid * shrink - - -@njit(cache=True) -def _inverse_small(a: np.ndarray) -> np.ndarray: - """Inverse of ``a`` by solving against the identity, all NaN if ``a`` is singular to working precision.""" - m = a.shape[0] - out = np.empty((m, m)) - for j in range(m): - unit = np.zeros(m) - unit[j] = 1.0 - out[:, j] = _solve_small(a, unit) - return out - - -@njit(cache=True) -def _random_intercept_slope(within: np.ndarray, sums: np.ndarray, sizes: np.ndarray, log_gamma: float) -> float: - """Derivative of the REML profile log-likelihood with respect to ``log(gamma)``.""" - gamma = np.exp(log_gamma) - p = within.shape[0] - 1 - _, rss, cross, resid = _random_intercept_gls(within, sums, sizes, gamma) - shrink = 1.0 / (1.0 + sizes * gamma) - design_sums = np.ascontiguousarray(sums[:, :p]) - quad = np.sum((design_sums @ _inverse_small(cross[:p, :p])) * design_sums, axis=1) - return ( - 0.5 * gamma * ((sizes.sum() - p) * np.sum(resid**2) / rss - np.sum(sizes * shrink) + np.sum(shrink**2 * quad)) - ) - - -@njit(cache=True) -def _random_intercept_log_gamma(within: np.ndarray, sums: np.ndarray, sizes: np.ndarray) -> float: - """Maximize the REML profile log-likelihood over ``log(gamma)`` by bracketing a root of its slope from 0, then bisection. - - Returns -inf if the maximum is on the boundary ``gamma = 0`` and NaN if the fit is degenerate. - """ - a, fa = 0.0, _random_intercept_slope(within, sums, sizes, 0.0) - b, fb = a, fa - step = 2.0 if fa > 0 else -2.0 - while fb != 0.0 and (fb > 0) == (fa > 0): - a, fa = b, fb - b += step - if b < -40.0: - return -np.inf - if b > 40.0 or not np.isfinite(fa): - return np.nan - fb = _random_intercept_slope(within, sums, sizes, b) - if not np.isfinite(fb): - return np.nan - mid = 0.5 * (a + b) - while fb != 0.0 and a != mid != b: - f_mid = _random_intercept_slope(within, sums, sizes, mid) - if (f_mid > 0) == (fa > 0): - a, fa = mid, f_mid - else: - b, fb = mid, f_mid - mid = 0.5 * (a + b) - return b if fb == 0.0 else mid - - -@njit(cache=True) -def _satterthwaite_df( # noqa: PLR0917 - sums: np.ndarray, sizes: np.ndarray, resid: np.ndarray, gamma: float, sigma2: float, inv: np.ndarray -) -> float: - """Satterthwaite denominator degrees of freedom of the coefficient at index 1, as lmerTest computes them. - - Uses the gradient of the coefficient's variance and the REML observed information, both with respect to ``(gamma, sigma2)``. - """ - p = inv.shape[0] - shrink = 1.0 / (1.0 + sizes * gamma) - design_sums = np.ascontiguousarray(sums[:, :p]) - scaled = design_sums @ inv - quad = np.sum(scaled * design_sums, axis=1) - v = np.zeros(p) - for g in range(sizes.shape[0]): - v += shrink[g] * resid[g] * design_sums[g] - pair_quad = np.outer(shrink, shrink) * (scaled @ np.ascontiguousarray(design_sums.T)) - log_det_curvature = -np.sum((sizes * shrink) ** 2) - np.sum(pair_quad**2) + 2.0 * np.sum(sizes * shrink**3 * quad) - rss_curvature = 2.0 * np.sum(sizes * shrink * resid**2) - 2.0 * np.dot(v, inv @ v) - info_gg = 0.5 * (log_det_curvature + rss_curvature / sigma2) - info_gs = np.sum(resid**2) / (2.0 * sigma2**2) - info_ss = (sizes.sum() - p) / (2.0 * sigma2**2) - grad_g = sigma2 * np.sum(shrink**2 * scaled[:, 1] ** 2) - grad_s = inv[1, 1] - det = info_gg * info_ss - info_gs**2 - denominator = (grad_g**2 * info_ss - 2.0 * grad_g * grad_s * info_gs + grad_s**2 * info_gg) / det - return 2.0 * (sigma2 * inv[1, 1]) ** 2 / denominator - - -@njit(parallel=True, cache=True) -def _random_intercept_fits( - x: np.ndarray, fixed: np.ndarray, groups: np.ndarray, n_groups: int -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """REML fits of ``fixed[:, -1] ~ fixed[:, 0] + x[i] + fixed[:, 1:-1] + (1 | groups)`` for every row ``x[i]``. - - Returns the coefficients of ``x[i]``, their standard errors and their Satterthwaite degrees of freedom, like lmerTest's ``summary``. - """ - n, n_fixed = fixed.shape - m = n_fixed + 1 - p = m - 1 - sizes = np.zeros(n_groups) - fixed_sums = np.zeros((n_groups, n_fixed)) - for i in range(n): - sizes[groups[i]] += 1.0 - for a in range(n_fixed): - fixed_sums[groups[i], a] += fixed[i, a] - centered = np.empty((n_fixed, n)) - for i in range(n): - for a in range(n_fixed): - centered[a, i] = fixed[i, a] - fixed_sums[groups[i], a] / sizes[groups[i]] - cols = np.concatenate((np.zeros(1, dtype=np.int64), np.arange(2, m))) - within = np.zeros((m, m)) - sums = np.zeros((n_groups, m)) - for a in range(n_fixed): - sums[:, cols[a]] = fixed_sums[:, a] - for b in range(n_fixed): - within[cols[a], cols[b]] = centered[a] @ centered[b] - estimates = np.full(x.shape[0], np.nan) - std_errors = np.full(x.shape[0], np.nan) - dfs = np.full(x.shape[0], np.nan) - for row in prange(x.shape[0]): - row_within = within.copy() - row_sums = sums.copy() - for i in range(n): - row_sums[groups[i], 1] += x[row, i] - means = row_sums[:, 1] / sizes - for i in range(n): - row_within[1, 1] += (x[row, i] - means[groups[i]]) ** 2 - for a in range(n_fixed): - row_within[1, cols[a]] = row_within[cols[a], 1] = x[row] @ centered[a] - log_gamma = _random_intercept_log_gamma(row_within, row_sums, sizes) - if np.isnan(log_gamma): - continue - gamma = np.exp(log_gamma) - beta, rss, cross, resid = _random_intercept_gls(row_within, row_sums, sizes, gamma) - inv = _inverse_small(cross[:p, :p]) - sigma2 = rss / (n - p) - estimates[row] = beta[1] - std_errors[row] = np.sqrt(sigma2 * inv[1, 1]) - dfs[row] = n - p if gamma == 0.0 else _satterthwaite_df(row_sums, sizes, resid, gamma, sigma2, inv) - return estimates, std_errors, dfs - - def _hlm_pvalue_per_row( expression: np.ndarray | pd.DataFrame, score: np.ndarray, @@ -630,14 +449,10 @@ def _hlm_pvalue_per_row( dfs = np.full(expression.shape[0], np.nan) rows = ~np.isnan(expression).any(axis=1) & ~np.isnan(fixed).any() if rows.any(): - n_threads = numba.get_num_threads() - numba.set_num_threads(min(n_threads, int(rows.sum()))) - try: + with _numba_threads(int(rows.sum())): estimates[rows], std_errors[rows], dfs[rows] = _random_intercept_fits( np.ascontiguousarray(expression[rows]), fixed, groups.astype(np.int64), len(labels) ) - finally: - numba.set_num_threads(n_threads) pvalues = 2 * stats.t.sf(np.abs(estimates / std_errors), dfs) return pd.DataFrame( {"estimate": estimates, "std_error": std_errors, "df": dfs, "pvalue": pvalues}, index=gene_index diff --git a/src/pertpy/tools/_differential_gene_expression/_base.py b/src/pertpy/tools/_differential_gene_expression/_base.py index 10213521..044b40f6 100644 --- a/src/pertpy/tools/_differential_gene_expression/_base.py +++ b/src/pertpy/tools/_differential_gene_expression/_base.py @@ -23,6 +23,7 @@ from pertpy._types import CSBase, cast_dense, cast_frame from pertpy.tools import PseudobulkSpace from pertpy.tools._differential_gene_expression._checks import check_is_numeric_matrix +from pertpy.tools._mixed_models import parse_random_effects # Colors are expressed relative to the surrounding text so that the summary renders on light and dark themes alike. # The table styles undo the borders, zebra stripes and centering that notebook frontends apply to rendered tables. @@ -1023,12 +1024,15 @@ def _get_significance(p_val): class LinearModelBase(MethodBase): + _supports_random_effects = False + def __init__(self, adata, design, *, mask=None, layer=None, **kwargs): """Initialize the method. Args: adata: AnnData object, usually pseudobulked. design: Model design. Can be either a design matrix, a formulaic formula.Formulaic formula in the format 'x + z' or '~x+z'. + Methods that support random effects accept random intercepts written as '(1 | variable)'. mask: A column in adata.var that contains a boolean mask with selected features. layer: Layer to use in fit(). If None, use the X array. **kwargs: Keyword arguments specific to the method implementation. @@ -1039,7 +1043,14 @@ def __init__(self, adata, design, *, mask=None, layer=None, **kwargs): from formulaic_contrasts import FormulaicContrasts self.formulaic_contrasts = None + self.random_effects: list[str] = [] if isinstance(design, str): + design, self.random_effects = parse_random_effects(design) + if self.random_effects and not self._supports_random_effects: + raise ValueError( + f"{type(self).__name__} does not support random effects such as '(1 | {self.random_effects[0]})'. " + "Statsmodels fits linear mixed models with a random intercept." + ) self.formulaic_contrasts = FormulaicContrasts(adata.obs, design) self.design = self.formulaic_contrasts.design_matrix else: @@ -1101,6 +1112,8 @@ def _summary(self) -> dict[str, str]: else: summary["Design"] = str(self.design.model_spec.formula) summary["Variables"] = _format_names(sorted(self.variables)) + if self.random_effects: + summary["Random effects"] = _format_names(self.random_effects) if hasattr(self.design, "columns"): summary["Coefficients"] = _format_names([str(column) for column in self.design.columns]) summary["Fitted"] = "yes" if self.is_fitted else "no" diff --git a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py index 5e5f33e5..e92a16e9 100644 --- a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py +++ b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py @@ -3,20 +3,22 @@ import numpy as np import pandas as pd -import statsmodels import statsmodels.api as sm +from anndata import AnnData from fast_array_utils.conv import to_dense from joblib import delayed, effective_n_jobs from scipy import stats from scipy.sparse import issparse -from statsmodels.tools.sm_exceptions import SingularMatrixWarning +from statsmodels.tools.sm_exceptions import ConvergenceWarning, SingularMatrixWarning from statsmodels.tools.tools import pinv_extended, recipr -from pertpy._parallel import _MAX_BLOCK_ELEMENTS, _block_slices, _parallelize_with_joblib +from pertpy._parallel import _MAX_BLOCK_ELEMENTS, _block_slices, _numba_threads, _parallelize_with_joblib from pertpy._types import CSBase +from pertpy.tools._mixed_models import _random_intercept_contrast, _random_intercept_response_fits, random_effect_codes from ._base import LinearModelBase from ._checks import check_is_numeric_matrix +from ._simple_tests import fdr_correction def _fit_block(regression_model, endog_block: np.ndarray, exog, kwargs: dict) -> list: @@ -60,8 +62,86 @@ def __getitem__(self, index): return sm.OLS(to_dense(self.endog[:, [index]])[:, 0], self.exog).fit() +class _MixedFits(Sequence): + """REML fits of every column of `endog` on the same design with a random intercept per level of `groups`, of which the statsmodels results are only built when accessed. + + Contrasts of the fixed effects are tested with Satterthwaite degrees of freedom like `lmerTest`. + Columns whose fit is degenerate, such as constant ones, hold NaN. + """ + + def __init__(self, endog, exog, groups: np.ndarray, n_groups: int): + self.endog, self.exog, self.groups = endog, exog, groups + design = np.ascontiguousarray(exog, dtype=np.float64) + n_vars, p = endog.shape[1], design.shape[1] + self.params = np.full((n_vars, p), np.nan) + self.inverses = np.full((n_vars, p, p), np.nan) + self.gammas = np.full(n_vars, np.nan) + self.scale = np.full(n_vars, np.nan) + self.residual_sums = np.full((n_vars, n_groups), np.nan) + self.design_sums = np.zeros((n_groups, p)) + self.sizes = np.zeros(n_groups) + for block in _block_slices(n_vars, max_block_size=_MAX_BLOCK_ELEMENTS // max(design.shape[0], 1)): + responses = np.ascontiguousarray(to_dense(endog[:, block]).T, dtype=np.float64) + with _numba_threads(responses.shape[0]): + fits = _random_intercept_response_fits(responses, design, groups, n_groups) + ( + self.params[block], + self.inverses[block], + self.gammas[block], + self.scale[block], + self.residual_sums[block], + ) = fits[:5] + self.design_sums, self.sizes = fits[5:] + + def __len__(self) -> int: + return len(self.params) + + def __getitem__(self, index): + if isinstance(index, slice): + return [self[i] for i in range(len(self))[index]] + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ConvergenceWarning) + return sm.MixedLM(to_dense(self.endog[:, [index]])[:, 0], self.exog, groups=self.groups).fit() + + def contrast(self, contrast: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Estimate, standard error and Satterthwaite degrees of freedom of `contrast` for every column.""" + with _numba_threads(len(self)): + return _random_intercept_contrast( + np.ascontiguousarray(contrast, dtype=np.float64), + self.params, + self.inverses, + self.gammas, + self.scale, + self.residual_sums, + self.design_sums, + self.sizes, + ) + + class Statsmodels(LinearModelBase): - """Differential expression test using a statsmodels linear regression.""" + """Differential expression test using a statsmodels linear regression. + + A random intercept written as ``(1 | variable)`` in the design, e.g. ``"~condition + (1 | donor)"``, turns the linear regression into a linear mixed model. + It is fitted by REML and its contrasts are tested with Satterthwaite degrees of freedom like ``lmerTest``, so effects that are constant within a donor are judged by the number of donors rather than the number of cells. + This is how to test single cells of a cell type while accounting for the samples they come from. + The per-gene results in ``models`` are then refits with ``statsmodels.MixedLM``, whose z-test p-values differ from the Satterthwaite p-values of ``test_contrasts``. + """ + + _supports_random_effects = True + + def __init__( + self, adata: AnnData, design: str | np.ndarray, *, mask: str | None = None, layer: str | None = None, **kwargs + ): + super().__init__(adata, design, mask=mask, layer=layer, **kwargs) + if len(self.random_effects) > 1: + raise ValueError( + f"Statsmodels supports a single random intercept, but the design has {len(self.random_effects)}: " + f"{self.random_effects}." + ) + self.groups: np.ndarray | None = None + self.n_groups = 0 + if self.random_effects: + self.groups, self.n_groups = random_effect_codes(self.adata.obs, self.random_effects[0]) def _check_counts(self): check_is_numeric_matrix(self.data) @@ -77,9 +157,10 @@ def fit( Args: regression_model: A statsmodels regression model class, either OLS or GLM. + A design with a random intercept requires OLS, which then fits the linear mixed model. n_jobs: Number of jobs to distribute the variables over, passed to `joblib.Parallel`. None means one job, -1 all available cores. - Unused by OLS without further arguments, which fits all variables at once. + Unused by OLS without further arguments and by the mixed model, which fit all variables at once. **kwargs: Additional arguments for fitting the specific method. In particular, this is where you can specify the family for GLM. @@ -95,8 +176,14 @@ def fit( if isinstance(data, CSBase): data = data.tocsc() + if self.groups is not None: + if regression_model is not sm.OLS or kwargs: + raise ValueError("A random intercept is only supported for OLS without further fit arguments.") + self.models: Sequence = _MixedFits(data, self.design, self.groups, self.n_groups) + self._fitted = True + return if regression_model is sm.OLS and not kwargs: - self.models: Sequence = _OLSFits(data, self.design) + self.models = _OLSFits(data, self.design) self._fitted = True return @@ -123,7 +210,10 @@ def fit( def _test_single_contrast(self, contrast, **kwargs) -> pd.DataFrame: r_matrix = np.atleast_2d(np.asarray(contrast, dtype=np.float64)) - if isinstance(self.models, _OLSFits): + if isinstance(self.models, _MixedFits): + effect, sd, df_resid = self.models.contrast(r_matrix[0]) + use_t = np.ones(len(effect), dtype=bool) + elif isinstance(self.models, _OLSFits): effect = self.models.params @ r_matrix[0] sd = np.sqrt(self.models.scale * (r_matrix @ self.models.normalized_cov_params @ r_matrix.T).item()) df_resid = np.full(len(effect), self.models.df_resid) @@ -136,7 +226,7 @@ def _test_single_contrast(self, contrast, **kwargs) -> pd.DataFrame: use_t = np.array([getattr(res, "use_t", False) for res in results], dtype=bool) t_value = effect * recipr(sd) p_value = np.where(use_t, stats.t.sf(np.abs(t_value), df_resid), stats.norm.sf(np.abs(t_value))) * 2 - return ( + return fdr_correction( pd.DataFrame( { "variable": self.adata.var_names, @@ -145,7 +235,5 @@ def _test_single_contrast(self, contrast, **kwargs) -> pd.DataFrame: "sd": sd, "log_fc": effect, } - ) - .sort_values("p_value") - .assign(adj_p_value=lambda x: statsmodels.stats.multitest.fdrcorrection(x["p_value"])[1]) + ).sort_values("p_value") ) diff --git a/src/pertpy/tools/_milo.py b/src/pertpy/tools/_milo.py index 7a4c2501..c720b4c4 100644 --- a/src/pertpy/tools/_milo.py +++ b/src/pertpy/tools/_milo.py @@ -20,7 +20,8 @@ from pertpy._doc import _doc_params, doc_common_plot_args from pertpy._logger import logger from pertpy._types import CSBase, cast_dense, cast_frame, cast_matrix -from pertpy.tools._milo_glmm import fit_nb_glmm_nhoods, log_cpm, parse_random_effects, random_effect_matrices +from pertpy.tools._milo_glmm import fit_nb_glmm_nhoods, log_cpm +from pertpy.tools._mixed_models import parse_random_effects, random_effect_matrices if TYPE_CHECKING: from collections.abc import Collection, Sequence diff --git a/src/pertpy/tools/_milo_glmm.py b/src/pertpy/tools/_milo_glmm.py index a9f5c2a3..0b29138e 100644 --- a/src/pertpy/tools/_milo_glmm.py +++ b/src/pertpy/tools/_milo_glmm.py @@ -1,7 +1,6 @@ from __future__ import annotations import math -import re from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, NamedTuple @@ -13,45 +12,10 @@ if TYPE_CHECKING: from collections.abc import Sequence -_RANDOM_EFFECT = re.compile(r"\(\s*1\s*\|\s*([^)]+?)\s*\)") _LOG_2PI = float(np.log(2 * np.pi)) _LOG_DISPERSION_BOUNDS = (float(np.log(1e-6)), float(np.log(1e3))) -def parse_random_effects(design: str) -> tuple[str, list[str]]: - """Split a formula into its fixed effects part and the variables entering as random intercepts. - - Random intercepts follow the ``(1 | variable)`` syntax of lme4 and R Milo. - - Returns: - The formula with the random effect terms removed and the random intercept variables. - """ - if re.search(r"\|", _RANDOM_EFFECT.sub("", design)): - raise ValueError(f"{design!r} is an invalid formula for random effects. Use the '(1 | variable)' format.") - - random_effects = [match.group(1).strip() for match in _RANDOM_EFFECT.finditer(design)] - fixed = _RANDOM_EFFECT.sub("", design) - fixed = re.sub(r"\+\s*(?=\+|$)", "", fixed).strip().rstrip("+").strip() - if fixed in {"", "~"}: - fixed = "~ 1" - return fixed, random_effects - - -def random_effect_matrices(obs: pd.DataFrame, random_effects: Sequence[str]) -> list[tuple[str, np.ndarray]]: - """Build one indicator matrix of shape samples x levels per random intercept variable.""" - matrices = [] - for variable in random_effects: - if variable not in obs.columns: - raise ValueError(f"Random effect variable {variable!r} is not a column of the sample metadata.") - dummies = pd.get_dummies(obs[variable].astype("category"), drop_first=False) - if dummies.shape[1] < 2: - raise ValueError( - f"Random effect variable {variable!r} has a single level, which cannot be told apart from the intercept." - ) - matrices.append((variable, dummies.to_numpy(dtype=float))) - return matrices - - class GLMMFit(NamedTuple): """Result of fitting a negative binomial GLMM to the counts of a single neighbourhood.""" diff --git a/src/pertpy/tools/_mixed_models.py b/src/pertpy/tools/_mixed_models.py new file mode 100644 index 00000000..c448cfac --- /dev/null +++ b/src/pertpy/tools/_mixed_models.py @@ -0,0 +1,513 @@ +"""Linear mixed models with one random intercept, shared by Milo, Dialogue and the differential expression interface. + +Random intercepts are written as ``(1 | variable)`` in formulas, the syntax of lme4 and R Milo. +The numba kernels fit ``y ~ X + (1 | group)`` by REML and compute Satterthwaite degrees of freedom like ``lmerTest``. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +import numpy as np +import pandas as pd +from numba import njit, prange + +if TYPE_CHECKING: + from collections.abc import Sequence + +_RANDOM_EFFECT = re.compile(r"\(\s*1\s*\|\s*([^)]+?)\s*\)") + + +def parse_random_effects(design: str) -> tuple[str, list[str]]: + """Split a formula into its fixed effects part and the variables entering as random intercepts. + + Random intercepts follow the ``(1 | variable)`` syntax of lme4 and R Milo. + + Returns: + The formula with the random effect terms removed and the random intercept variables. + """ + if re.search(r"\|", _RANDOM_EFFECT.sub("", design)): + raise ValueError(f"{design!r} is an invalid formula for random effects. Use the '(1 | variable)' format.") + + random_effects = [match.group(1).strip() for match in _RANDOM_EFFECT.finditer(design)] + fixed = _RANDOM_EFFECT.sub("", design) + fixed = re.sub(r"\+\s*(?=\+|$)", "", fixed).strip().rstrip("+").strip() + if fixed in {"", "~"}: + fixed = "~ 1" + return fixed, random_effects + + +def random_effect_codes(obs: pd.DataFrame, variable: str) -> tuple[np.ndarray, int]: + """Integer level codes of a random intercept variable and its number of levels.""" + if variable not in obs.columns: + raise ValueError(f"Random effect variable {variable!r} is not a column of the sample metadata.") + levels = pd.Categorical(obs[variable]).remove_unused_categories() + if (levels.codes < 0).any(): + raise ValueError(f"Random effect variable {variable!r} contains missing values.") + if len(levels.categories) < 2: + raise ValueError( + f"Random effect variable {variable!r} has a single level, which cannot be told apart from the intercept." + ) + return levels.codes.astype(np.int64), len(levels.categories) + + +def random_effect_matrices(obs: pd.DataFrame, random_effects: Sequence[str]) -> list[tuple[str, np.ndarray]]: + """Build one indicator matrix of shape samples x levels per random intercept variable.""" + matrices = [] + for variable in random_effects: + codes, n_levels = random_effect_codes(obs, variable) + matrices.append((variable, np.eye(n_levels)[codes])) + return matrices + + +@njit(cache=True) +def _work_buffers(p: int, n_groups: int) -> tuple[np.ndarray, ...]: + """Scratch arrays of one random-intercept fit: ``cross``, ``chol``, ``inv``, ``beta``, ``vector``, ``resid``, ``shrink`` and ``scaled``.""" + return ( + np.empty((p + 1, p + 1)), + np.empty((p, p)), + np.empty((p, p)), + np.empty(p), + np.empty(p), + np.empty(n_groups), + np.empty(n_groups), + np.empty((n_groups, p)), + ) + + +@njit(cache=True, error_model="numpy") +def _cholesky(a: np.ndarray, p: int, chol: np.ndarray) -> bool: + """Lower Cholesky factor of the leading ``p`` x ``p`` block of ``a`` into ``chol``, False if that block is singular to working precision.""" + for j in range(p): + diagonal = a[j, j] + for k in range(j): + diagonal -= chol[j, k] ** 2 + if not diagonal > 1e-7 * a[j, j]: + return False + chol[j, j] = np.sqrt(diagonal) + for i in range(j + 1, p): + value = a[i, j] + for k in range(j): + value -= chol[i, k] * chol[j, k] + chol[i, j] = value / chol[j, j] + return True + + +@njit(cache=True, error_model="numpy") +def _cholesky_solve(chol: np.ndarray, p: int, x: np.ndarray) -> None: + """Solve ``chol @ chol.T @ solution = x`` for the solution, in place.""" + for i in range(p): + value = x[i] + for k in range(i): + value -= chol[i, k] * x[k] + x[i] = value / chol[i, i] + for i in range(p - 1, -1, -1): + value = x[i] + for k in range(i + 1, p): + value -= chol[k, i] * x[k] + x[i] = value / chol[i, i] + + +@njit(cache=True, error_model="numpy") +def _cholesky_inverse(chol: np.ndarray, p: int, inv: np.ndarray) -> None: + """Inverse of the matrix factored by ``chol`` into ``inv``.""" + for j in range(p): + for i in range(p): + inv[i, j] = 1.0 if i == j else 0.0 + _cholesky_solve(chol, p, inv[:, j]) + + +@njit(cache=True, error_model="numpy") +def _random_intercept_gls( + within: np.ndarray, sums: np.ndarray, sizes: np.ndarray, gamma: float, work: tuple[np.ndarray, ...] +) -> float: + """GLS fit of a random-intercept model, given the ratio ``gamma`` of the random-intercept to the residual variance. + + ``within`` holds the within-group cross-products and ``sums`` the per-group sums of the design columns followed by the response. + Fills the generalized cross-products, their Cholesky factor, the fixed effects and the per-group residual sums divided by ``1 + n_g * gamma`` of ``work``. + Returns the generalized residual sum of squares, NaN if the design is singular. + """ + cross, chol, _, beta, _, resid, shrink, _ = work + p = beta.shape[0] + for g in range(sizes.shape[0]): + shrink[g] = 1.0 / (1.0 + sizes[g] * gamma) + for a in range(p + 1): + for b in range(a, p + 1): + value = within[a, b] + for g in range(sizes.shape[0]): + value += shrink[g] / sizes[g] * sums[g, a] * sums[g, b] + cross[a, b] = cross[b, a] = value + if not _cholesky(cross, p, chol): + return np.nan + for a in range(p): + beta[a] = cross[a, p] + _cholesky_solve(chol, p, beta) + rss = float(cross[p, p]) + for a in range(p): + rss -= cross[a, p] * beta[a] + for g in range(sizes.shape[0]): + value = sums[g, p] + for a in range(p): + value -= sums[g, a] * beta[a] + resid[g] = value * shrink[g] + return rss + + +@njit(cache=True, error_model="numpy") +def _random_intercept_slope( + within: np.ndarray, sums: np.ndarray, sizes: np.ndarray, log_gamma: float, work: tuple[np.ndarray, ...] +) -> float: + """Derivative of the REML profile log-likelihood with respect to ``log(gamma)``.""" + _, chol, inv, beta, _, resid, shrink, _ = work + gamma = np.exp(log_gamma) + p = beta.shape[0] + rss = _random_intercept_gls(within, sums, sizes, gamma, work) + if np.isnan(rss): + return np.nan + _cholesky_inverse(chol, p, inv) + residual_term = 0.0 + size_term = 0.0 + quad_term = 0.0 + for g in range(sizes.shape[0]): + residual_term += resid[g] ** 2 + size_term += sizes[g] * shrink[g] + quad = 0.0 + for a in range(p): + for b in range(p): + quad += sums[g, a] * inv[a, b] * sums[g, b] + quad_term += shrink[g] ** 2 * quad + return 0.5 * gamma * ((sizes.sum() - p) * residual_term / rss - size_term + quad_term) + + +@njit(cache=True, error_model="numpy") +def _random_intercept_log_gamma( + within: np.ndarray, sums: np.ndarray, sizes: np.ndarray, work: tuple[np.ndarray, ...] +) -> float: + """Maximize the REML profile log-likelihood over ``log(gamma)`` by bracketing a root of its slope from 0, then bisection. + + Returns -inf if the maximum is on the boundary ``gamma = 0`` and NaN if the fit is degenerate. + """ + a, fa = 0.0, _random_intercept_slope(within, sums, sizes, 0.0, work) + b, fb = a, fa + step = 2.0 if fa > 0 else -2.0 + while fb != 0.0 and (fb > 0) == (fa > 0): + a, fa = b, fb + b += step + if b < -40.0: + return -np.inf + if b > 40.0 or not np.isfinite(fa): + return np.nan + fb = _random_intercept_slope(within, sums, sizes, b, work) + if not np.isfinite(fb): + return np.nan + mid = 0.5 * (a + b) + while fb != 0.0 and a != mid != b: + f_mid = _random_intercept_slope(within, sums, sizes, mid, work) + if (f_mid > 0) == (fa > 0): + a, fa = mid, f_mid + else: + b, fb = mid, f_mid + mid = 0.5 * (a + b) + return b if fb == 0.0 else mid + + +@njit(cache=True, error_model="numpy") +def _satterthwaite_df( # noqa: PLR0917 + sums: np.ndarray, sizes: np.ndarray, gamma: float, sigma2: float, contrast: np.ndarray, work: tuple[np.ndarray, ...] +) -> float: + """Satterthwaite denominator degrees of freedom of a contrast of the fixed effects, as lmerTest computes them. + + Uses the gradient of the contrast's variance and the REML observed information, both with respect to ``(gamma, sigma2)``. + Reads the inverse generalized cross-products and the shrunk residual sums of the fit from ``work``. + """ + _, _, inv, beta, vector, resid, shrink, scaled = work + p = beta.shape[0] + n_groups = sizes.shape[0] + for g in range(n_groups): + shrink[g] = 1.0 / (1.0 + sizes[g] * gamma) + for a in range(p): + value = 0.0 + for b in range(p): + value += sums[g, b] * inv[b, a] + scaled[g, a] = value + for a in range(p): + vector[a] = 0.0 + size_term = 0.0 + pair_term = 0.0 + quad_term = 0.0 + residual_term = 0.0 + weighted_residual_term = 0.0 + for g in range(n_groups): + quad = 0.0 + for a in range(p): + quad += scaled[g, a] * sums[g, a] + vector[a] += shrink[g] * resid[g] * sums[g, a] + size_term += (sizes[g] * shrink[g]) ** 2 + quad_term += sizes[g] * shrink[g] ** 3 * quad + residual_term += resid[g] ** 2 + weighted_residual_term += sizes[g] * shrink[g] * resid[g] ** 2 + for h in range(n_groups): + pair = 0.0 + for a in range(p): + pair += scaled[g, a] * sums[h, a] + pair_term += (shrink[g] * shrink[h] * pair) ** 2 + vector_quad = 0.0 + for a in range(p): + for b in range(p): + vector_quad += vector[a] * inv[a, b] * vector[b] + info_gg = 0.5 * ( + -size_term - pair_term + 2.0 * quad_term + (2.0 * weighted_residual_term - 2.0 * vector_quad) / sigma2 + ) + info_gs = residual_term / (2.0 * sigma2**2) + info_ss = (sizes.sum() - p) / (2.0 * sigma2**2) + grad_s = 0.0 + for a in range(p): + vector[a] = 0.0 + for b in range(p): + vector[a] += inv[a, b] * contrast[b] + grad_s += contrast[a] * vector[a] + grad_g = 0.0 + for g in range(n_groups): + projection = 0.0 + for a in range(p): + projection += sums[g, a] * vector[a] + grad_g += shrink[g] ** 2 * projection**2 + grad_g *= sigma2 + det = info_gg * info_ss - info_gs**2 + denominator = (grad_g**2 * info_ss - 2.0 * grad_g * grad_s * info_gs + grad_s**2 * info_gg) / det + return 2.0 * (sigma2 * grad_s) ** 2 / denominator + + +@njit(cache=True, error_model="numpy") +def _predictor_fit( # noqa: PLR0917 + x: np.ndarray, + centered: np.ndarray, + within: np.ndarray, + sums: np.ndarray, + sizes: np.ndarray, + groups: np.ndarray, + cols: np.ndarray, + unit: np.ndarray, +) -> tuple[float, float, float]: + """Fit the model of :func:`_random_intercept_fits` for one predictor ``x`` and return its coefficient, standard error and Satterthwaite degrees of freedom.""" + n = x.shape[0] + p = within.shape[0] - 1 + row_within = within.copy() + row_sums = sums.copy() + for i in range(n): + row_sums[groups[i], 1] += x[i] + for i in range(n): + row_within[1, 1] += (x[i] - row_sums[groups[i], 1] / sizes[groups[i]]) ** 2 + for a in range(cols.shape[0]): + value = 0.0 + for i in range(n): + value += x[i] * centered[a, i] + row_within[1, cols[a]] = row_within[cols[a], 1] = value + work = _work_buffers(p, sizes.shape[0]) + log_gamma = _random_intercept_log_gamma(row_within, row_sums, sizes, work) + if np.isnan(log_gamma): + return np.nan, np.nan, np.nan + gamma = np.exp(log_gamma) + rss = _random_intercept_gls(row_within, row_sums, sizes, gamma, work) + _, chol, inv, beta, _, _, _, _ = work + _cholesky_inverse(chol, p, inv) + sigma2 = rss / (n - p) + df = n - p if gamma == 0.0 else _satterthwaite_df(row_sums, sizes, gamma, sigma2, unit, work) + return beta[1], np.sqrt(sigma2 * inv[1, 1]), df + + +@njit(parallel=True, cache=True, error_model="numpy") +def _random_intercept_fits( + x: np.ndarray, fixed: np.ndarray, groups: np.ndarray, n_groups: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """REML fits of ``fixed[:, -1] ~ fixed[:, 0] + x[i] + fixed[:, 1:-1] + (1 | groups)`` for every row ``x[i]``. + + Returns the coefficients of ``x[i]``, their standard errors and their Satterthwaite degrees of freedom, like lmerTest's ``summary``. + """ + n, n_fixed = fixed.shape + m = n_fixed + 1 + sizes = np.zeros(n_groups) + fixed_sums = np.zeros((n_groups, n_fixed)) + for i in range(n): + sizes[groups[i]] += 1.0 + for a in range(n_fixed): + fixed_sums[groups[i], a] += fixed[i, a] + centered = np.empty((n_fixed, n)) + for i in range(n): + for a in range(n_fixed): + centered[a, i] = fixed[i, a] - fixed_sums[groups[i], a] / sizes[groups[i]] + cols = np.concatenate((np.zeros(1, dtype=np.int64), np.arange(2, m))) + within = np.zeros((m, m)) + sums = np.zeros((n_groups, m)) + for a in range(n_fixed): + for g in range(n_groups): + sums[g, cols[a]] = fixed_sums[g, a] + for b in range(n_fixed): + value = 0.0 + for i in range(n): + value += centered[a, i] * centered[b, i] + within[cols[a], cols[b]] = value + unit = np.zeros(m - 1) + unit[1] = 1.0 + estimates = np.full(x.shape[0], np.nan) + std_errors = np.full(x.shape[0], np.nan) + dfs = np.full(x.shape[0], np.nan) + for row in prange(x.shape[0]): + estimates[row], std_errors[row], dfs[row] = _predictor_fit( + x[row], centered, within, sums, sizes, groups, cols, unit + ) + return estimates, std_errors, dfs + + +@njit(cache=True, error_model="numpy") +def _response_fit( # noqa: PLR0917 + y: np.ndarray, + centered: np.ndarray, + within_design: np.ndarray, + design_sums: np.ndarray, + sizes: np.ndarray, + groups: np.ndarray, + coefficients: np.ndarray, + inverse: np.ndarray, + residual_sums: np.ndarray, +) -> tuple[float, float]: + """Fit the model of :func:`_random_intercept_response_fits` for one response ``y``. + + Writes its fixed effects, inverse generalized cross-products and shrunk per-group residual sums into the given arrays. + Returns the variance ratio and the residual variance, both NaN if the fit is degenerate. + """ + n = y.shape[0] + p = design_sums.shape[1] + n_groups = sizes.shape[0] + within = np.zeros((p + 1, p + 1)) + within[:p, :p] = within_design + sums = np.zeros((n_groups, p + 1)) + sums[:, :p] = design_sums + total = 0.0 + for i in range(n): + sums[groups[i], p] += y[i] + total += y[i] * y[i] + for a in range(p): + within[a, p] += y[i] * centered[i, a] + for a in range(p): + within[p, a] = within[a, p] + for g in range(n_groups): + total -= sums[g, p] ** 2 / sizes[g] + within[p, p] = total + work = _work_buffers(p, n_groups) + log_gamma = _random_intercept_log_gamma(within, sums, sizes, work) + if np.isnan(log_gamma): + return np.nan, np.nan + gamma = np.exp(log_gamma) + rss = _random_intercept_gls(within, sums, sizes, gamma, work) + _, chol, inv, beta, _, resid, _, _ = work + _cholesky_inverse(chol, p, inv) + coefficients[:] = beta + inverse[:] = inv + residual_sums[:] = resid + return gamma, rss / (n - p) + + +@njit(parallel=True, cache=True, error_model="numpy") +def _random_intercept_response_fits( + responses: np.ndarray, design: np.ndarray, groups: np.ndarray, n_groups: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """REML fits of ``responses[i] ~ design + (1 | groups)`` for every row of ``responses``. + + Returns the fixed effects, the inverse generalized cross-products of the design, the variance ratio, the residual variance and the shrunk per-group residual sums of every fit, followed by the per-group design sums and group sizes shared by all fits. + Rows whose fit is degenerate hold NaN. + """ + n_rows, n = responses.shape + p = design.shape[1] + sizes = np.zeros(n_groups) + design_sums = np.zeros((n_groups, p)) + for i in range(n): + sizes[groups[i]] += 1.0 + for a in range(p): + design_sums[groups[i], a] += design[i, a] + centered = np.empty((n, p)) + for i in range(n): + for a in range(p): + centered[i, a] = design[i, a] - design_sums[groups[i], a] / sizes[groups[i]] + within_design = np.zeros((p, p)) + for i in range(n): + for a in range(p): + for b in range(p): + within_design[a, b] += centered[i, a] * centered[i, b] + coefficients = np.full((n_rows, p), np.nan) + inverses = np.full((n_rows, p, p), np.nan) + gammas = np.full(n_rows, np.nan) + variances = np.full(n_rows, np.nan) + residual_sums = np.full((n_rows, n_groups), np.nan) + for row in prange(n_rows): + gammas[row], variances[row] = _response_fit( + responses[row], + centered, + within_design, + design_sums, + sizes, + groups, + coefficients[row], + inverses[row], + residual_sums[row], + ) + return coefficients, inverses, gammas, variances, residual_sums, design_sums, sizes + + +@njit(cache=True, error_model="numpy") +def _contrast_fit( # noqa: PLR0917 + contrast: np.ndarray, + coefficients: np.ndarray, + inverse: np.ndarray, + gamma: float, + variance: float, + residual_sums: np.ndarray, + design_sums: np.ndarray, + sizes: np.ndarray, +) -> tuple[float, float, float]: + """Estimate, standard error and Satterthwaite degrees of freedom of ``contrast`` for one fit of :func:`_response_fit`.""" + p = coefficients.shape[0] + work = _work_buffers(p, sizes.shape[0]) + _, _, inv, _, _, resid, _, _ = work + inv[:] = inverse + resid[:] = residual_sums + effect = 0.0 + quad = 0.0 + for a in range(p): + effect += contrast[a] * coefficients[a] + for b in range(p): + quad += contrast[a] * inv[a, b] * contrast[b] + df = sizes.sum() - p if gamma == 0.0 else _satterthwaite_df(design_sums, sizes, gamma, variance, contrast, work) + return effect, np.sqrt(variance * quad), df + + +@njit(parallel=True, cache=True, error_model="numpy") +def _random_intercept_contrast( # noqa: PLR0917 + contrast: np.ndarray, + coefficients: np.ndarray, + inverses: np.ndarray, + gammas: np.ndarray, + variances: np.ndarray, + residual_sums: np.ndarray, + design_sums: np.ndarray, + sizes: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Estimate, standard error and Satterthwaite degrees of freedom of one contrast for every fit of :func:`_random_intercept_response_fits`.""" + n_rows = coefficients.shape[0] + effects = np.full(n_rows, np.nan) + standard_errors = np.full(n_rows, np.nan) + dfs = np.full(n_rows, np.nan) + for row in prange(n_rows): + if not np.isnan(gammas[row]): + effects[row], standard_errors[row], dfs[row] = _contrast_fit( + contrast, + coefficients[row], + inverses[row], + gammas[row], + variances[row], + residual_sums[row], + design_sums, + sizes, + ) + return effects, standard_errors, dfs diff --git a/tests/tools/_differential_gene_expression/test_statsmodels.py b/tests/tools/_differential_gene_expression/test_statsmodels.py index 47752fb0..e1cefb62 100644 --- a/tests/tools/_differential_gene_expression/test_statsmodels.py +++ b/tests/tools/_differential_gene_expression/test_statsmodels.py @@ -1,6 +1,8 @@ from importlib.util import find_spec +import anndata as ad import numpy as np +import pandas as pd import pytest import statsmodels.api as sm @@ -8,6 +10,25 @@ pytestmark = pytest.mark.skip(reason="formulaic_contrasts and formulaic not available") +def _cells_of_donors(n_genes: int, *, seed: int = 0) -> ad.AnnData: + """Log-scale expression of cells from eight donors, with the condition assigned per donor and a cell-level covariate.""" + rng = np.random.default_rng(seed) + donors = np.repeat(np.arange(8), [12, 20, 15, 30, 9, 25, 18, 22]) + covariate = rng.normal(size=donors.size) + donor_effect = rng.normal(scale=0.8, size=(8, n_genes)) + donor_effect[:, 2::3] = 0 + X = rng.normal(size=(donors.size, n_genes)) + donor_effect[donors] + 0.3 * covariate[:, None] + obs = pd.DataFrame( + { + "donor": [f"D{donor}" for donor in donors], + "condition": np.where(donors % 2 == 0, "ctrl", "stim"), + "covariate": covariate, + }, + index=[f"cell{i}" for i in range(donors.size)], + ) + return ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=[f"gene{i}" for i in range(n_genes)])) + + @pytest.mark.parametrize( "kwargs", [{}, {"regression_model": sm.GLM, "family": sm.families.NegativeBinomial(alpha=1.0)}] ) @@ -26,3 +47,86 @@ def test_statsmodels(test_adata, kwargs): # TODO: there should be a test checking if, for a concrete example, the output p-values and effect sizes are what # we expect (-> frozen snapshot, that way we also get a heads-up if something changes upstream) + + +def test_random_intercept_matches_lmertest(): + """Reference values from lmerTest 3.2.1 (REML, Satterthwaite) on this exact data, including a fit on the boundary.""" + from pertpy.tools._differential_gene_expression import Statsmodels + + adata = _cells_of_donors(n_genes=3) + method = Statsmodels(adata, design="~condition + covariate + (1 | donor)") + method.fit() + contrasts = { + "stim": method.contrast(column="condition", baseline="ctrl", group_to_compare="stim"), + "covariate": np.array([0.0, 0.0, 1.0]), + "both": method.cond(condition="stim", covariate=1.0) - method.cond(condition="ctrl", covariate=0.0), + } + res = method.test_contrasts(contrasts).set_index(["contrast", "variable"]).sort_index() + expected = pd.DataFrame( + [ + ("stim", "gene0", 0.0479424000659, 0.40487594424, 0.909669910934), + ("covariate", "gene0", 0.206583084507, 0.0930645872892, 0.0280035797359), + ("both", "gene0", 0.254525484573, 0.418902673772, 0.563310062404), + ("stim", "gene1", -0.536172106102, 0.696514357941, 0.471370865701), + ("covariate", "gene1", 0.39318982076, 0.0850237207331, 8.35745118657e-06), + ("both", "gene1", -0.142982285342, 0.703391140188, 0.845551760376), + ("stim", "gene2", 0.0209764095948, 0.168433420581, 0.90105807902), + ("covariate", "gene2", 0.286895356907, 0.0844553833793, 0.000875172448353), + ("both", "gene2", 0.307871766502, 0.194457306979, 0.115502310835), + ], + columns=["contrast", "variable", "log_fc", "sd", "p_value"], + ).set_index(["contrast", "variable"]) + np.testing.assert_allclose(res[expected.columns].to_numpy(), expected.sort_index().to_numpy(), rtol=1e-5) + + +def test_random_intercept_is_calibrated(): + """A condition assigned per donor is judged by the eight donors, not the cells, so null genes are not called.""" + from pertpy.tools._differential_gene_expression import Statsmodels + + adata = _cells_of_donors(n_genes=300, seed=1) + method = Statsmodels(adata, design="~condition + (1 | donor)") + method.fit() + res = method.test_contrasts(method.contrast(column="condition", baseline="ctrl", group_to_compare="stim")) + assert (res["p_value"] < 0.05).mean() < 0.1 + + +def test_random_intercept_constant_variable(): + """Variables whose mixed model is degenerate get NaN statistics and are excluded from the FDR correction.""" + from pertpy.tools._differential_gene_expression import Statsmodels + + adata = _cells_of_donors(n_genes=4) + adata.X[:, 0] = 1.0 + method = Statsmodels(adata, design="~condition + (1 | donor)") + method.fit() + res = method.test_contrasts(np.array([0.0, 1.0])).set_index("variable") + assert res.loc["gene0", ["log_fc", "sd", "p_value", "adj_p_value"]].isna().all() + assert res.drop(index="gene0")["adj_p_value"].notna().all() + + +@pytest.mark.parametrize( + ("design", "match"), + [ + ("~condition + (1 | donor) + (1 | batch)", "single random intercept"), + ("~condition + (1 | missing)", "not a column"), + ("~condition + (1 | batch)", "single level"), + ("~condition + (1 | incomplete)", "missing values"), + ], +) +def test_random_intercept_rejects_invalid_designs(design, match): + from pertpy.tools._differential_gene_expression import Statsmodels + + adata = _cells_of_donors(n_genes=2) + adata.obs["batch"] = "b1" + adata.obs["incomplete"] = adata.obs["donor"].where(adata.obs.index != "cell0") + with pytest.raises(ValueError, match=match): + Statsmodels(adata, design=design) + + +@pytest.mark.skipif(find_spec("pydeseq2") is None, reason="pydeseq2 not available") +def test_random_intercept_rejected_by_other_methods(): + from pertpy.tools._differential_gene_expression import PyDESeq2 + + adata = _cells_of_donors(n_genes=2) + adata.X = np.ones_like(adata.X, dtype=int) + with pytest.raises(ValueError, match="does not support random effects"): + PyDESeq2(adata, design="~condition + (1 | donor)") diff --git a/tests/tools/test_milo_glmm.py b/tests/tools/test_milo_glmm.py index 4a5d5cc8..24554eaa 100644 --- a/tests/tools/test_milo_glmm.py +++ b/tests/tools/test_milo_glmm.py @@ -2,14 +2,8 @@ import pandas as pd import pytest -from pertpy.tools._milo_glmm import ( - between_within_df, - fit_nb_glmm, - fit_nb_glmm_nhoods, - has_separation, - parse_random_effects, - random_effect_matrices, -) +from pertpy.tools._milo_glmm import between_within_df, fit_nb_glmm, fit_nb_glmm_nhoods, has_separation +from pertpy.tools._mixed_models import parse_random_effects, random_effect_matrices @pytest.mark.parametrize(