diff --git a/src/pertpy/tools/_milo.py b/src/pertpy/tools/_milo.py index 40aa251e..7647989d 100644 --- a/src/pertpy/tools/_milo.py +++ b/src/pertpy/tools/_milo.py @@ -34,20 +34,22 @@ from sklearn.metrics.pairwise import euclidean_distances -def _contrast_vector(columns: list[str], model_contrasts: str) -> np.ndarray: +def _contrast_vector(columns: list[str], model_contrasts: str, reference_levels: Collection[str] = ()) -> np.ndarray: """Turn an R style contrast such as ``conditionB-conditionA`` into weights over formulaic design columns. Formulaic names a coefficient ``condition[T.B]`` where R names it ``conditionB``, so the columns are matched on their R spelling. + A term naming a reference level in ``reference_levels`` has no coefficient under treatment coding and weighs zero. """ r_names = {column.replace("[T.", "").replace("[", "").replace("]", ""): column for column in columns} weights = pd.Series(0.0, index=columns) for sign, term in re.findall(r"([+-]?)\s*([^+-]+)", model_contrasts): name = term.strip() - if name not in r_names: + if name in r_names: + weights[r_names[name]] += -1.0 if sign == "-" else 1.0 + elif name not in reference_levels: raise ValueError( f"Contrast term {name!r} does not match any coefficient of the design. Available: {sorted(r_names)}." ) - weights[r_names[name]] += -1.0 if sign == "-" else 1.0 return weights.to_numpy() @@ -153,7 +155,8 @@ def make_nhoods( Otherwise: nhoods: :class:`scipy.sparse.csr_matrix` in `adata.obsm['nhoods']`. - A binary matrix of cell to neighbourhood assignments. Neighbourhoods in the columns are ordered by the order of the index cell in adata.obs_names + A binary matrix of cell to neighbourhood assignments, in which every index cell belongs to its own neighbourhood. + Neighbourhoods in the columns are ordered by the order of the index cell in adata.obs_names nhood_ixs_refined: pandas.Series in `adata.obs['nhood_ixs_refined']`. A boolean indicating whether a cell is an index for a neighbourhood @@ -223,6 +226,7 @@ def make_nhoods( refined_vertices = np.unique(refined_vertices) refined_vertices.sort() + knn_graph.setdiag(1) # type: ignore[union-attr] nhoods = knn_graph[:, refined_vertices] adata.obsm["nhoods"] = nhoods @@ -351,9 +355,9 @@ def da_nhoods( max_iter: Maximum number of iterations of a mixed model fit. tol: Convergence tolerance of a mixed model fit. solver: The solver to fit the model to, ignored for a mixed model. - The "edger" solver requires R, rpy2 and edgeR to be installed and is the closest to the R implementation. + The "edger" solver requires R, rpy2 and edgeR to be installed and reproduces the R implementation. The "pydeseq2" requires pydeseq2 to be installed. - It is still very comparable to the "edger" solver but might be a bit slower. + It uses the Wald test of DESeq2 instead of the quasi-likelihood F-test of edgeR, so its results are close to but not identical with those of R Milo. Returns: None, modifies `milo_mdata['milo']` in place, adding the results of the DA test to `.var`: @@ -439,6 +443,11 @@ def da_nhoods( if isinstance(design_df[column].dtype, pd.CategoricalDtype): design_df[column] = design_df[column].cat.remove_unused_categories() + fixed = fixed_design if add_intercept and model_contrasts is None else fixed_design + " + 0" + reference_levels = { + f"{column}{design_df[column].astype('category').cat.categories[0]}" + for column in design_df.select_dtypes(exclude="number").columns + } if random_effects: if find_spec("formulaic_contrasts") is None: raise ImportError( @@ -446,7 +455,6 @@ def da_nhoods( ) from formulaic_contrasts import FormulaicContrasts - fixed = fixed_design if add_intercept and model_contrasts is None else fixed_design + " + 0" design_matrix = FormulaicContrasts(design_df, fixed).design_matrix _check_residual_df(design_matrix, design) counts_filtered = count_mat[np.ix_(keep_nhoods, keep_smp)] @@ -457,7 +465,7 @@ def da_nhoods( np.asarray(design_matrix, dtype=float), random_effect_matrices(design_df, random_effects), np.log(lib_size_filtered), - contrast=_contrast_vector(list(design_matrix.columns), model_contrasts) + contrast=_contrast_vector(list(design_matrix.columns), model_contrasts, reference_levels) if model_contrasts is not None else None, reml=reml, @@ -485,12 +493,10 @@ def da_nhoods( from rpy2.robjects.vectors import FloatVector # Define model matrix - if not add_intercept or model_contrasts is not None: - design = design + " + 0" design_df = design_df.astype(dict.fromkeys(design_df.select_dtypes(exclude=["number"]).columns, "category")) with localconverter(ro.default_converter + pandas2ri.converter): design_r = pandas2ri.py2rpy(design_df) - formula_r = stats.formula(design) + formula_r = stats.formula(fixed) model = stats.model_matrix(object=formula_r, data=design_r) model_np = np.array(model) _check_residual_df(model_np, design) @@ -504,7 +510,7 @@ def da_nhoods( dge = edgeR.DGEList(counts=count_mat_r, lib_size=lib_size_r) dge = edgeR.calcNormFactors(dge, method="TMM") dge = edgeR.estimateDisp(dge, model) - fit = edgeR.glmQLFit(dge, model, robust=True) + fit = edgeR.glmQLFit(dge, model, robust=True, legacy=True) # Test n_coef = model_np.shape[1] if model_contrasts is not None: @@ -518,7 +524,7 @@ def da_nhoods( get_model_cols = STAP(r_str, "get_model_cols") with localconverter(ro.default_converter + numpy2ri.converter + pandas2ri.converter): - model_mat_cols = get_model_cols.get_model_cols(design_df, design) + model_mat_cols = get_model_cols.get_model_cols(design_df, fixed) with localconverter(ro.default_converter + pandas2ri.converter + numpy2ri.converter): model_df = pandas2ri.rpy2py(model) model_df = pd.DataFrame(model_df) @@ -568,54 +574,24 @@ def da_nhoods( dict.fromkeys(design_df_filtered.select_dtypes(exclude=["number"]).columns, "category") ) - design_clean = design if design.startswith("~") else f"~{design}" - dds = DeseqDataSet( counts=pd.DataFrame(counts_filtered.T, index=design_df_filtered.index), metadata=design_df_filtered, - design=design_clean, + design=fixed if fixed.startswith("~") else f"~{fixed}", refit_cooks=True, size_factors_fit_type="poscounts", ) - _check_residual_df(dds.obsm["design_matrix"], design) # type: ignore[arg-type] + design_matrix = cast_frame(dds.obsm["design_matrix"]) + _check_residual_df(design_matrix, design) dds.deseq2() - if model_contrasts is not None and "-" in model_contrasts: - if "(" in model_contrasts or "+" in model_contrasts.split("-")[1]: - raise ValueError( - f"Complex contrasts like '{model_contrasts}' are not supported by pydeseq2. " - "Use simple pairwise contrasts (e.g., 'GroupA-GroupB') or switch to solver='edger'." - ) - - parts = model_contrasts.split("-") - factor_name = design_clean.replace("~", "").split("+")[-1].strip() - group1 = parts[0].replace(factor_name, "").strip() - group2 = parts[1].replace(factor_name, "").strip() - if factor_name not in design_df_filtered.columns: - raise ValueError( - f"Contrast factor {factor_name!r} is not a column of the design dataframe. " - f"Available columns: {list(design_df_filtered.columns)}." - ) - if not isinstance(design_df_filtered[factor_name].dtype, pd.CategoricalDtype): - design_df_filtered[factor_name] = design_df_filtered[factor_name].astype("category") - available_levels = list(design_df_filtered[factor_name].cat.categories) - missing = [g for g in (group1, group2) if g not in available_levels] - if missing: - raise ValueError( - f"Contrast levels {missing!r} not found in factor {factor_name!r}. " - f"Available levels: {available_levels}. " - f"Contrasts must follow the form '{factor_name}-{factor_name}' " - "with both levels present in the data." - ) - stat_res = DeseqStats(dds, contrast=[factor_name, group1, group2]) - else: - factor_name = design_clean.replace("~", "").split("+")[-1].strip() - if not isinstance(design_df_filtered[factor_name], pd.CategoricalDtype): - design_df_filtered[factor_name] = design_df_filtered[factor_name].astype("category") - categories = design_df_filtered[factor_name].cat.categories - stat_res = DeseqStats(dds, contrast=[factor_name, categories[-1], categories[0]]) - + contrast = ( + _contrast_vector(list(design_matrix.columns), model_contrasts, reference_levels) + if model_contrasts is not None + else np.eye(design_matrix.shape[1])[-1] + ) + stat_res = DeseqStats(dds, contrast=contrast) stat_res.summary() res = stat_res.results_df diff --git a/tests/tools/test_milo.py b/tests/tools/test_milo.py index 1742eaca..c5c04d25 100644 --- a/tests/tools/test_milo.py +++ b/tests/tools/test_milo.py @@ -70,6 +70,13 @@ def test_make_nhoods_sizes(adata, milo): assert knn_graph.sum(0).min() <= adata.obsm["nhoods"].sum(0).min() +def test_make_nhoods_contains_index_cells(adata, milo): + adata = adata.copy() + milo.make_nhoods(adata) + index_cells = np.flatnonzero(adata.obs["nhood_ixs_refined"] == 1) + assert np.all(np.asarray(adata.obsm["nhoods"][index_cells, np.arange(len(index_cells))]) == 1) + + def test_make_nhoods_neighbors_key(adata, milo): adata = adata.copy() k = adata.uns["neighbors"]["params"]["n_neighbors"] @@ -231,6 +238,63 @@ def test_da_nhoods_default_contrast(da_nhoods_mdata, milo, solver): assert np.corrcoef(contr_results["logFC"], default_results["logFC"])[0, 1] > 0.99 +@pytest.fixture +def three_condition_mdata(adata, milo): + adata = adata.copy() + milo.make_nhoods(adata) + rng = np.random.default_rng(seed=42) + conditions = ["ConditionA", "ConditionB", "ConditionC"] + adata.obs["condition"] = rng.choice(conditions, size=adata.n_obs) + da_cells = adata.obs["louvain"] == "1" + adata.obs.loc[da_cells, "condition"] = rng.choice(conditions, size=da_cells.sum(), p=[0.1, 0.8, 0.1]) + adata.obs["replicate"] = rng.choice(["R1", "R2", "R3"], size=adata.n_obs) + adata.obs["sample"] = adata.obs["replicate"] + adata.obs["condition"] + return milo.count_nhoods(adata, sample_col="sample") + + +def test_da_nhoods_contrast_of_single_coefficient(three_condition_mdata, milo, solver): + """A contrast naming a single coefficient tests that coefficient instead of the last level of the design.""" + mdata = three_condition_mdata + index_cells = mdata["milo"].var["index_cell"] + enriched = (mdata["rna"].obs.loc[index_cells, "louvain"] == "1").to_numpy() + + milo.da_nhoods(mdata, design="~replicate+condition", model_contrasts="conditionConditionB", solver=solver) + b_vs_a = mdata["milo"].var["logFC"].to_numpy() + milo.da_nhoods(mdata, design="~replicate+condition", model_contrasts="conditionConditionC", solver=solver) + c_vs_a = mdata["milo"].var["logFC"].to_numpy() + + assert np.nanmean(b_vs_a[enriched]) > 1 + assert np.nanmean(b_vs_a[enriched]) > np.nanmean(c_vs_a[enriched]) + 1 + + +def test_da_nhoods_contrast_against_reference_level(three_condition_mdata, milo): + """The reference level has no coefficient, so subtracting it leaves the contrast unchanged.""" + mdata = three_condition_mdata + milo.da_nhoods(mdata, design="~replicate+condition", model_contrasts="conditionConditionB", solver="pydeseq2") + coefficient = mdata["milo"].var["logFC"].to_numpy() + milo.da_nhoods( + mdata, + design="~replicate+condition", + model_contrasts="conditionConditionB-conditionConditionA", + solver="pydeseq2", + ) + np.testing.assert_allclose(mdata["milo"].var["logFC"].to_numpy(), coefficient) + + +def test_da_nhoods_continuous_covariate_per_unit(da_nhoods_mdata, milo, solver): + """The log fold change of a continuous covariate is per unit, so rescaling the covariate rescales it.""" + mdata = da_nhoods_mdata.copy() + obs = mdata["rna"].obs + obs["dose"] = (obs["condition"] == "ConditionB") + obs["replicate"].str[1].astype(float) / 10 + + milo.da_nhoods(mdata, design="~dose", solver=solver) + per_unit = mdata["milo"].var["logFC"].to_numpy() + obs["dose"] *= 10 + milo.da_nhoods(mdata, design="~dose", solver=solver) + + np.testing.assert_allclose(mdata["milo"].var["logFC"].to_numpy() * 10, per_unit, rtol=1e-3, atol=1e-3) + + @pytest.mark.skipif(find_spec("formulaic_contrasts") is None, reason="formulaic-contrasts not available") def test_da_nhoods_glmm(da_nhoods_mdata, milo): mdata = da_nhoods_mdata.copy()