From 393767bfee25ced1852752a28dc3602596c1f9b7 Mon Sep 17 00:00:00 2001 From: Sizerta Date: Thu, 1 Oct 2026 20:29:45 +0400 Subject: [PATCH 1/6] Add linear mixed model to the differential expression interface --- docs/api/tools_index.md | 1 + src/pertpy/tools/__init__.py | 3 +- .../_differential_gene_expression/__init__.py | 3 + .../_differential_gene_expression/_base.py | 12 +- .../_linear_mixed_model.py | 220 ++++++++++++++++++ .../test_linear_mixed_model.py | 166 +++++++++++++ 6 files changed, 401 insertions(+), 4 deletions(-) create mode 100644 src/pertpy/tools/_differential_gene_expression/_linear_mixed_model.py create mode 100644 tests/tools/_differential_gene_expression/test_linear_mixed_model.py diff --git a/docs/api/tools_index.md b/docs/api/tools_index.md index 8d956e2c..0590e1a3 100644 --- a/docs/api/tools_index.md +++ b/docs/api/tools_index.md @@ -25,6 +25,7 @@ Pertpy enables differential gene expression tests through a common interface tha tools.TTest tools.PermutationTest tools.Statsmodels + tools.LinearMixedModel ``` Example implementation: diff --git a/src/pertpy/tools/__init__.py b/src/pertpy/tools/__init__.py index f5939b62..5ae73bea 100644 --- a/src/pertpy/tools/__init__.py +++ b/src/pertpy/tools/__init__.py @@ -44,7 +44,7 @@ def __getattr__(name: str): ) from None with jax_import(name): return import_module("pertpy.tools._coda._tasccoda").Tasccoda - elif name in ["EdgeR", "PermutationTest", "PyDESeq2", "Statsmodels", "TTest", "WilcoxonTest"]: + elif name in ["EdgeR", "LinearMixedModel", "PermutationTest", "PyDESeq2", "Statsmodels", "TTest", "WilcoxonTest"]: module = import_module("pertpy.tools._differential_gene_expression") return getattr(module, name) @@ -67,6 +67,7 @@ def __dir__(): "TTest", "PermutationTest", "Statsmodels", + "LinearMixedModel", "DistanceTest", "Distance", "Enrichment", diff --git a/src/pertpy/tools/_differential_gene_expression/__init__.py b/src/pertpy/tools/_differential_gene_expression/__init__.py index cbf591d6..03941bd3 100644 --- a/src/pertpy/tools/_differential_gene_expression/__init__.py +++ b/src/pertpy/tools/_differential_gene_expression/__init__.py @@ -12,6 +12,7 @@ def __getattr__(name: str): "PyDESeq2": ["pydeseq2", "formulaic_contrasts", "formulaic"], "EdgeR": ["rpy2", "formulaic_contrasts", "formulaic"], "Statsmodels": ["formulaic_contrasts", "formulaic"], + "LinearMixedModel": ["formulaic_contrasts", "formulaic"], } if name in deps: @@ -23,6 +24,7 @@ def __getattr__(name: str): "PyDESeq2": "pertpy.tools._differential_gene_expression._pydeseq2", "EdgeR": "pertpy.tools._differential_gene_expression._edger", "Statsmodels": "pertpy.tools._differential_gene_expression._statsmodels", + "LinearMixedModel": "pertpy.tools._differential_gene_expression._linear_mixed_model", } module = import_module(module_map[name]) @@ -53,6 +55,7 @@ def _get_available_methods(): "EdgeR", "PyDESeq2", "Statsmodels", + "LinearMixedModel", "SimpleComparisonBase", "WilcoxonTest", "TTest", diff --git a/src/pertpy/tools/_differential_gene_expression/_base.py b/src/pertpy/tools/_differential_gene_expression/_base.py index 10213521..2118a704 100644 --- a/src/pertpy/tools/_differential_gene_expression/_base.py +++ b/src/pertpy/tools/_differential_gene_expression/_base.py @@ -1060,9 +1060,7 @@ def compare_groups( fit_kwargs=MappingProxyType({}), test_kwargs=MappingProxyType({}), ): - design = f"~{column}" - if paired_by is not None: - design += f"+{paired_by}" + design = cls._comparison_design(column, paired_by) if isinstance(groups_to_compare, str): groups_to_compare = [groups_to_compare] model = cls(adata, design=design, mask=mask, layer=layer) @@ -1079,6 +1077,14 @@ def compare_groups( return de_res + @classmethod + def _comparison_design(cls, column: str, paired_by: str | None) -> str: + """Build the design formula used by `compare_groups`.""" + design = f"~{column}" + if paired_by is not None: + design += f"+{paired_by}" + return design + @property def variables(self): """Get the names of the variables used in the model definition.""" diff --git a/src/pertpy/tools/_differential_gene_expression/_linear_mixed_model.py b/src/pertpy/tools/_differential_gene_expression/_linear_mixed_model.py new file mode 100644 index 00000000..66b50d99 --- /dev/null +++ b/src/pertpy/tools/_differential_gene_expression/_linear_mixed_model.py @@ -0,0 +1,220 @@ +import warnings + +import numpy as np +import pandas as pd +import statsmodels.api as sm +from fast_array_utils.conv import to_dense +from joblib import delayed, effective_n_jobs +from scipy import stats +from statsmodels.stats.multitest import fdrcorrection +from tqdm.auto import tqdm + +from pertpy._parallel import _MAX_BLOCK_ELEMENTS, _block_slices, _parallelize_with_joblib +from pertpy._types import CSBase +from pertpy.tools._milo_glmm import parse_random_effects + +from ._base import LinearModelBase +from ._checks import check_is_numeric_matrix + + +def _fit_block(endog_block: np.ndarray, exog: np.ndarray, groups: np.ndarray, kwargs: dict) -> list: + """Fit one mixed model per column of `endog_block`. + + The default gradient-based optimizers of statsmodels regularly reach the optimum but fail their gradient + tolerance check. Unless an optimizer was requested explicitly, such fits are repeated with derivative-free + optimizers, as the statsmodels convergence warning suggests. + Per-variable warnings are silenced here and reported once, via the `converged` column of the results. + Variables whose model cannot be fit at all (e.g. constant expression) yield `None`. + """ + models = [] + for i in range(endog_block.shape[1]): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + model = sm.MixedLM(endog_block[:, i], exog, groups=groups) + result = model.fit(**kwargs) + if not result.converged and "method" not in kwargs: + result = model.fit(**kwargs, method=["powell", "nm"]) + models.append(result) + except (np.linalg.LinAlgError, ValueError): + models.append(None) + return models + + +def _between_within_df(design: np.ndarray, groups: np.ndarray) -> tuple[np.ndarray, int, int]: + """Classify design columns and compute between-within denominator degrees of freedom. + + A column is *between-group* if it is constant within every group (e.g. a condition assigned per donor, + or the intercept), otherwise *within-group* (e.g. a condition that varies across the cells of a donor). + + Returns: + A boolean mask of the between-group columns, the between-group degrees of freedom + (number of groups minus the rank of the between-group columns) and the within-group degrees of freedom + (number of observations minus the rank of the group indicators together with the within-group columns). + """ + codes = pd.factorize(groups)[0] + n_groups = int(codes.max()) + 1 + group_means = np.zeros((n_groups, design.shape[1])) + np.add.at(group_means, codes, design) + group_means /= np.bincount(codes, minlength=n_groups)[:, None] + is_between = np.all(np.isclose(design, group_means[codes]), axis=0) + + df_between = n_groups - np.linalg.matrix_rank(group_means[:, is_between]) + indicators = np.eye(n_groups)[codes] + df_within = design.shape[0] - np.linalg.matrix_rank(np.column_stack([indicators, design[:, ~is_between]])) + return is_between, int(df_between), int(df_within) + + +class LinearMixedModel(LinearModelBase): + """Differential expression test using a linear mixed model with a random intercept. + + Fits one :class:`statsmodels.regression.mixed_linear_model.MixedLM` per variable. + This allows testing all cells while accounting for the sample they come from, + or testing pseudobulk data whose design requires a random effect. + As for :class:`~pertpy.tools.Statsmodels`, the data should be on a continuous scale, + e.g. log-normalized expression. + + The random intercept is specified in the formula using the ``(1 | variable)`` syntax of lme4, + for example ``"~condition + (1 | donor)"``. Exactly one random intercept is supported. + + Fixed effects are tested with a t-test on the requested contrast. Because statsmodels only provides + a normal approximation, which is anti-conservative when there are few groups, the denominator degrees + of freedom follow the between-within method: effects that are constant within each group + (e.g. a condition assigned per donor) are tested against the number of groups minus the number of such + effects, effects varying within groups against the within-group residual degrees of freedom. + Contrasts mixing both kinds use the smaller of the two. + + Examples: + >>> import pertpy as pt + >>> model = pt.tl.LinearMixedModel(adata, design="~condition + (1 | donor)") # doctest: +SKIP + >>> model.fit() # doctest: +SKIP + >>> res = model.test_contrasts( + ... model.contrast(column="condition", baseline="control", group_to_compare="treated") + ... ) # doctest: +SKIP + """ + + def __init__(self, adata, design, *, mask=None, layer=None, **kwargs): + if not isinstance(design, str): + raise TypeError( + "LinearMixedModel requires a formula with a random intercept, e.g. '~condition + (1 | donor)'." + ) + fixed_design, random_effects = parse_random_effects(design) + if len(random_effects) != 1: + raise ValueError( + f"LinearMixedModel requires exactly one random intercept in the formula, e.g. '~condition + (1 | donor)', " + f"but {design!r} has {len(random_effects)}." + ) + (random_effect,) = random_effects + if random_effect not in adata.obs.columns: + raise ValueError(f"Random effect variable {random_effect!r} is not a column of adata.obs.") + if adata.obs[random_effect].isna().any(): + raise ValueError(f"Random effect variable {random_effect!r} contains missing values.") + if adata.obs[random_effect].nunique() < 2: + raise ValueError(f"Random effect variable {random_effect!r} needs at least two levels.") + + super().__init__(adata, fixed_design, mask=mask, layer=layer, **kwargs) + self.random_effect = random_effect + self.groups = adata.obs[random_effect].astype(str).to_numpy() + self.is_between, self.df_between, self.df_within = _between_within_df( + np.asarray(self.design, dtype=float), self.groups + ) + + @classmethod + def _comparison_design(cls, column: str, paired_by: str | None) -> str: + if paired_by is None: + raise ValueError( + "LinearMixedModel.compare_groups requires `paired_by`, which enters the model as the random intercept." + ) + return f"~{column} + (1 | {paired_by})" + + def _check_counts(self): + check_is_numeric_matrix(self.data) + + def fit(self, *, reml: bool = True, n_jobs: int | None = None, **kwargs) -> None: + """Fit one linear mixed model per variable. + + Args: + reml: Whether to fit by restricted maximum likelihood (recommended) instead of maximum likelihood. + n_jobs: Number of jobs to distribute the variables over, passed to `joblib.Parallel`. + None means one job, -1 all available cores. + **kwargs: Additional arguments passed to :meth:`statsmodels.regression.mixed_linear_model.MixedLM.fit`, + e.g. `method`. + + Examples: + >>> import pertpy as pt + >>> model = pt.tl.LinearMixedModel(adata, design="~condition + (1 | donor)") # doctest: +SKIP + >>> model.fit() # doctest: +SKIP + """ + data = self.data + # Slicing out variables is significantly more efficient in csc format. + if isinstance(data, CSBase): + data = data.tocsc() + + exog = np.asarray(self.design, dtype=float) + fit_kwargs = {"reml": reml, **kwargs} + n_workers = effective_n_jobs(n_jobs if n_jobs is not None else 1) + blocks = _block_slices( + self.adata.n_vars, + n_blocks=self.adata.n_vars if n_workers == 1 else 4 * n_workers, + max_block_size=_MAX_BLOCK_ELEMENTS // max(self.adata.n_obs, 1), + ) + self.models = [ + model + for block_models in _parallelize_with_joblib( + (delayed(_fit_block)(to_dense(data[:, block]), exog, self.groups, fit_kwargs) for block in blocks), + total=len(blocks), + n_jobs=n_jobs, + ) + for model in block_models + ] + n_failed = sum(model is None or not model.converged for model in self.models) + if n_failed: + warnings.warn( + f"The mixed model did not converge (or could not be fit) for {n_failed} of {len(self.models)} variables. " + "See the `converged` column of the results.", + UserWarning, + stacklevel=2, + ) + self._fitted = True + + def _contrast_df(self, contrast: np.ndarray) -> int: + """Denominator degrees of freedom for a contrast, following the between-within method.""" + involved = ~np.isclose(contrast, 0) + df = min( + self.df_between if (involved & self.is_between).any() else np.inf, + self.df_within if (involved & ~self.is_between).any() else np.inf, + ) + if df < 1: + raise ValueError( + f"The contrast cannot be tested: it involves effects that are constant within `{self.random_effect}`, " + f"but there are not enough levels of `{self.random_effect}` left to estimate them " + f"(between-group degrees of freedom: {self.df_between})." + ) + return int(df) + + def _test_single_contrast(self, contrast, **kwargs) -> pd.DataFrame: + contrast = np.asarray(contrast, dtype=float) + df = self._contrast_df(contrast) + res = [] + for var, mod in zip(tqdm(self.adata.var_names), self.models, strict=False): + row = {"variable": var, "p_value": np.nan, "t_value": np.nan, "sd": np.nan, "log_fc": np.nan} + if mod is not None: + # MixedLMResults.t_test expects a 2D restriction matrix over the fixed effects. + t_test = mod.t_test(np.atleast_2d(contrast)) + effect, sd = float(np.ravel(t_test.effect)[0]), float(np.ravel(t_test.sd)[0]) + t_value = effect / sd if sd > 0 else np.nan + row |= { + "p_value": 2 * stats.t.sf(abs(t_value), df) if np.isfinite(t_value) else np.nan, + "t_value": t_value, + "sd": sd, + "log_fc": effect, + } + row["converged"] = mod is not None and bool(mod.converged) + res.append(row) + + res_df = pd.DataFrame(res).sort_values("p_value") + tested = res_df["p_value"].notna().to_numpy() + adj = np.full(len(res_df), np.nan) + if tested.any(): + adj[tested] = fdrcorrection(res_df["p_value"].to_numpy()[tested])[1] + return res_df.assign(adj_p_value=adj) diff --git a/tests/tools/_differential_gene_expression/test_linear_mixed_model.py b/tests/tools/_differential_gene_expression/test_linear_mixed_model.py new file mode 100644 index 00000000..c6e54ed7 --- /dev/null +++ b/tests/tools/_differential_gene_expression/test_linear_mixed_model.py @@ -0,0 +1,166 @@ +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 +from scipy import stats + +if find_spec("formulaic_contrasts") is None or find_spec("formulaic") is None: + pytestmark = pytest.mark.skip(reason="formulaic_contrasts and formulaic not available") + + +def _donor_level_adata(n_donors: int = 6, n_cells: int = 30, n_genes: int = 5, *, seed: int = 0) -> ad.AnnData: + """Log-scale expression with a random donor effect; the condition is assigned per donor (between-donor).""" + rng = np.random.default_rng(seed) + donor = np.repeat([f"D{i}" for i in range(n_donors)], n_cells) + condition = np.where(np.repeat(np.arange(n_donors), n_cells) < n_donors // 2, "A", "B") + donor_effect = rng.normal(0, 0.7, (n_donors, n_genes)) + X = rng.normal(size=(n_donors * n_cells, n_genes)) + np.repeat(donor_effect, n_cells, axis=0) + obs = pd.DataFrame({"donor": donor, "condition": condition}, index=[f"c{i}" for i in range(len(donor))]) + return ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=[f"gene{i}" for i in range(n_genes)])) + + +def test_linear_mixed_model(test_adata_minimal): + """Check that the method can be initialized and fitted, and perform basic checks on the result of test_contrasts.""" + from pertpy.tools._differential_gene_expression import LinearMixedModel + + method = LinearMixedModel(adata=test_adata_minimal, design="~condition + (1 | donor)") + assert not method.is_fitted + method.fit() + assert method.is_fitted + res_df = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) + assert len(res_df) == test_adata_minimal.n_vars + assert {"variable", "p_value", "t_value", "sd", "log_fc", "adj_p_value", "converged"} <= set(res_df.columns) + # gene1 has a large difference between conditions + assert res_df.set_index("variable").loc["gene1", "p_value"] < 0.05 + + +def test_matches_statsmodels(): + """Effect sizes come from statsmodels' MixedLM; p-values use the between-group degrees of freedom.""" + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata(n_genes=3) + method = LinearMixedModel(adata, design="~condition + (1 | donor)") + method.fit() + contrast = method.contrast(column="condition", baseline="A", group_to_compare="B") + res = method.test_contrasts(contrast).set_index("variable") + + exog = np.column_stack([np.ones(adata.n_obs), (adata.obs["condition"] == "B").to_numpy(float)]) + for i, gene in enumerate(adata.var_names): + expected = sm.MixedLM(adata.X[:, i], exog, groups=adata.obs["donor"].to_numpy()).fit(reml=True) + np.testing.assert_allclose(res.loc[gene, "log_fc"], expected.fe_params[1], rtol=1e-6) + np.testing.assert_allclose(res.loc[gene, "sd"], expected.bse_fe[1], rtol=1e-6) + t_value = expected.fe_params[1] / expected.bse_fe[1] + # 6 donors minus intercept and condition, both constant within donors + np.testing.assert_allclose(res.loc[gene, "p_value"], 2 * stats.t.sf(abs(t_value), df=4), rtol=1e-6) + + +def test_between_within_df(): + from pertpy.tools._differential_gene_expression._linear_mixed_model import _between_within_df + + groups = np.repeat(["D0", "D1", "D2", "D3"], 5) + between = np.repeat([0.0, 0.0, 1.0, 1.0], 5) # constant within each donor + within = np.tile([0.0, 1.0, 0.0, 1.0, 1.0], 4) # varies within donors + design = np.column_stack([np.ones(20), between, within]) + + is_between, df_between, df_within = _between_within_df(design, groups) + + np.testing.assert_array_equal(is_between, [True, True, False]) + assert df_between == 4 - 2 # donors minus intercept and between-donor effect + assert df_within == 20 - 4 - 1 # cells minus donor indicators and within-donor effect + + +def test_type_i_error_between_donor_effect(): + """With few donors, the default normal approximation would be anti-conservative; the t-test must not be.""" + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata(n_genes=200, seed=1) + method = LinearMixedModel(adata, design="~condition + (1 | donor)") + method.fit() + res = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) + assert (res["p_value"] < 0.05).mean() < 0.1 + + +def test_optimizer_fallback_converges(recwarn): + """Fits that fail statsmodels' gradient check are refit with derivative-free optimizers, so clean data converges.""" + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata(n_genes=200, seed=1) + method = LinearMixedModel(adata, design="~condition + (1 | donor)") + method.fit() + res = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) + assert res["converged"].all() + assert not [w for w in recwarn if "did not converge" in str(w.message)] + + +def test_compare_groups(): + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata() + res = LinearMixedModel.compare_groups( + adata, column="condition", baseline="A", groups_to_compare="B", paired_by="donor" + ) + assert len(res) == adata.n_vars + assert set(res["contrast"]) == {"B"} + + with pytest.raises(ValueError, match="requires `paired_by`"): + LinearMixedModel.compare_groups(adata, column="condition", baseline="A", groups_to_compare="B") + + +@pytest.mark.parametrize( + ("design", "error", "match"), + [ + ("~condition", ValueError, "exactly one random intercept"), + ("~condition + (1 | donor) + (1 | batch)", ValueError, "exactly one random intercept"), + ("~condition + (1 | missing)", ValueError, "not a column of adata.obs"), + (np.ones((180, 1)), TypeError, "requires a formula"), + ], +) +def test_invalid_design(design, error, match): + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata() + adata.obs["batch"] = "b1" + with pytest.raises(error, match=match): + LinearMixedModel(adata, design=design) + + +def test_missing_random_effect_values(): + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata() + adata.obs["donor"] = adata.obs["donor"].astype(object) + adata.obs.iloc[0, adata.obs.columns.get_loc("donor")] = None + with pytest.raises(ValueError, match="missing values"): + LinearMixedModel(adata, design="~condition + (1 | donor)") + + +def test_untestable_contrast(): + """One donor per condition leaves no degrees of freedom for a between-donor effect.""" + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata(n_donors=2) + method = LinearMixedModel(adata, design="~condition + (1 | donor)") + method.fit() + with pytest.raises(ValueError, match="cannot be tested"): + method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) + + +def test_constant_variable(): + """Variables that cannot be fit get NaN statistics and are excluded from the FDR correction.""" + from pertpy.tools._differential_gene_expression import LinearMixedModel + + adata = _donor_level_adata(n_genes=4) + adata.X[:, 0] = 0 + method = LinearMixedModel(adata, design="~condition + (1 | donor)") + with pytest.warns(UserWarning, match="did not converge"): + method.fit() + res = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) + res = res.set_index("variable") + + assert np.isnan(res.loc["gene0", "p_value"]) + assert np.isnan(res.loc["gene0", "adj_p_value"]) + assert not res.loc["gene0", "converged"] + assert res.drop(index="gene0")["adj_p_value"].notna().all() From cdd619ef9d9e719740461f9d3a1cb201de2e1c72 Mon Sep 17 00:00:00 2001 From: Lukas Heumos Date: Wed, 21 Oct 2026 16:57:01 +0200 Subject: [PATCH 2/6] Fit random intercepts in Statsmodels through a shared mixed-model kernel Random intercepts enter the differential expression interface as a "(1 | donor)" term in the design of pt.tl.Statsmodels instead of a separate LinearMixedModel class, so the familiar formula, cond and contrast API and Milo's random effect syntax carry over unchanged. LinearModelBase parses the term for every method and PyDESeq2 and EdgeR raise a clear error, while compare_groups keeps paired_by as a fixed effect. The fit uses the REML random-intercept kernel that Dialogue already had, generalized from one predictor to many responses on a shared design and from the coefficient at index 1 to arbitrary contrasts. Estimates, standard errors, Satterthwaite degrees of freedom and p-values match lmerTest::contest1D to 1e-6 on between-donor, within-donor and mixed contrasts and on boundary fits. The kernels, parse_random_effects and random_effect_matrices move to pertpy/tools/_mixed_models.py; random_effect_matrices now drops unused categories and rejects missing levels, which the kernel cannot index. The kernels are restructured so that the prange body is a single call into a serial njit worker with explicit loops and preallocated buffers. Array expressions inside the prange body spawn one nested parallel region per gene, which made 64 threads 15 to 60 times slower than one; the Statsmodels fit on 15k cells x 2000 genes drops from 17 s to 0.2 s and Dialogue's test_celltype_pairs on the example data from 17 s to 6 s, with identical results. The numba thread cap is shared through _parallel._numba_threads. Degenerate genes (constant expression) get NaN statistics and are excluded from the FDR correction through the existing fdr_correction helper. --- docs/api/tools_index.md | 2 +- docs/changelog.md | 4 + src/pertpy/_parallel.py | 13 + src/pertpy/tools/__init__.py | 3 +- src/pertpy/tools/_dialogue.py | 191 +------ .../_differential_gene_expression/__init__.py | 3 - .../_differential_gene_expression/_base.py | 25 +- .../_linear_mixed_model.py | 220 -------- .../_statsmodels.py | 105 +++- src/pertpy/tools/_milo.py | 3 +- src/pertpy/tools/_milo_glmm.py | 36 -- src/pertpy/tools/_mixed_models.py | 513 ++++++++++++++++++ .../test_linear_mixed_model.py | 166 ------ .../test_statsmodels.py | 104 ++++ tests/tools/test_milo_glmm.py | 10 +- 15 files changed, 754 insertions(+), 644 deletions(-) delete mode 100644 src/pertpy/tools/_differential_gene_expression/_linear_mixed_model.py create mode 100644 src/pertpy/tools/_mixed_models.py delete mode 100644 tests/tools/_differential_gene_expression/test_linear_mixed_model.py diff --git a/docs/api/tools_index.md b/docs/api/tools_index.md index 0590e1a3..3d1efcab 100644 --- a/docs/api/tools_index.md +++ b/docs/api/tools_index.md @@ -13,6 +13,7 @@ Differential gene expression involves the quantitative comparison of gene expression levels between two or more groups, such as different cell types, tissues, or conditions to discern genes that are significantly up- or downregulated in response to specific biological contexts or stimuli. Pertpy enables differential gene expression tests through a common interface that supports complex designs. +{class}`~pertpy.tools.Statsmodels` also fits linear mixed models with a random intercept, written as `(1 | donor)` in the design, to test single cells while accounting for the samples they come from. ```{eval-rst} .. autosummary:: @@ -25,7 +26,6 @@ Pertpy enables differential gene expression tests through a common interface tha tools.TTest tools.PermutationTest tools.Statsmodels - tools.LinearMixedModel ``` Example implementation: diff --git a/docs/changelog.md b/docs/changelog.md index 67186e55..52e33506 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/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/__init__.py b/src/pertpy/tools/__init__.py index 5ae73bea..f5939b62 100644 --- a/src/pertpy/tools/__init__.py +++ b/src/pertpy/tools/__init__.py @@ -44,7 +44,7 @@ def __getattr__(name: str): ) from None with jax_import(name): return import_module("pertpy.tools._coda._tasccoda").Tasccoda - elif name in ["EdgeR", "LinearMixedModel", "PermutationTest", "PyDESeq2", "Statsmodels", "TTest", "WilcoxonTest"]: + elif name in ["EdgeR", "PermutationTest", "PyDESeq2", "Statsmodels", "TTest", "WilcoxonTest"]: module = import_module("pertpy.tools._differential_gene_expression") return getattr(module, name) @@ -67,7 +67,6 @@ def __dir__(): "TTest", "PermutationTest", "Statsmodels", - "LinearMixedModel", "DistanceTest", "Distance", "Enrichment", 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/__init__.py b/src/pertpy/tools/_differential_gene_expression/__init__.py index 03941bd3..cbf591d6 100644 --- a/src/pertpy/tools/_differential_gene_expression/__init__.py +++ b/src/pertpy/tools/_differential_gene_expression/__init__.py @@ -12,7 +12,6 @@ def __getattr__(name: str): "PyDESeq2": ["pydeseq2", "formulaic_contrasts", "formulaic"], "EdgeR": ["rpy2", "formulaic_contrasts", "formulaic"], "Statsmodels": ["formulaic_contrasts", "formulaic"], - "LinearMixedModel": ["formulaic_contrasts", "formulaic"], } if name in deps: @@ -24,7 +23,6 @@ def __getattr__(name: str): "PyDESeq2": "pertpy.tools._differential_gene_expression._pydeseq2", "EdgeR": "pertpy.tools._differential_gene_expression._edger", "Statsmodels": "pertpy.tools._differential_gene_expression._statsmodels", - "LinearMixedModel": "pertpy.tools._differential_gene_expression._linear_mixed_model", } module = import_module(module_map[name]) @@ -55,7 +53,6 @@ def _get_available_methods(): "EdgeR", "PyDESeq2", "Statsmodels", - "LinearMixedModel", "SimpleComparisonBase", "WilcoxonTest", "TTest", diff --git a/src/pertpy/tools/_differential_gene_expression/_base.py b/src/pertpy/tools/_differential_gene_expression/_base.py index 2118a704..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: @@ -1060,7 +1071,9 @@ def compare_groups( fit_kwargs=MappingProxyType({}), test_kwargs=MappingProxyType({}), ): - design = cls._comparison_design(column, paired_by) + design = f"~{column}" + if paired_by is not None: + design += f"+{paired_by}" if isinstance(groups_to_compare, str): groups_to_compare = [groups_to_compare] model = cls(adata, design=design, mask=mask, layer=layer) @@ -1077,14 +1090,6 @@ def compare_groups( return de_res - @classmethod - def _comparison_design(cls, column: str, paired_by: str | None) -> str: - """Build the design formula used by `compare_groups`.""" - design = f"~{column}" - if paired_by is not None: - design += f"+{paired_by}" - return design - @property def variables(self): """Get the names of the variables used in the model definition.""" @@ -1107,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/_linear_mixed_model.py b/src/pertpy/tools/_differential_gene_expression/_linear_mixed_model.py deleted file mode 100644 index 66b50d99..00000000 --- a/src/pertpy/tools/_differential_gene_expression/_linear_mixed_model.py +++ /dev/null @@ -1,220 +0,0 @@ -import warnings - -import numpy as np -import pandas as pd -import statsmodels.api as sm -from fast_array_utils.conv import to_dense -from joblib import delayed, effective_n_jobs -from scipy import stats -from statsmodels.stats.multitest import fdrcorrection -from tqdm.auto import tqdm - -from pertpy._parallel import _MAX_BLOCK_ELEMENTS, _block_slices, _parallelize_with_joblib -from pertpy._types import CSBase -from pertpy.tools._milo_glmm import parse_random_effects - -from ._base import LinearModelBase -from ._checks import check_is_numeric_matrix - - -def _fit_block(endog_block: np.ndarray, exog: np.ndarray, groups: np.ndarray, kwargs: dict) -> list: - """Fit one mixed model per column of `endog_block`. - - The default gradient-based optimizers of statsmodels regularly reach the optimum but fail their gradient - tolerance check. Unless an optimizer was requested explicitly, such fits are repeated with derivative-free - optimizers, as the statsmodels convergence warning suggests. - Per-variable warnings are silenced here and reported once, via the `converged` column of the results. - Variables whose model cannot be fit at all (e.g. constant expression) yield `None`. - """ - models = [] - for i in range(endog_block.shape[1]): - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - try: - model = sm.MixedLM(endog_block[:, i], exog, groups=groups) - result = model.fit(**kwargs) - if not result.converged and "method" not in kwargs: - result = model.fit(**kwargs, method=["powell", "nm"]) - models.append(result) - except (np.linalg.LinAlgError, ValueError): - models.append(None) - return models - - -def _between_within_df(design: np.ndarray, groups: np.ndarray) -> tuple[np.ndarray, int, int]: - """Classify design columns and compute between-within denominator degrees of freedom. - - A column is *between-group* if it is constant within every group (e.g. a condition assigned per donor, - or the intercept), otherwise *within-group* (e.g. a condition that varies across the cells of a donor). - - Returns: - A boolean mask of the between-group columns, the between-group degrees of freedom - (number of groups minus the rank of the between-group columns) and the within-group degrees of freedom - (number of observations minus the rank of the group indicators together with the within-group columns). - """ - codes = pd.factorize(groups)[0] - n_groups = int(codes.max()) + 1 - group_means = np.zeros((n_groups, design.shape[1])) - np.add.at(group_means, codes, design) - group_means /= np.bincount(codes, minlength=n_groups)[:, None] - is_between = np.all(np.isclose(design, group_means[codes]), axis=0) - - df_between = n_groups - np.linalg.matrix_rank(group_means[:, is_between]) - indicators = np.eye(n_groups)[codes] - df_within = design.shape[0] - np.linalg.matrix_rank(np.column_stack([indicators, design[:, ~is_between]])) - return is_between, int(df_between), int(df_within) - - -class LinearMixedModel(LinearModelBase): - """Differential expression test using a linear mixed model with a random intercept. - - Fits one :class:`statsmodels.regression.mixed_linear_model.MixedLM` per variable. - This allows testing all cells while accounting for the sample they come from, - or testing pseudobulk data whose design requires a random effect. - As for :class:`~pertpy.tools.Statsmodels`, the data should be on a continuous scale, - e.g. log-normalized expression. - - The random intercept is specified in the formula using the ``(1 | variable)`` syntax of lme4, - for example ``"~condition + (1 | donor)"``. Exactly one random intercept is supported. - - Fixed effects are tested with a t-test on the requested contrast. Because statsmodels only provides - a normal approximation, which is anti-conservative when there are few groups, the denominator degrees - of freedom follow the between-within method: effects that are constant within each group - (e.g. a condition assigned per donor) are tested against the number of groups minus the number of such - effects, effects varying within groups against the within-group residual degrees of freedom. - Contrasts mixing both kinds use the smaller of the two. - - Examples: - >>> import pertpy as pt - >>> model = pt.tl.LinearMixedModel(adata, design="~condition + (1 | donor)") # doctest: +SKIP - >>> model.fit() # doctest: +SKIP - >>> res = model.test_contrasts( - ... model.contrast(column="condition", baseline="control", group_to_compare="treated") - ... ) # doctest: +SKIP - """ - - def __init__(self, adata, design, *, mask=None, layer=None, **kwargs): - if not isinstance(design, str): - raise TypeError( - "LinearMixedModel requires a formula with a random intercept, e.g. '~condition + (1 | donor)'." - ) - fixed_design, random_effects = parse_random_effects(design) - if len(random_effects) != 1: - raise ValueError( - f"LinearMixedModel requires exactly one random intercept in the formula, e.g. '~condition + (1 | donor)', " - f"but {design!r} has {len(random_effects)}." - ) - (random_effect,) = random_effects - if random_effect not in adata.obs.columns: - raise ValueError(f"Random effect variable {random_effect!r} is not a column of adata.obs.") - if adata.obs[random_effect].isna().any(): - raise ValueError(f"Random effect variable {random_effect!r} contains missing values.") - if adata.obs[random_effect].nunique() < 2: - raise ValueError(f"Random effect variable {random_effect!r} needs at least two levels.") - - super().__init__(adata, fixed_design, mask=mask, layer=layer, **kwargs) - self.random_effect = random_effect - self.groups = adata.obs[random_effect].astype(str).to_numpy() - self.is_between, self.df_between, self.df_within = _between_within_df( - np.asarray(self.design, dtype=float), self.groups - ) - - @classmethod - def _comparison_design(cls, column: str, paired_by: str | None) -> str: - if paired_by is None: - raise ValueError( - "LinearMixedModel.compare_groups requires `paired_by`, which enters the model as the random intercept." - ) - return f"~{column} + (1 | {paired_by})" - - def _check_counts(self): - check_is_numeric_matrix(self.data) - - def fit(self, *, reml: bool = True, n_jobs: int | None = None, **kwargs) -> None: - """Fit one linear mixed model per variable. - - Args: - reml: Whether to fit by restricted maximum likelihood (recommended) instead of maximum likelihood. - n_jobs: Number of jobs to distribute the variables over, passed to `joblib.Parallel`. - None means one job, -1 all available cores. - **kwargs: Additional arguments passed to :meth:`statsmodels.regression.mixed_linear_model.MixedLM.fit`, - e.g. `method`. - - Examples: - >>> import pertpy as pt - >>> model = pt.tl.LinearMixedModel(adata, design="~condition + (1 | donor)") # doctest: +SKIP - >>> model.fit() # doctest: +SKIP - """ - data = self.data - # Slicing out variables is significantly more efficient in csc format. - if isinstance(data, CSBase): - data = data.tocsc() - - exog = np.asarray(self.design, dtype=float) - fit_kwargs = {"reml": reml, **kwargs} - n_workers = effective_n_jobs(n_jobs if n_jobs is not None else 1) - blocks = _block_slices( - self.adata.n_vars, - n_blocks=self.adata.n_vars if n_workers == 1 else 4 * n_workers, - max_block_size=_MAX_BLOCK_ELEMENTS // max(self.adata.n_obs, 1), - ) - self.models = [ - model - for block_models in _parallelize_with_joblib( - (delayed(_fit_block)(to_dense(data[:, block]), exog, self.groups, fit_kwargs) for block in blocks), - total=len(blocks), - n_jobs=n_jobs, - ) - for model in block_models - ] - n_failed = sum(model is None or not model.converged for model in self.models) - if n_failed: - warnings.warn( - f"The mixed model did not converge (or could not be fit) for {n_failed} of {len(self.models)} variables. " - "See the `converged` column of the results.", - UserWarning, - stacklevel=2, - ) - self._fitted = True - - def _contrast_df(self, contrast: np.ndarray) -> int: - """Denominator degrees of freedom for a contrast, following the between-within method.""" - involved = ~np.isclose(contrast, 0) - df = min( - self.df_between if (involved & self.is_between).any() else np.inf, - self.df_within if (involved & ~self.is_between).any() else np.inf, - ) - if df < 1: - raise ValueError( - f"The contrast cannot be tested: it involves effects that are constant within `{self.random_effect}`, " - f"but there are not enough levels of `{self.random_effect}` left to estimate them " - f"(between-group degrees of freedom: {self.df_between})." - ) - return int(df) - - def _test_single_contrast(self, contrast, **kwargs) -> pd.DataFrame: - contrast = np.asarray(contrast, dtype=float) - df = self._contrast_df(contrast) - res = [] - for var, mod in zip(tqdm(self.adata.var_names), self.models, strict=False): - row = {"variable": var, "p_value": np.nan, "t_value": np.nan, "sd": np.nan, "log_fc": np.nan} - if mod is not None: - # MixedLMResults.t_test expects a 2D restriction matrix over the fixed effects. - t_test = mod.t_test(np.atleast_2d(contrast)) - effect, sd = float(np.ravel(t_test.effect)[0]), float(np.ravel(t_test.sd)[0]) - t_value = effect / sd if sd > 0 else np.nan - row |= { - "p_value": 2 * stats.t.sf(abs(t_value), df) if np.isfinite(t_value) else np.nan, - "t_value": t_value, - "sd": sd, - "log_fc": effect, - } - row["converged"] = mod is not None and bool(mod.converged) - res.append(row) - - res_df = pd.DataFrame(res).sort_values("p_value") - tested = res_df["p_value"].notna().to_numpy() - adj = np.full(len(res_df), np.nan) - if tested.any(): - adj[tested] = fdrcorrection(res_df["p_value"].to_numpy()[tested])[1] - return res_df.assign(adj_p_value=adj) diff --git a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py index 5e5f33e5..b823477a 100644 --- a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py +++ b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py @@ -3,8 +3,8 @@ 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 @@ -12,11 +12,13 @@ from statsmodels.tools.sm_exceptions import 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,83 @@ 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]] + 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. + """ + + _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 +154,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 +173,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 +207,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 +223,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 +232,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 e42a6bce..91e08e14 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_linear_mixed_model.py b/tests/tools/_differential_gene_expression/test_linear_mixed_model.py deleted file mode 100644 index c6e54ed7..00000000 --- a/tests/tools/_differential_gene_expression/test_linear_mixed_model.py +++ /dev/null @@ -1,166 +0,0 @@ -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 -from scipy import stats - -if find_spec("formulaic_contrasts") is None or find_spec("formulaic") is None: - pytestmark = pytest.mark.skip(reason="formulaic_contrasts and formulaic not available") - - -def _donor_level_adata(n_donors: int = 6, n_cells: int = 30, n_genes: int = 5, *, seed: int = 0) -> ad.AnnData: - """Log-scale expression with a random donor effect; the condition is assigned per donor (between-donor).""" - rng = np.random.default_rng(seed) - donor = np.repeat([f"D{i}" for i in range(n_donors)], n_cells) - condition = np.where(np.repeat(np.arange(n_donors), n_cells) < n_donors // 2, "A", "B") - donor_effect = rng.normal(0, 0.7, (n_donors, n_genes)) - X = rng.normal(size=(n_donors * n_cells, n_genes)) + np.repeat(donor_effect, n_cells, axis=0) - obs = pd.DataFrame({"donor": donor, "condition": condition}, index=[f"c{i}" for i in range(len(donor))]) - return ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=[f"gene{i}" for i in range(n_genes)])) - - -def test_linear_mixed_model(test_adata_minimal): - """Check that the method can be initialized and fitted, and perform basic checks on the result of test_contrasts.""" - from pertpy.tools._differential_gene_expression import LinearMixedModel - - method = LinearMixedModel(adata=test_adata_minimal, design="~condition + (1 | donor)") - assert not method.is_fitted - method.fit() - assert method.is_fitted - res_df = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) - assert len(res_df) == test_adata_minimal.n_vars - assert {"variable", "p_value", "t_value", "sd", "log_fc", "adj_p_value", "converged"} <= set(res_df.columns) - # gene1 has a large difference between conditions - assert res_df.set_index("variable").loc["gene1", "p_value"] < 0.05 - - -def test_matches_statsmodels(): - """Effect sizes come from statsmodels' MixedLM; p-values use the between-group degrees of freedom.""" - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata(n_genes=3) - method = LinearMixedModel(adata, design="~condition + (1 | donor)") - method.fit() - contrast = method.contrast(column="condition", baseline="A", group_to_compare="B") - res = method.test_contrasts(contrast).set_index("variable") - - exog = np.column_stack([np.ones(adata.n_obs), (adata.obs["condition"] == "B").to_numpy(float)]) - for i, gene in enumerate(adata.var_names): - expected = sm.MixedLM(adata.X[:, i], exog, groups=adata.obs["donor"].to_numpy()).fit(reml=True) - np.testing.assert_allclose(res.loc[gene, "log_fc"], expected.fe_params[1], rtol=1e-6) - np.testing.assert_allclose(res.loc[gene, "sd"], expected.bse_fe[1], rtol=1e-6) - t_value = expected.fe_params[1] / expected.bse_fe[1] - # 6 donors minus intercept and condition, both constant within donors - np.testing.assert_allclose(res.loc[gene, "p_value"], 2 * stats.t.sf(abs(t_value), df=4), rtol=1e-6) - - -def test_between_within_df(): - from pertpy.tools._differential_gene_expression._linear_mixed_model import _between_within_df - - groups = np.repeat(["D0", "D1", "D2", "D3"], 5) - between = np.repeat([0.0, 0.0, 1.0, 1.0], 5) # constant within each donor - within = np.tile([0.0, 1.0, 0.0, 1.0, 1.0], 4) # varies within donors - design = np.column_stack([np.ones(20), between, within]) - - is_between, df_between, df_within = _between_within_df(design, groups) - - np.testing.assert_array_equal(is_between, [True, True, False]) - assert df_between == 4 - 2 # donors minus intercept and between-donor effect - assert df_within == 20 - 4 - 1 # cells minus donor indicators and within-donor effect - - -def test_type_i_error_between_donor_effect(): - """With few donors, the default normal approximation would be anti-conservative; the t-test must not be.""" - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata(n_genes=200, seed=1) - method = LinearMixedModel(adata, design="~condition + (1 | donor)") - method.fit() - res = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) - assert (res["p_value"] < 0.05).mean() < 0.1 - - -def test_optimizer_fallback_converges(recwarn): - """Fits that fail statsmodels' gradient check are refit with derivative-free optimizers, so clean data converges.""" - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata(n_genes=200, seed=1) - method = LinearMixedModel(adata, design="~condition + (1 | donor)") - method.fit() - res = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) - assert res["converged"].all() - assert not [w for w in recwarn if "did not converge" in str(w.message)] - - -def test_compare_groups(): - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata() - res = LinearMixedModel.compare_groups( - adata, column="condition", baseline="A", groups_to_compare="B", paired_by="donor" - ) - assert len(res) == adata.n_vars - assert set(res["contrast"]) == {"B"} - - with pytest.raises(ValueError, match="requires `paired_by`"): - LinearMixedModel.compare_groups(adata, column="condition", baseline="A", groups_to_compare="B") - - -@pytest.mark.parametrize( - ("design", "error", "match"), - [ - ("~condition", ValueError, "exactly one random intercept"), - ("~condition + (1 | donor) + (1 | batch)", ValueError, "exactly one random intercept"), - ("~condition + (1 | missing)", ValueError, "not a column of adata.obs"), - (np.ones((180, 1)), TypeError, "requires a formula"), - ], -) -def test_invalid_design(design, error, match): - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata() - adata.obs["batch"] = "b1" - with pytest.raises(error, match=match): - LinearMixedModel(adata, design=design) - - -def test_missing_random_effect_values(): - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata() - adata.obs["donor"] = adata.obs["donor"].astype(object) - adata.obs.iloc[0, adata.obs.columns.get_loc("donor")] = None - with pytest.raises(ValueError, match="missing values"): - LinearMixedModel(adata, design="~condition + (1 | donor)") - - -def test_untestable_contrast(): - """One donor per condition leaves no degrees of freedom for a between-donor effect.""" - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata(n_donors=2) - method = LinearMixedModel(adata, design="~condition + (1 | donor)") - method.fit() - with pytest.raises(ValueError, match="cannot be tested"): - method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) - - -def test_constant_variable(): - """Variables that cannot be fit get NaN statistics and are excluded from the FDR correction.""" - from pertpy.tools._differential_gene_expression import LinearMixedModel - - adata = _donor_level_adata(n_genes=4) - adata.X[:, 0] = 0 - method = LinearMixedModel(adata, design="~condition + (1 | donor)") - with pytest.warns(UserWarning, match="did not converge"): - method.fit() - res = method.test_contrasts(method.contrast(column="condition", baseline="A", group_to_compare="B")) - res = res.set_index("variable") - - assert np.isnan(res.loc["gene0", "p_value"]) - assert np.isnan(res.loc["gene0", "adj_p_value"]) - assert not res.loc["gene0", "converged"] - assert res.drop(index="gene0")["adj_p_value"].notna().all() 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( From 6a32240996e48a6037e0328e0d69a99f52056904 Mon Sep 17 00:00:00 2001 From: Lukas Heumos Date: Wed, 21 Oct 2026 17:02:32 +0200 Subject: [PATCH 3/6] Drop the mixed-model sentence from the tools index --- docs/api/tools_index.md | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/api/tools_index.md b/docs/api/tools_index.md index 3d1efcab..8d956e2c 100644 --- a/docs/api/tools_index.md +++ b/docs/api/tools_index.md @@ -13,7 +13,6 @@ Differential gene expression involves the quantitative comparison of gene expression levels between two or more groups, such as different cell types, tissues, or conditions to discern genes that are significantly up- or downregulated in response to specific biological contexts or stimuli. Pertpy enables differential gene expression tests through a common interface that supports complex designs. -{class}`~pertpy.tools.Statsmodels` also fits linear mixed models with a random intercept, written as `(1 | donor)` in the design, to test single cells while accounting for the samples they come from. ```{eval-rst} .. autosummary:: From 2c6413ea9d4b17b94d76b69c7dc67c7ca3850c63 Mon Sep 17 00:00:00 2001 From: Lukas Heumos Date: Wed, 21 Oct 2026 17:15:34 +0200 Subject: [PATCH 4/6] Point the tutorials at the differential expression notebook with the random-intercept section --- docs/tutorials/notebooks | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From ff4df19bda594258ceef86e2b8756281d9471d7c Mon Sep 17 00:00:00 2001 From: Lukas Heumos Date: Thu, 22 Oct 2026 14:14:17 +0200 Subject: [PATCH 5/6] Note that the per-gene mixed model refits report z-test p-values statsmodels.MixedLM has no Satterthwaite degrees of freedom, so models[i] matches test_contrasts in estimates but reports smaller z-test p-values. Its ConvergenceWarnings are left visible because they are statsmodels' own diagnostics for the object the user asked for. --- src/pertpy/tools/_differential_gene_expression/_statsmodels.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py index b823477a..e7f05d5a 100644 --- a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py +++ b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py @@ -122,6 +122,7 @@ class Statsmodels(LinearModelBase): 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 From f92840fc19e7e5e2f7bece929afa85c2b74af12c Mon Sep 17 00:00:00 2001 From: Lukas Heumos Date: Thu, 22 Oct 2026 14:15:51 +0200 Subject: [PATCH 6/6] Silence ConvergenceWarnings of the per-gene mixed model refits MixedLM warns on most fits near the boundary, which made inspecting models[i] noisy without telling the user anything actionable about the Satterthwaite results. --- .../tools/_differential_gene_expression/_statsmodels.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py index e7f05d5a..e92a16e9 100644 --- a/src/pertpy/tools/_differential_gene_expression/_statsmodels.py +++ b/src/pertpy/tools/_differential_gene_expression/_statsmodels.py @@ -9,7 +9,7 @@ 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, _numba_threads, _parallelize_with_joblib @@ -99,7 +99,9 @@ def __len__(self) -> int: def __getitem__(self, index): if isinstance(index, slice): return [self[i] for i in range(len(self))[index]] - return sm.MixedLM(to_dense(self.endog[:, [index]])[:, 0], self.exog, groups=self.groups).fit() + 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."""