diff --git a/.github/workflows/pr_validation.yml b/.github/workflows/pr_validation.yml index 178fa458..d6bae5b1 100644 --- a/.github/workflows/pr_validation.yml +++ b/.github/workflows/pr_validation.yml @@ -46,8 +46,23 @@ jobs: run: | uv run ruff format --check pydeseq2 - name: Type check with mypy + if: matrix.python == '3.11' run: | uv run mypy -p pydeseq2 + - name: Test SciPy 1.13 sparse compatibility + if: matrix.python == '3.11' + run: | + uv run --no-project --isolated --python 3.11 \ + --with-editable . \ + --with 'numpy==2.0.0' \ + --with 'scipy==1.13.0' \ + --with 'anndata==0.11.0' \ + --with 'pandas==2.2.2' \ + --with 'zarr<3' \ + --with 'pytest==8.4.0' \ + python -m pytest -q \ + tests/test_transcript_length_normalization.py \ + tests/test_utils.py - name: Test with pytest run: | uv run coverage run -m pytest diff --git a/README.md b/README.md index 4b539440..779d8344 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ In case there is a feature you would particularly like to be implemented, feel f - [Installation](#installation) - [Requirements](#requirements) - [Getting started](#getting-started) + - [Transcript-length normalization](#transcript-length-normalization) - [Documentation](#documentation) - [Data](#data) - [Contributing](#contributing) @@ -93,6 +94,63 @@ Please don't hesitate to open an issue in case you encounter any issue due to po The [Getting Started](https://pydeseq2.readthedocs.io/en/latest/auto_examples/index.html) section of the documentation contains downloadable examples on how to use PyDESeq2. +### Transcript-length normalization + +For unscaled estimated gene counts imported from transcript-level quantifiers +(`countsFromAbundance="no"` in tximport), pass the matching average transcript lengths +alongside the counts. Both matrices must use samples as rows and genes as columns. When +using data frames, their labels and ordering must be identical; arrays must already use +the same ordering: + +```python +dds = DeseqDataSet( + counts=estimated_counts, + metadata=metadata, + transcript_lengths=average_transcript_lengths, + design="~condition", +) +dds.deseq2() +``` + +[`pytximport`](https://github.com/complextissue/pytximport), described in [4], can +produce a compatible AnnData object directly: + +```python +from pydeseq2.dds import DeseqDataSet +from pytximport import tximport + +sample_names = ["sample_1", "sample_2"] +txi = tximport( + [f"{sample}/quant.sf" for sample in sample_names], + data_type="salmon", + transcript_gene_map=transcript_gene_map, + counts_from_abundance=None, + output_type="anndata", +) +txi.obs = metadata.loc[sample_names].copy() +dds = DeseqDataSet(adata=txi, design="~condition") +dds.deseq2() +``` + +Use only unscaled estimated counts (`counts_from_abundance=None`); the file list and +metadata rows must use the same sample order. For backed AnnData, call +`adata.to_memory()` first and ensure sufficient memory. Length-source precedence is +deterministic: explicit `transcript_lengths`, `adata.layers["avg_tx_length"]`, then +compatible pytximport fields; PyDESeq2 copies `adata.obsm["length"]` into the canonical +layer. Because that matrix has no gene labels and AnnData does not align its second +`obsm` dimension with `adata.var`, pass it before gene-axis subsetting/reordering or +apply the same selection and order to `adata.obsm["length"]`. + +Following the +[`tximport`](https://bioconductor.org/packages/release/bioc/html/tximport.html) and +[`DESeq2`](https://bioconductor.org/packages/release/bioc/html/DESeq2.html) +workflow [1, 3], PyDESeq2 rounds estimated counts and combines transcript-length +offsets with median-of-ratios library-size normalization. Average lengths and the +resulting sample-by-gene factors are available in `dds.layers["avg_tx_length"]` +and `dds.layers["normalization_factors"]`. The `ratio` and `poscounts` +size-factor methods are supported; `iterative` is not currently compatible with +transcript-length offsets. + ### Documentation @@ -159,6 +217,16 @@ In Dec 2025, the maintenance of PyDESeq2 was taken over by the scverse community Bioinformatics, 35(12), 2084-2092. +[3] Soneson, C., Love, M. I., & Robinson, M. D. (2015). "Differential analyses + for RNA-seq: transcript-level estimates improve gene-level inferences." + F1000Research, 4:1521. + + +[4] Kuehl, M., Wong, M. N., Wanner, N., Bonn, S., & Puelles, V. G. (2024). + "Gene count estimation with pytximport enables reproducible analysis of + bulk RNA sequencing data in Python." Bioinformatics, 40(12), btae700. + + ## License PyDESeq2 is released under an [MIT license](https://github.com/owkin/PyDESeq2/blob/main/LICENSE). diff --git a/docs/source/conf.py b/docs/source/conf.py index 94c13f37..1c5b4309 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -56,6 +56,7 @@ "python": ("https://docs.python.org/3", None), "numpy": ("https://numpy.org/doc/stable/", None), "pandas": ("https://pandas.pydata.org/docs/", None), + "scipy": ("https://docs.scipy.org/doc/scipy/", None), "anndata": ("https://anndata.readthedocs.io/en/latest/", None), } diff --git a/pydeseq2/dds.py b/pydeseq2/dds.py index f0298057..a933c17d 100644 --- a/pydeseq2/dds.py +++ b/pydeseq2/dds.py @@ -1,6 +1,7 @@ import sys import time import warnings +from typing import Any from typing import Literal from typing import cast @@ -9,6 +10,7 @@ import pandas as pd from formulaic_contrasts import FormulaicContrasts # type: ignore[import-untyped] from scipy.optimize import minimize +from scipy.sparse import issparse # type: ignore from scipy.special import polygamma # type: ignore from scipy.stats import f # type: ignore from scipy.stats import trim_mean # type: ignore @@ -30,6 +32,22 @@ warnings.simplefilter("ignore", FutureWarning) +def _rounded_counts(counts: Any) -> Any: + """Round dense or sparse estimated counts without mutating the input.""" + if isinstance(counts, pd.DataFrame): + return counts.round() + if isinstance(counts, np.ndarray): + return np.rint(counts) + if counts.format in {"csr", "csc"}: + rounded = counts.copy() + else: + rounded = counts.tocsr() + rounded.sum_duplicates() + rounded.data = np.rint(rounded.data) + rounded.eliminate_zeros() + return rounded + + class DeseqDataSet(ad.AnnData): r"""A class to implement dispersion and log fold-change (LFC) estimation. @@ -45,7 +63,11 @@ class DeseqDataSet(ad.AnnData): adata : anndata.AnnData AnnData from which to initialize the DeseqDataSet. Must have counts ('X') and sample metadata ('obs') fields. If ``None``, both ``counts`` and ``metadata`` - arguments must be provided. + arguments must be provided. Compatible pytximport objects with unscaled + estimated counts are detected from ``obsm["length"]`` and + ``uns["counts_from_abundance"] is None``. + Backed AnnData objects are not supported; call ``adata.to_memory()`` before + initialization. counts : pandas.DataFrame Raw counts. One column per gene, rows are indexed by sample barcodes. @@ -54,6 +76,18 @@ class DeseqDataSet(ad.AnnData): DataFrame containing sample metadata. Must be indexed by sample barcodes. + transcript_lengths : pandas.DataFrame or numpy.ndarray, optional + Average transcript lengths for each sample and gene, typically imported from + transcript-level quantification with unscaled estimated counts (tximport's + ``countsFromAbundance="no"`` mode). Scaled TPM-derived counts must not be paired + with these offsets. Must have the same sample-by-gene shape and ordering as + ``counts``. When provided, PyDESeq2 + constructs gene- and sample-specific normalization factors that correct for + transcript-length changes as well as library size. Estimated counts are rounded + to the nearest integer, matching DESeq2. Explicit lengths take precedence over + ``adata.layers["avg_tx_length"]``, which takes precedence over compatible + pytximport fields. (default: ``None``). + design : str or pandas.DataFrame Model design. Can be either a pandas DataFrame representing a design matrix, or a formulaic formula in the format ``'x + z'`` or ``'~x+z'``. @@ -161,6 +195,8 @@ class DeseqDataSet(ad.AnnData): layers Key-indexed multi-dimensional arrays aligned to dimensions of `X`, e.g. "cooks". + Average transcript lengths and the resulting factors are stored as + ``"avg_tx_length"`` and ``"normalization_factors"``, respectively. n_processes : int Number of cpus to use for multiprocessing. @@ -209,6 +245,7 @@ def __init__( adata: ad.AnnData | None = None, counts: pd.DataFrame | None = None, metadata: pd.DataFrame | None = None, + transcript_lengths: pd.DataFrame | np.ndarray | None = None, design: str | pd.DataFrame = "~condition", design_factors: str | list[str] | None = None, continuous_factors: list[str] | None = None, @@ -228,7 +265,48 @@ def __init__( low_memory: bool = False, ) -> None: # Initialize the AnnData part + has_stored_transcript_lengths = ( + adata is not None and "avg_tx_length" in adata.layers + ) + selected_transcript_lengths: Any = transcript_lengths + has_transcript_lengths = ( + selected_transcript_lengths is not None or has_stored_transcript_lengths + ) + preserve_inherited_fit = False if adata is not None: + has_pytximport_fields = ( + "length" in adata.obsm and "counts_from_abundance" in adata.uns + ) + has_transcript_lengths = has_transcript_lengths or has_pytximport_fields + counts_from_abundance = adata.uns.get("counts_from_abundance") + if counts_from_abundance is not None and has_transcript_lengths: + raise ValueError( + "pytximport transcript-length offsets require unscaled " + "estimated counts with " + "adata.uns['counts_from_abundance'] set to None; got " + f"{counts_from_abundance!r}. Abundance-scaled counts must " + "not be combined with transcript-length offsets." + ) + if ( + has_pytximport_fields + and selected_transcript_lengths is None + and not has_stored_transcript_lengths + ): + selected_transcript_lengths = adata.obsm["length"] + + preserve_inherited_fit = ( + selected_transcript_lengths is None + and has_stored_transcript_lengths + and "normalization_factors" in adata.layers + ) + + if adata.isbacked: + raise ValueError( + "DeseqDataSet requires an in-memory AnnData object. Call " + "adata.to_memory() before initialization." + ) + expected_length_index = adata.obs_names + expected_length_columns = adata.var_names if counts is not None: warnings.warn( "adata was provided; ignoring counts.", UserWarning, stacklevel=2 @@ -237,21 +315,97 @@ def __init__( warnings.warn( "adata was provided; ignoring metadata.", UserWarning, stacklevel=2 ) + prepared_counts: Any = adata.X + if has_transcript_lengths: + prepared_counts = _rounded_counts(prepared_counts) + elif issparse(prepared_counts): + sparse_counts = cast(Any, prepared_counts) + if sparse_counts.format not in {"csr", "csc"}: + prepared_counts = sparse_counts.tocsr() + prepared_counts.sum_duplicates() + elif not sparse_counts.has_canonical_format: + prepared_counts = sparse_counts.copy() + prepared_counts.sum_duplicates() # Test counts before going further - test_valid_counts(adata.X) - # Copy fields from original AnnData - self.__dict__.update(adata.__dict__) - # Cast counts to ints to avoid any issue - self.X = adata.X.astype(int) + test_valid_counts(prepared_counts) + integer_counts = prepared_counts.astype(int) + if has_transcript_lengths: + # Own the containers PyDESeq2 writes while retaining references to + # existing large aligned matrices. + # AnnData migrates legacy neighbor matrices out of uns in place. + owned_uns = dict(adata.uns) + if "neighbors" in owned_uns: + owned_uns["neighbors"] = dict(owned_uns["neighbors"]) + super().__init__( + X=integer_counts, + obs=cast(pd.DataFrame, adata.obs).copy(), + var=cast(pd.DataFrame, adata.var).copy(), + uns=None, + obsm=cast(Any, dict(adata.obsm)), + varm=cast(Any, dict(adata.varm)), + obsp=cast(Any, dict(adata.obsp)), + varp=cast(Any, dict(adata.varp)), + # AnnData 0.13 can expose X as a None-keyed layer item. + layers={ + key: value + for key, value in adata.layers.items() + if key is not None + }, + raw=cast(Any, adata.raw), + ) + self.uns.update(owned_uns) + else: + self.__dict__.update(adata.__dict__) + # AnnData 0.13 stores X under the None layer key. Detach the + # container so assigning self.X does not modify the input AnnData. + if None in self.__dict__.get("_layers", {}): + self.__dict__["_layers"] = self.__dict__["_layers"].copy() + self.X = integer_counts elif counts is not None and metadata is not None: + expected_length_index = counts.index + expected_length_columns = counts.columns + prepared_counts = ( + _rounded_counts(counts) if has_transcript_lengths else counts + ) # Test counts before going further - test_valid_counts(counts) - super().__init__(X=counts.astype(int), obs=metadata) + test_valid_counts(prepared_counts) + super().__init__(X=prepared_counts.astype(int), obs=metadata) else: raise ValueError( "Either adata or both counts and metadata arguments must be provided." ) + if selected_transcript_lengths is not None: + if isinstance(selected_transcript_lengths, pd.DataFrame): + if not selected_transcript_lengths.index.equals(expected_length_index): + raise ValueError( + "transcript_lengths must have the same sample index as counts." + ) + if not selected_transcript_lengths.columns.equals( + expected_length_columns + ): + raise ValueError( + "transcript_lengths must have the same gene columns as counts." + ) + transcript_lengths_array = selected_transcript_lengths.to_numpy( + dtype=float, copy=True + ) + else: + transcript_lengths_array = np.array( + selected_transcript_lengths, dtype=float, copy=True + ) + self._validate_transcript_lengths(transcript_lengths_array) + self.layers["avg_tx_length"] = transcript_lengths_array + elif "avg_tx_length" in self.layers: + stored_lengths = self.layers["avg_tx_length"] + transcript_lengths_array = ( + cast(Any, stored_lengths).toarray() + if hasattr(stored_lengths, "toarray") + else np.asarray(stored_lengths, dtype=float) + ) + self._validate_transcript_lengths(transcript_lengths_array) + self.layers["avg_tx_length"] = transcript_lengths_array + self.fit_type = fit_type self.design = design @@ -307,6 +461,87 @@ def __init__( # Check that the design matrix has full rank self._check_full_rank_design() + if preserve_inherited_fit and adata is not None: + source_design = adata.obsm.get("design_matrix") + normalization_control_mask = self._make_control_mask(control_genes) + fit_settings = ( + ("fit_type", fit_type), + ("min_mu", min_mu), + ("min_disp", min_disp), + ("max_disp", np.maximum(max_disp, self.n_obs)), + ("refit_cooks", refit_cooks), + ("min_replicates", min_replicates), + ("beta_tol", beta_tol), + ) + preserve_inherited_fit = ( + isinstance(source_design, pd.DataFrame) + and cast(pd.DataFrame, self.obsm["design_matrix"]).equals(source_design) + and all( + hasattr(adata, attr) and np.array_equal(getattr(adata, attr), value) + for attr, value in fit_settings + ) + and getattr(adata, "_normalization_fit_type", None) + == size_factors_fit_type + and hasattr(adata, "_normalization_control_mask") + and np.array_equal( + adata._normalization_control_mask, + normalization_control_mask, + ) + and hasattr(adata, "inference") + and ( + type(adata.inference) is DefaultInference + if inference is None + else inference is adata.inference + ) + ) + + if adata is not None and has_transcript_lengths and not preserve_inherited_fit: + # Invalidate fitting state derived from copied normalization factors. + for column in ("size_factors", "replaceable"): + if column in self.obs: + del self.obs[column] + for column in ( + "_normed_means", + "non_zero", + "_MoM_dispersions", + "genewise_dispersions", + "vst_genewise_dispersions", + "_genewise_converged", + "fitted_dispersions", + "MAP_dispersions", + "_MAP_converged", + "dispersions", + "_outlier_genes", + "_LFC_converged", + "replaced", + "refitted", + "_pvalue_cooks_outlier", + ): + if column in self.var: + del self.var[column] + for layer in ( + "normalization_factors", + "normed_counts", + "_mu_hat", + "_vst_mu_hat", + "vst_counts", + "cooks", + "replace_cooks", + ): + self.layers.pop(layer, None) + for key in ("_mu_LFC", "_hat_diagonals"): + self.obsm.pop(key, None) + self.varm.pop("LFC", None) + for key in ( + "trend_coeffs", + "vst_trend_coeffs", + "mean_disp", + "disp_function_type", + "_squared_logres", + "prior_disp_var", + ): + self.uns.pop(key, None) + self.min_mu = min_mu self.min_disp = min_disp self.max_disp = np.maximum(max_disp, self.n_obs) @@ -335,6 +570,108 @@ def __init__( # Initialize the inference object. self.inference = inference or DefaultInference(n_cpus=n_cpus) + if preserve_inherited_fit and adata is not None: + if "LFC" in self.varm: + self.varm["LFC"] = self.varm["LFC"].copy() + + self._normalization_fit_type = adata._normalization_fit_type + self._normalization_control_mask = adata._normalization_control_mask.copy() + + for attr in ( + "non_zero_idx", + "non_zero_genes", + "counts_to_refit", + "new_all_zeroes_genes", + ): + if hasattr(adata, attr): + setattr(self, attr, getattr(adata, attr).copy()) + + if hasattr(adata, "vst_fit_type"): + self.vst_fit_type = adata.vst_fit_type + + def _validate_transcript_lengths(self, transcript_lengths: np.ndarray) -> None: + """Validate a sample-by-gene average transcript-length matrix.""" + if transcript_lengths.shape != self.shape: + raise ValueError( + "transcript_lengths must have the same shape as counts " + f"({self.shape}), got {transcript_lengths.shape}." + ) + if not np.isfinite(transcript_lengths).all(): + raise ValueError("transcript_lengths must only contain finite values.") + if (transcript_lengths <= 0).any(): + raise ValueError("transcript_lengths must contain only positive values.") + + def _make_control_mask(self, control_genes: Any) -> np.ndarray: + """Return a gene mask for the selected normalization controls.""" + if control_genes is None: + return np.ones(self.n_vars, dtype=bool) + control_mask = np.zeros(self.n_vars, dtype=bool) + control_mask[self._normalize_indices((slice(None), control_genes))[1]] = True + return control_mask + + def _get_normalization_factors(self, gene_idx: Any = None) -> np.ndarray: + """Return normalization factors, optionally restricted to genes.""" + if "normalization_factors" in self.layers: + factors = np.asarray(self.layers["normalization_factors"]) + return factors if gene_idx is None else factors[:, gene_idx] + return self.obs["size_factors"].to_numpy() + + def _fit_transcript_length_factors( + self, + fit_type: Literal["ratio", "poscounts"], + control_mask: np.ndarray, + ) -> None: + """Fit DESeq2-style normalization factors from average transcript lengths.""" + counts = cast(Any, self.X).toarray() if issparse(self.X) else np.asarray(self.X) + transcript_lengths = np.asarray(self.layers["avg_tx_length"]) + + # DESeq2 centers each gene's average transcript lengths around a geometric + # mean of one before estimating library-size factors on adjusted counts. + length_factors = transcript_lengths / np.exp( + np.mean(np.log(transcript_lengths), axis=0) + ) + adjusted_counts = counts / length_factors + + if fit_type == "ratio": + self.logmeans, self.filtered_genes = deseq2_norm_fit(adjusted_counts) + else: + log_counts = np.zeros_like(counts, dtype=float) + np.log(counts, out=log_counts, where=counts != 0) + self.logmeans = log_counts.mean(axis=0) + self.filtered_genes = counts.sum(axis=0) > 0 + + usable_genes = control_mask & self.filtered_genes + if not usable_genes.any(): + raise ValueError( + "No genes are available to estimate size factors after applying " + "transcript-length and control-gene filters." + ) + + log_size_factors = np.empty(self.n_obs) + for sample_idx in range(self.n_obs): + sample_genes = usable_genes & (adjusted_counts[sample_idx] > 0) + if not sample_genes.any(): + raise ValueError( + "At least one sample has no positive counts among the genes " + "available for transcript-length normalization." + ) + log_size_factors[sample_idx] = np.median( + np.log(adjusted_counts[sample_idx, sample_genes]) + - self.logmeans[sample_genes] + ) + + # estimateNormFactors() in DESeq2 returns a matrix whose gene-wise geometric + # means are one. Retain the library-size component separately for backwards + # compatibility, while using the full matrix throughout model fitting. + size_factors = np.exp(log_size_factors) + size_factors /= np.exp(np.mean(np.log(size_factors))) + normalization_factors = length_factors * size_factors[:, None] + normalization_factors /= np.exp(np.mean(np.log(normalization_factors), axis=0)) + + self.obs["size_factors"] = size_factors + self.layers["normalization_factors"] = normalization_factors + self.layers["normed_counts"] = counts / normalization_factors + @property def variables(self): """Get the names of the variables used in the model definition.""" @@ -401,7 +738,7 @@ def vst_fit( """ # Start by fitting median-of-ratio size factors if not already present, # or if they were computed iteratively - if "size_factors" not in self.obsm or self.logmeans is None: + if "size_factors" not in self.obs or self.logmeans is None: self.fit_size_factors( fit_type=self.size_factors_fit_type ) # by default, fit_type != "iterative" @@ -461,6 +798,12 @@ def vst_transform(self, counts: np.ndarray | None = None) -> np.ndarray: "The vst_fit method should be called prior to vst_transform." ) + if counts is not None and "avg_tx_length" in self.layers: + raise ValueError( + "Transforming external counts with transcript-length normalization " + "requires matching transcript lengths, which are not yet supported." + ) + if counts is None: # the transformed counts will be the current ones normed_counts = self.layers["normed_counts"] @@ -649,28 +992,34 @@ def fit_size_factors( " DeseqDataSet initialization" ) - # If control genes are provided, set a mask where those genes are True - # This will override self.control_genes - if control_genes is not None: - _control_mask = np.zeros(self.X.shape[1], dtype=bool) + _control_mask = self._make_control_mask(control_genes) + normalization_control_mask = _control_mask.copy() - # Use AnnData internal indexing to get gene index array - # Allows bool/int/var_name to be provided - _control_mask[self._normalize_indices((slice(None), control_genes))[1]] = ( - True - ) + if "avg_tx_length" not in self.layers and "normalization_factors" in self.layers: + del self.layers["normalization_factors"] - # Otherwise mask all genes to be True - else: - _control_mask = np.ones(self.X.shape[1], dtype=bool) + if "avg_tx_length" in self.layers: + if fit_type == "iterative": + raise ValueError( + "The iterative size-factor method does not support " + "transcript-length normalization. Use 'ratio' or 'poscounts'." + ) + self._fit_transcript_length_factors( + fit_type=fit_type, + control_mask=_control_mask, + ) - if fit_type == "iterative": + elif fit_type == "iterative": self._fit_iterate_size_factors() elif fit_type == "poscounts": + counts = ( + cast(Any, self.X).toarray() if issparse(self.X) else np.asarray(self.X) + ) + # Calculate logcounts for x > 0 and take the mean for each gene - log_counts = np.zeros_like(self.X, dtype=float) - np.log(self.X, out=log_counts, where=self.X != 0) + log_counts = np.zeros_like(counts, dtype=float) + np.log(counts, out=log_counts, where=counts != 0) logmeans = log_counts.mean(axis=0) # Determine which genes are usable (finite logmeans) @@ -682,18 +1031,22 @@ def sizeFactor(x): _mask = np.logical_and(_control_mask, x > 0) return np.exp(np.median(np.log(x[_mask]) - logmeans[_mask])) - sf = np.apply_along_axis(sizeFactor, 1, self.X) + sf = np.apply_along_axis(sizeFactor, 1, counts) del log_counts # Normalize size factors to a geometric mean of 1 to match DESeq self.obs["size_factors"] = sf / (np.exp(np.mean(np.log(sf)))) self.layers["normed_counts"] = ( - self.X / self.obs["size_factors"].values[:, None] + counts / self.obs["size_factors"].values[:, None] ) self.logmeans = logmeans # Test whether it is possible to use median-of-ratios. - elif (self.X == 0).any(0).all(): + elif ( + (np.asarray((cast(Any, self.X) != 0).sum(axis=0)).ravel() < self.n_obs).all() + if issparse(self.X) + else (self.X == 0).any(0).all() + ): # There is at least a zero for each gene warnings.warn( "Every gene contains at least one zero, " @@ -704,22 +1057,25 @@ def sizeFactor(x): self._fit_iterate_size_factors() elif self.X is not None: - self.logmeans, self.filtered_genes = deseq2_norm_fit( - self.X.toarray() if not isinstance(self.X, np.ndarray) else self.X + counts = ( + cast(Any, self.X).toarray() if issparse(self.X) else np.asarray(self.X) ) + self.logmeans, self.filtered_genes = deseq2_norm_fit(counts) _control_mask &= self.filtered_genes ( self.layers["normed_counts"], self.obs["size_factors"], ) = deseq2_norm_transform( - self.X, cast(np.ndarray, self.logmeans), _control_mask + counts, cast(np.ndarray, self.logmeans), _control_mask ) else: raise ValueError("Counts matrix 'X' is None, cannot fit size factors.") end = time.time() self.var["_normed_means"] = self.layers["normed_counts"].mean(axis=0) + self._normalization_fit_type = fit_type + self._normalization_control_mask = normalization_control_mask if not self.quiet: print(f"... done in {end - start:.2f} seconds.\n", file=sys.stderr) @@ -740,7 +1096,7 @@ def fit_genewise_dispersions(self, vst=False) -> None: self.fit_size_factors(fit_type=self.size_factors_fit_type) # Exclude genes with all zeroes - self.var["non_zero"] = ~(self.X == 0).all(axis=0) + self.var["non_zero"] = np.asarray((self.X != 0).sum(axis=0)).ravel() > 0 self.non_zero_idx = np.arange(self.n_vars)[self.var["non_zero"]] self.non_zero_genes = self.var_names[self.var["non_zero"]] @@ -752,7 +1108,7 @@ def fit_genewise_dispersions(self, vst=False) -> None: # Convert design_matrix to numpy for speed design_matrix = self.obsm["design_matrix"].values - size_factors = self.obs["size_factors"].values + size_factors = self._get_normalization_factors(self.non_zero_idx) # mu_hat is initialized differently depending on the number of different factor # groups. If there are as many different factor combinations as design factors @@ -966,9 +1322,10 @@ def fit_LFC(self) -> None: if not self.quiet: print("Fitting LFCs...", file=sys.stderr) start = time.time() + normalization_factors = self._get_normalization_factors(self.non_zero_idx) mle_lfcs_, mu_, hat_diagonals_, converged_ = self.inference.irls( counts=self.X[:, self.non_zero_idx], - size_factors=self.obs["size_factors"].values, + size_factors=normalization_factors, design_matrix=design_matrix, disp=self.var.loc[self.var["non_zero"], "dispersions"].values, min_mu=self.min_mu, @@ -1021,7 +1378,13 @@ def calculate_cooks(self) -> None: ) # Calculate the squared pearson residuals for non-zero features - squared_pearson_res = self.X[:, self.var["non_zero"]] - self.obsm["_mu_LFC"] + counts = self.X[:, self.non_zero_idx] + if hasattr(counts, "tocoo"): + counts = counts.tocoo() + squared_pearson_res = -self.obsm["_mu_LFC"] + np.add.at(squared_pearson_res, (counts.row, counts.col), counts.data) + else: + squared_pearson_res = counts - self.obsm["_mu_LFC"] squared_pearson_res **= 2 # Calculate the overdispersion parameter tau @@ -1110,11 +1473,14 @@ def cooks_outlier(self): axis=0 ) - pos = self.layers["cooks"][:, cooks_outlier].argmax(0) - - cooks_outlier[cooks_outlier] = ( - self.X[:, cooks_outlier] > self.X[:, cooks_outlier][pos, np.arange(len(pos))] - ).sum(axis=0) < 3 + pos = np.asarray(self.layers["cooks"][:, cooks_outlier].argmax(0)).ravel() + for gene_idx, sample_idx in zip(np.flatnonzero(cooks_outlier), pos, strict=True): + gene_counts = ( + cast(Any, self.X)[:, [gene_idx]].toarray().ravel() + if issparse(self.X) + else np.asarray(self.X[:, gene_idx]).ravel() + ) + cooks_outlier[gene_idx] = (gene_counts > gene_counts[sample_idx]).sum() < 3 if self.low_memory: del self.layers["cooks"] @@ -1167,8 +1533,9 @@ def _fit_MoM_dispersions(self) -> None: normed_counts, self.obsm["design_matrix"].values, ) + normalization_factors = self._get_normalization_factors(self.non_zero_idx) mde = self.inference.fit_moments_dispersions( - normed_counts, self.obs["size_factors"] + normed_counts, normalization_factors ) alpha_hat = np.minimum(rde, mde) @@ -1342,34 +1709,31 @@ def _replace_outliers(self) -> None: self.var["replaced"] = idx.any(axis=0) if sum(self.var["replaced"] > 0): - # Compute replacement counts: trimmed means * size_factors + # Compute replacement counts: trimmed means * normalization factors self.counts_to_refit = self[:, self.var["replaced"]].copy() - - trim_base_mean = pd.DataFrame( - np.asarray( - trimmed_mean( - self.counts_to_refit.X - / self.obs["size_factors"].values[:, None], - trim=0.2, - axis=0, - ) - ), - index=self.counts_to_refit.var_names, + if hasattr(self.counts_to_refit.X, "toarray"): + self.counts_to_refit.X = self.counts_to_refit.X.toarray() + normalization_factors = self._get_normalization_factors( + self.var["replaced"].to_numpy() ) - - replacement_counts = ( - pd.DataFrame( - trim_base_mean.values * self.obs["size_factors"].values, - index=self.counts_to_refit.var_names, - columns=self.counts_to_refit.obs_names, + if normalization_factors.ndim == 1: + normalization_factors = normalization_factors[:, None] + + trim_base_mean = np.asarray( + trimmed_mean( + self.counts_to_refit.X / normalization_factors, + trim=0.2, + axis=0, ) - .astype(int) - .T + ) + replacement_counts = np.asarray( + trim_base_mean[None, :] * normalization_factors, + dtype=int, ) self.counts_to_refit.X[ self.obs["replaceable"].values[:, None] & idx[:, self.var["replaced"]] - ] = replacement_counts.values[ + ] = replacement_counts[ self.obs["replaceable"].values[:, None] & idx[:, self.var["replaced"]] ] @@ -1423,11 +1787,16 @@ def _refit_without_outliers( quiet=self.quiet, ) - # Use the same size factors + # Use the same normalization factors sub_dds.obs["size_factors"] = self.counts_to_refit.obs["size_factors"] - sub_dds.layers["normed_counts"] = ( - sub_dds.X / sub_dds.obs["size_factors"].values[:, None] - ) + if "normalization_factors" in self.counts_to_refit.layers: + sub_dds.layers["normalization_factors"] = self.counts_to_refit.layers[ + "normalization_factors" + ] + normalization_factors = sub_dds._get_normalization_factors() + if normalization_factors.ndim == 1: + normalization_factors = normalization_factors[:, None] + sub_dds.layers["normed_counts"] = sub_dds.X / normalization_factors # Estimate gene-wise dispersions. sub_dds.fit_genewise_dispersions() @@ -1489,9 +1858,13 @@ def _fit_iterate_size_factors(self, niter: int = 10, quant: float = 0.95) -> Non (default: ``0.95``). """ + self.logmeans = None + self.filtered_genes = None + counts = cast(Any, self.X).toarray() if issparse(self.X) else np.asarray(self.X) + # Initialize size factors and normed counts fields self.obs["size_factors"] = np.ones(self.n_obs) - self.layers["normed_counts"] = self.X + self.layers["normed_counts"] = counts # Reduce the design matrix to an intercept and reconstruct at the end self.obsm["design_matrix_buffer"] = self.obsm["design_matrix"].copy() @@ -1503,7 +1876,7 @@ def _fit_iterate_size_factors(self, niter: int = 10, quant: float = 0.95) -> Non def objective(p): sf = np.exp(p - np.mean(p)) nll = nb_nll( - counts=self[:, self.non_zero_genes].X, + counts=counts[:, self.non_zero_idx], mu=self[:, self.non_zero_genes].layers["_mu_hat"] / self.obs["size_factors"].values[:, None] * sf[:, None], @@ -1561,7 +1934,7 @@ def objective(p): del self.obsm["design_matrix_buffer"] # Store normalized counts - self.layers["normed_counts"] = self.X / self.obs["size_factors"].values[:, None] + self.layers["normed_counts"] = counts / self.obs["size_factors"].values[:, None] def _check_full_rank_design(self): """Check that the design matrix has full column rank.""" diff --git a/pydeseq2/default_inference.py b/pydeseq2/default_inference.py index fd735f5e..4ef32fab 100644 --- a/pydeseq2/default_inference.py +++ b/pydeseq2/default_inference.py @@ -1,4 +1,6 @@ +from typing import Any from typing import Literal +from typing import cast import numpy as np import pandas as pd @@ -11,6 +13,24 @@ from pydeseq2 import utils +def _column_sliceable_counts(counts: inference.CountMatrix) -> Any: + """Use CSC storage for efficient repeated sparse column slices.""" + if isinstance(counts, np.ndarray): + return counts + return cast(Any, counts).tocsc(copy=False) + + +def _gene_factors(factors: np.ndarray, gene_idx: int) -> np.ndarray: + return factors[:, gene_idx] if factors.ndim == 2 else factors + + +def _gene_counts(counts: Any, gene_idx: int) -> np.ndarray: + """Return one gene's counts as a dense sample vector.""" + if isinstance(counts, np.ndarray): + return np.asarray(counts[:, gene_idx]).ravel() + return counts[:, [gene_idx]].toarray().ravel() + + class DefaultInference(inference.Inference): """Default DESeq2-related inference methods, using scipy/sklearn/numpy. @@ -57,11 +77,12 @@ def n_cpus(self, n_cpus: int) -> None: def lin_reg_mu( # noqa: D102 self, - counts: np.ndarray, + counts: inference.CountMatrix, size_factors: np.ndarray, design_matrix: np.ndarray, min_mu: float, ) -> np.ndarray: + column_counts = _column_sliceable_counts(counts) with parallel_backend(self._backend, inner_max_num_threads=1): mu_hat_ = np.array( Parallel( @@ -70,19 +91,19 @@ def lin_reg_mu( # noqa: D102 batch_size=self._batch_size, )( delayed(utils.fit_lin_mu)( - counts=counts[:, i], - size_factors=size_factors, + counts=_gene_counts(column_counts, i), + size_factors=_gene_factors(size_factors, i), design_matrix=design_matrix, min_mu=min_mu, ) - for i in range(counts.shape[1]) + for i in range(column_counts.shape[1]) ) ) return mu_hat_.T def irls( # noqa: D102 self, - counts: np.ndarray, + counts: inference.CountMatrix, size_factors: np.ndarray, design_matrix: np.ndarray, disp: np.ndarray, @@ -93,6 +114,7 @@ def irls( # noqa: D102 optimizer: Literal["BFGS", "L-BFGS-B"] = "L-BFGS-B", maxiter: int = 250, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + column_counts = _column_sliceable_counts(counts) with parallel_backend(self._backend, inner_max_num_threads=1): res = Parallel( n_jobs=self.n_cpus, @@ -100,8 +122,8 @@ def irls( # noqa: D102 batch_size=self._batch_size, )( delayed(utils.irls_solver)( - counts=counts[:, i], - size_factors=size_factors, + counts=_gene_counts(column_counts, i), + size_factors=_gene_factors(size_factors, i), design_matrix=design_matrix, disp=disp[i], min_mu=min_mu, @@ -111,7 +133,7 @@ def irls( # noqa: D102 optimizer=optimizer, maxiter=maxiter, ) - for i in range(counts.shape[1]) + for i in range(column_counts.shape[1]) ) res = zip(*res, strict=False) MLE_lfcs_, mu_hat_, hat_diagonals_, converged_ = (np.array(m) for m in res) @@ -125,7 +147,7 @@ def irls( # noqa: D102 def alpha_mle( # noqa: D102 self, - counts: np.ndarray, + counts: inference.CountMatrix, design_matrix: np.ndarray, mu: np.ndarray, alpha_hat: np.ndarray, @@ -136,6 +158,7 @@ def alpha_mle( # noqa: D102 prior_reg: bool = False, optimizer: Literal["BFGS", "L-BFGS-B"] = "L-BFGS-B", ) -> tuple[np.ndarray, np.ndarray]: + column_counts = _column_sliceable_counts(counts) with parallel_backend(self._backend, inner_max_num_threads=1): res = Parallel( n_jobs=self.n_cpus, @@ -143,7 +166,7 @@ def alpha_mle( # noqa: D102 batch_size=self._batch_size, )( delayed(utils.fit_alpha_mle)( - counts=counts[:, i], + counts=_gene_counts(column_counts, i), design_matrix=design_matrix, mu=mu[:, i], alpha_hat=alpha_hat[i], @@ -154,7 +177,7 @@ def alpha_mle( # noqa: D102 prior_reg=prior_reg, optimizer=optimizer, ) - for i in range(counts.shape[1]) + for i in range(column_counts.shape[1]) ) res = zip(*res, strict=False) dispersions_, l_bfgs_b_converged_ = (np.array(m) for m in res) @@ -232,7 +255,7 @@ def grad(coeffs): def lfc_shrink_nbinom_glm( # noqa: D102 self, design_matrix: np.ndarray, - counts: np.ndarray, + counts: inference.CountMatrix, size: np.ndarray, offset: np.ndarray, prior_no_shrink_scale: float, @@ -240,8 +263,8 @@ def lfc_shrink_nbinom_glm( # noqa: D102 optimizer: str, shrink_index: int, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + column_counts = _column_sliceable_counts(counts) with parallel_backend(self._backend, inner_max_num_threads=1): - num_genes = counts.shape[1] res = Parallel( n_jobs=self.n_cpus, verbose=self._joblib_verbosity, @@ -249,15 +272,15 @@ def lfc_shrink_nbinom_glm( # noqa: D102 )( delayed(utils.nbinomGLM)( design_matrix=design_matrix, - counts=counts[:, i], + counts=_gene_counts(column_counts, i), size=size[i], - offset=offset, + offset=_gene_factors(offset, i), prior_no_shrink_scale=prior_no_shrink_scale, prior_scale=prior_scale, optimizer=optimizer, shrink_index=shrink_index, ) - for i in range(num_genes) + for i in range(column_counts.shape[1]) ) res = zip(*res, strict=False) lfcs, inv_hessians, l_bfgs_b_converged_ = (np.array(m) for m in res) diff --git a/pydeseq2/ds.py b/pydeseq2/ds.py index 6e22fde7..300a8e21 100644 --- a/pydeseq2/ds.py +++ b/pydeseq2/ds.py @@ -317,11 +317,10 @@ def run_wald_test(self) -> None: file=sys.stderr, ) - mu = ( - np.exp(self.design_matrix @ self.LFC.T) - .multiply(self.dds.obs["size_factors"], 0) - .values - ) + normalization_factors = self.dds._get_normalization_factors() + if normalization_factors.ndim == 1: + normalization_factors = normalization_factors[:, None] + mu = np.exp(self.design_matrix @ self.LFC.T).to_numpy() * normalization_factors # Set regularization factors. if self.prior_LFC_var is not None: @@ -385,7 +384,7 @@ def lfc_shrink(self, coeff: str, adapt: bool = True) -> None: design_matrix = self.design_matrix.values size = 1.0 / self.dds.var["dispersions"].values - offset = np.log(self.dds.obs["size_factors"]).values + offset = np.log(self.dds._get_normalization_factors(self.dds.non_zero_idx)) # Set priors prior_no_shrink_scale = 15 diff --git a/pydeseq2/inference.py b/pydeseq2/inference.py index 40bff73b..272c94db 100644 --- a/pydeseq2/inference.py +++ b/pydeseq2/inference.py @@ -1,9 +1,14 @@ from abc import ABC from abc import abstractmethod from typing import Literal +from typing import TypeAlias import numpy as np import pandas as pd +from scipy.sparse import sparray +from scipy.sparse import spmatrix + +CountMatrix: TypeAlias = np.ndarray | spmatrix | sparray class Inference(ABC): @@ -12,7 +17,7 @@ class Inference(ABC): @abstractmethod def lin_reg_mu( self, - counts: np.ndarray, + counts: CountMatrix, size_factors: np.ndarray, design_matrix: np.ndarray, min_mu: float, @@ -23,11 +28,14 @@ def lin_reg_mu( Parameters ---------- - counts : ndarray - Raw counts. + counts : ndarray or scipy.sparse.spmatrix or scipy.sparse.sparray + Raw count matrix. size_factors : ndarray - Sample-wise scaling factors (obtained from median-of-ratios). + Sample-wise scaling factors with shape ``(n_samples,)``, or + sample-by-gene normalization factors with shape + ``(n_samples, n_genes)`` (obtained from median-of-ratios and optional + gene-specific offsets). design_matrix : ndarray Design matrix. @@ -45,7 +53,7 @@ def lin_reg_mu( @abstractmethod def irls( self, - counts: np.ndarray, + counts: CountMatrix, size_factors: np.ndarray, design_matrix: np.ndarray, disp: np.ndarray, @@ -62,11 +70,14 @@ def irls( Parameters ---------- - counts : ndarray - Raw counts. + counts : ndarray or scipy.sparse.spmatrix or scipy.sparse.sparray + Raw count matrix. size_factors : ndarray - Sample-wise scaling factors (obtained from median-of-ratios). + Sample-wise scaling factors with shape ``(n_samples,)``, or + sample-by-gene normalization factors with shape + ``(n_samples, n_genes)`` (obtained from median-of-ratios and optional + gene-specific offsets). design_matrix : ndarray Design matrix. @@ -120,7 +131,7 @@ def irls( @abstractmethod def alpha_mle( self, - counts: np.ndarray, + counts: CountMatrix, design_matrix: np.ndarray, mu: np.ndarray, alpha_hat: np.ndarray, @@ -135,8 +146,8 @@ def alpha_mle( Parameters ---------- - counts : ndarray - Raw counts. + counts : ndarray or scipy.sparse.spmatrix or scipy.sparse.sparray + Raw count matrix. design_matrix : ndarray Design matrix. @@ -272,7 +283,9 @@ def fit_moments_dispersions( Array of deseq2-normalized read counts. Rows: samples, columns: genes. size_factors : ndarray - DESeq2 normalization factors. + Sample-wise scaling factors with shape ``(n_samples,)``, or + sample-by-gene normalization factors with shape + ``(n_samples, n_genes)``. Returns ------- @@ -310,7 +323,7 @@ def dispersion_trend_gamma_glm( def lfc_shrink_nbinom_glm( self, design_matrix: np.ndarray, - counts: np.ndarray, + counts: CountMatrix, size: np.ndarray, offset: np.ndarray, prior_no_shrink_scale: float, @@ -327,14 +340,16 @@ def lfc_shrink_nbinom_glm( design_matrix : ndarray Design matrix. - counts : ndarray - Raw counts. + counts : ndarray or scipy.sparse.spmatrix or scipy.sparse.sparray + Raw count matrix. size : ndarray Size parameter of NB family (inverse of dispersion). offset : ndarray - Natural logarithm of size factor. + Natural logarithm of sample-wise size factors with shape + ``(n_samples,)``, or sample-by-gene normalization factors with shape + ``(n_samples, n_genes)``. prior_no_shrink_scale : float Prior variance for the intercept. diff --git a/pydeseq2/utils.py b/pydeseq2/utils.py index 5fcb50e6..3bfe3312 100644 --- a/pydeseq2/utils.py +++ b/pydeseq2/utils.py @@ -2,6 +2,7 @@ from math import ceil from math import floor from pathlib import Path +from typing import Any from typing import Literal from typing import cast @@ -10,6 +11,8 @@ from matplotlib import pyplot as plt from scipy.linalg import solve # type: ignore from scipy.optimize import minimize # type: ignore +from scipy.sparse import sparray # type: ignore +from scipy.sparse import spmatrix # type: ignore from scipy.special import gammaln # type: ignore from scipy.special import polygamma # type: ignore from scipy.stats import norm # type: ignore @@ -107,14 +110,16 @@ def load_example_data( return df -def test_valid_counts(counts: pd.DataFrame | np.ndarray) -> None: +def test_valid_counts( + counts: pd.DataFrame | np.ndarray | spmatrix | sparray, +) -> None: """Test that the count matrix contains valid inputs. More precisely, test that inputs are non-negative integers. Parameters ---------- - counts : pandas.DataFrame or ndarray + counts : pandas.DataFrame, ndarray, scipy.sparse.spmatrix or scipy.sparse.sparray Raw counts. One column per gene, rows are indexed by sample barcodes. """ if isinstance(counts, pd.DataFrame): @@ -122,6 +127,17 @@ def test_valid_counts(counts: pd.DataFrame | np.ndarray) -> None: raise ValueError("NaNs are not allowed in the count matrix.") if not np.issubdtype(counts.to_numpy().dtype, np.number): raise ValueError("The count matrix should only contain numbers.") + elif isinstance(counts, (spmatrix, sparray)): + sparse_counts = cast(Any, counts) + if sparse_counts.format not in {"coo", "csc", "csr"}: + values = np.asarray(sparse_counts.tocoo(copy=False).data) + else: + values = np.asarray(sparse_counts.data) + if not np.issubdtype(values.dtype, np.number): + raise ValueError("The count matrix should only contain numbers.") + if np.isnan(values).any(): + raise ValueError("NaNs are not allowed in the count matrix.") + counts = values else: if np.isnan(counts).any().any(): raise ValueError("NaNs are not allowed in the count matrix.") @@ -875,9 +891,14 @@ def fit_moments_dispersions( Estimated dispersion parameter for each gene. """ # Exclude genes with all zeroes - normed_counts = normed_counts[:, ~(normed_counts == 0).all(axis=0)] - # mean inverse size factor - s_mean_inv = (1 / size_factors).mean(axis=0) + nonzero_genes = ~(normed_counts == 0).all(axis=0) + normed_counts = normed_counts[:, nonzero_genes] + # Mean inverse size factor. For gene-specific normalization factors, DESeq2 + # first averages factors across genes within each sample, then averages the + # inverse sample means into one scalar shared by all genes. + if size_factors.ndim != 1: + size_factors = size_factors[:, nonzero_genes].mean(axis=1) + s_mean_inv = (1 / size_factors).mean() mu = normed_counts.mean(0) sigma = normed_counts.var(0, ddof=1) # ddof=1 is to use an unbiased estimator, as in R diff --git a/tests/test_transcript_length_normalization.py b/tests/test_transcript_length_normalization.py new file mode 100644 index 00000000..f5bb1f0e --- /dev/null +++ b/tests/test_transcript_length_normalization.py @@ -0,0 +1,1166 @@ +import warnings + +import anndata as ad +import numpy as np +import pandas as pd +import pytest +from scipy.sparse import bsr_array +from scipy.sparse import bsr_matrix +from scipy.sparse import coo_array +from scipy.sparse import coo_matrix +from scipy.sparse import csc_array +from scipy.sparse import csc_matrix +from scipy.sparse import csr_array +from scipy.sparse import csr_matrix +from scipy.sparse import dia_array +from scipy.sparse import dia_matrix +from scipy.sparse import dok_array +from scipy.sparse import dok_matrix +from scipy.sparse import lil_array +from scipy.sparse import lil_matrix + +from pydeseq2.dds import DeseqDataSet +from pydeseq2.default_inference import DefaultInference +from pydeseq2.ds import DeseqStats +from pydeseq2.utils import fit_moments_dispersions +from pydeseq2.utils import load_example_data + + +def small_tximport_data(): + samples = pd.Index([f"sample{i}" for i in range(1, 7)]) + genes = pd.Index(["gene1", "gene2", "gene3"]) + counts = pd.DataFrame( + [ + [100.2, 201.7, 302.1], + [121.8, 181.2, 330.9], + [89.6, 240.4, 281.2], + [170.1, 260.8, 410.3], + [160.7, 290.2, 389.9], + [190.4, 250.6, 430.8], + ], + index=samples, + columns=genes, + ) + metadata = pd.DataFrame( + {"condition": ["A", "A", "A", "B", "B", "B"]}, + index=samples, + ) + transcript_lengths = pd.DataFrame( + [ + [900.0, 1400.0, 2100.0], + [950.0, 1350.0, 2050.0], + [1000.0, 1300.0, 2000.0], + [1250.0, 1100.0, 1850.0], + [1300.0, 1050.0, 1800.0], + [1350.0, 1000.0, 1750.0], + ], + index=samples, + columns=genes, + ) + return counts, metadata, transcript_lengths + + +def small_pytximport_adata(): + counts, metadata, transcript_lengths = small_tximport_data() + adata = ad.AnnData( + X=counts.to_numpy(), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + adata.obsm["length"] = transcript_lengths.to_numpy() + adata.uns["counts_from_abundance"] = None + return adata, counts, metadata, transcript_lengths + + +def varying_transcript_lengths(counts): + sample_effect = np.linspace(-1.0, 1.0, counts.shape[0]) + gene_effect = np.linspace(-1.0, 1.0, counts.shape[1]) + lengths = (800.0 + 100.0 * np.arange(counts.shape[1]))[None, :] + lengths = lengths * (1.0 + 0.2 * np.outer(sample_effect, gene_effect)) + return pd.DataFrame(lengths, index=counts.index, columns=counts.columns) + + +def anndata_with_sparse_counts(sparse_counts, *, obs=None, var=None): + try: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="AnnData previously had undefined behavior around matrices", + category=FutureWarning, + ) + return ad.AnnData(X=sparse_counts, obs=obs, var=var) + except ValueError as error: + if not str(error).startswith("Only CSR and CSC"): + raise + pytest.skip(f"{type(sparse_counts).__name__} is not supported by AnnData") + + +def assert_fit_state_cleared(dds): + assert set(dds.obs).isdisjoint({"size_factors", "replaceable"}) + assert set(dds.var).isdisjoint( + { + "_normed_means", + "non_zero", + "_MoM_dispersions", + "genewise_dispersions", + "vst_genewise_dispersions", + "_genewise_converged", + "fitted_dispersions", + "MAP_dispersions", + "_MAP_converged", + "dispersions", + "_outlier_genes", + "_LFC_converged", + "_pvalue_cooks_outlier", + "replaced", + "refitted", + } + ) + assert set(dds.obsm).isdisjoint({"_mu_LFC", "_hat_diagonals"}) + assert "LFC" not in dds.varm + assert set(dds.layers).isdisjoint( + { + "normalization_factors", + "normed_counts", + "_mu_hat", + "_vst_mu_hat", + "cooks", + "replace_cooks", + "vst_counts", + } + ) + assert set(dds.uns).isdisjoint( + { + "trend_coeffs", + "vst_trend_coeffs", + "mean_disp", + "disp_function_type", + "_squared_logres", + "prior_disp_var", + } + ) + for attr in ( + "non_zero_idx", + "non_zero_genes", + "counts_to_refit", + "new_all_zeroes_genes", + ): + assert not hasattr(dds, attr) + + +def test_transcript_length_normalization_matches_deseq2_formula(): + counts, metadata, transcript_lengths = small_tximport_data() + dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + design="~condition", + quiet=True, + ) + dds.fit_size_factors() + + rounded_counts = np.rint(counts.to_numpy()) + lengths = transcript_lengths.to_numpy() + length_factors = lengths / np.exp(np.mean(np.log(lengths), axis=0)) + adjusted_counts = rounded_counts / length_factors + logmeans = np.mean(np.log(adjusted_counts), axis=0) + size_factors = np.exp(np.median(np.log(adjusted_counts) - logmeans, axis=1)) + size_factors /= np.exp(np.mean(np.log(size_factors))) + expected_factors = length_factors * size_factors[:, None] + expected_factors /= np.exp(np.mean(np.log(expected_factors), axis=0)) + + np.testing.assert_array_equal(dds.X, rounded_counts) + np.testing.assert_allclose(dds.layers["avg_tx_length"], lengths) + np.testing.assert_allclose(dds.obs["size_factors"], size_factors) + np.testing.assert_allclose(dds.layers["normalization_factors"], expected_factors) + np.testing.assert_allclose( + dds.layers["normed_counts"], rounded_counts / expected_factors + ) + np.testing.assert_allclose( + np.exp(np.mean(np.log(dds.layers["normalization_factors"]), axis=0)), + np.ones(dds.n_vars), + ) + + +def test_poscounts_transcript_length_normalization_matches_deseq2_formula(): + counts, metadata, transcript_lengths = small_tximport_data() + for idx in range(3): + counts.iat[idx, idx] = 0 + dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + quiet=True, + ) + dds.fit_size_factors("poscounts") + + rounded_counts = np.rint(counts.to_numpy()) + lengths = transcript_lengths.to_numpy() + length_factors = lengths / np.exp(np.mean(np.log(lengths), axis=0)) + adjusted_counts = rounded_counts / length_factors + log_counts = np.zeros_like(rounded_counts) + np.log(rounded_counts, out=log_counts, where=rounded_counts != 0) + logmeans = log_counts.mean(axis=0) + size_factors = np.empty(dds.n_obs) + for sample_idx in range(dds.n_obs): + positive = adjusted_counts[sample_idx] > 0 + size_factors[sample_idx] = np.exp( + np.median(np.log(adjusted_counts[sample_idx, positive]) - logmeans[positive]) + ) + size_factors /= np.exp(np.mean(np.log(size_factors))) + expected_factors = length_factors * size_factors[:, None] + expected_factors /= np.exp(np.mean(np.log(expected_factors), axis=0)) + + np.testing.assert_allclose(dds.layers["normalization_factors"], expected_factors) + + +@pytest.mark.parametrize( + ("length_source", "bad_value"), + [ + ("explicit", 0.0), + ("explicit", np.nan), + ("pytximport", np.inf), + ], +) +def test_transcript_lengths_must_be_positive_and_finite(length_source, bad_value): + if length_source == "explicit": + counts, metadata, transcript_lengths = small_tximport_data() + transcript_lengths.iloc[0, 0] = bad_value + kwargs = { + "counts": counts, + "metadata": metadata, + "transcript_lengths": transcript_lengths, + } + else: + adata, _, _, _ = small_pytximport_adata() + lengths = np.asarray(adata.obsm["length"]).copy() + lengths[0, 0] = bad_value + adata.obsm["length"] = lengths + kwargs = {"adata": adata} + + with pytest.raises(ValueError, match="transcript_lengths"): + DeseqDataSet(**kwargs) + + +def test_transcript_length_labels_must_match_counts(): + counts, metadata, transcript_lengths = small_tximport_data() + transcript_lengths = transcript_lengths.rename(index={"sample1": "wrong_sample"}) + + with pytest.raises(ValueError, match="same sample index"): + DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + ) + + +def test_transcript_lengths_reject_iterative_size_factors(): + counts, metadata, transcript_lengths = small_tximport_data() + dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + ) + + with pytest.raises(ValueError, match="iterative"): + dds.fit_size_factors("iterative") + + +def test_external_vst_transform_requires_matching_transcript_lengths(): + counts, metadata, transcript_lengths = small_tximport_data() + dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + ) + dds.fit_size_factors() + + with pytest.raises(ValueError, match="matching transcript lengths"): + dds.vst_transform(dds.X) + + +def test_transcript_length_factors_are_used_throughout_pipeline(monkeypatch): + counts = load_example_data("raw_counts", "synthetic", debug=False) + metadata = load_example_data("metadata", "synthetic", debug=False) + transcript_lengths = varying_transcript_lengths(counts) + adata = ad.AnnData( + X=csc_matrix(counts.to_numpy()), + obs=metadata, + var=pd.DataFrame(index=counts.columns), + ) + adata.obsm["length"] = transcript_lengths.to_numpy() + adata.uns["counts_from_abundance"] = None + + dds = DeseqDataSet( + adata=adata, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + reference_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + dds.deseq2() + reference_dds.deseq2() + + np.testing.assert_allclose( + dds.layers["normalization_factors"], + reference_dds.layers["normalization_factors"], + ) + dispersion_columns = ["_MoM_dispersions", "genewise_dispersions", "dispersions"] + np.testing.assert_allclose( + dds.var[dispersion_columns].to_numpy(), + reference_dds.var[dispersion_columns].to_numpy(), + ) + np.testing.assert_allclose(dds.varm["LFC"], reference_dds.varm["LFC"]) + + non_zero = dds.var["non_zero"].to_numpy() + normalization_factors = dds.layers["normalization_factors"][:, non_zero] + design_matrix = dds.obsm["design_matrix"].to_numpy() + coefficients = dds.varm["LFC"].loc[dds.var_names[non_zero]].to_numpy() + expected_mu = np.maximum( + normalization_factors * np.exp(design_matrix @ coefficients.T), + dds.min_mu, + ) + np.testing.assert_allclose(dds.obsm["_mu_LFC"], expected_mu) + + stats = DeseqStats(dds, contrast=["condition", "B", "A"], quiet=True) + reference_stats = DeseqStats( + reference_dds, contrast=["condition", "B", "A"], quiet=True + ) + stats.summary() + reference_stats.summary() + assert stats.results_df["pvalue"].notna().any() + np.testing.assert_allclose( + stats.results_df[["stat", "pvalue", "lfcSE"]], + reference_stats.results_df[["stat", "pvalue", "lfcSE"]], + equal_nan=True, + ) + + coefficient = stats.LFC.columns[-1] + stats.lfc_shrink(coeff=coefficient, adapt=False) + reference_stats.lfc_shrink(coeff=coefficient, adapt=False) + assert stats.shrunk_LFCs + np.testing.assert_allclose( + stats.LFC[coefficient], reference_stats.LFC[coefficient], equal_nan=True + ) + monkeypatch.setattr(dds, "fit_size_factors", lambda *args, **kwargs: pytest.fail()) + dds.vst(fit_type="mean") + reference_dds.vst(fit_type="mean") + np.testing.assert_allclose( + dds.layers["vst_counts"], reference_dds.layers["vst_counts"] + ) + assert isinstance(dds.X, csc_matrix) + + restored_dds = DeseqDataSet( + adata=dds, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + assert restored_dds.vst_fit_type == dds.vst_fit_type + np.testing.assert_allclose(restored_dds.vst_transform(), dds.vst_transform()) + + +@pytest.mark.parametrize("sparse_constructor", [csr_matrix, csr_array]) +def test_sparse_anndata_ratio_pipeline_matches_dense(sparse_constructor): + counts = load_example_data("raw_counts", "synthetic", debug=False) + metadata = load_example_data("metadata", "synthetic", debug=False) + adata = ad.AnnData( + X=sparse_constructor(counts.to_numpy()), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + + dds = DeseqDataSet( + adata=adata, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + reference_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + dds.deseq2() + reference_dds.deseq2() + + np.testing.assert_allclose( + dds.obs["size_factors"], reference_dds.obs["size_factors"] + ) + np.testing.assert_allclose( + dds.layers["normed_counts"], reference_dds.layers["normed_counts"] + ) + dispersion_columns = ["_MoM_dispersions", "genewise_dispersions", "dispersions"] + np.testing.assert_allclose( + dds.var[dispersion_columns], reference_dds.var[dispersion_columns] + ) + np.testing.assert_allclose(dds.varm["LFC"], reference_dds.varm["LFC"]) + assert isinstance(dds.X, sparse_constructor) + + +@pytest.mark.parametrize( + "sparse_constructor", + [ + bsr_matrix, + coo_matrix, + dia_matrix, + dok_matrix, + lil_matrix, + bsr_array, + coo_array, + dia_array, + dok_array, + lil_array, + ], +) +def test_general_sparse_anndata_pipeline_matches_dense(sparse_constructor): + counts, metadata, _ = small_tximport_data() + counts = counts.round().astype(int) + adata = anndata_with_sparse_counts( + sparse_constructor(counts.to_numpy()), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + + dds = DeseqDataSet(adata=adata, design="~condition", quiet=True) + assert dds.X.format == "csr" + np.testing.assert_array_equal(dds.X.toarray(), counts.to_numpy()) + + +@pytest.mark.parametrize( + "sparse_constructor", + [coo_matrix, csr_array, csc_matrix], +) +def test_sparse_anndata_sums_duplicates_before_validation(sparse_constructor): + data = np.array([0.5, 0.5, 2.0]) + if sparse_constructor is coo_matrix: + sparse_counts = sparse_constructor( + (data, (np.array([0, 0, 1]), np.array([0, 0, 1]))), + shape=(2, 2), + ) + else: + sparse_counts = sparse_constructor( + (data, np.array([0, 0, 1]), np.array([0, 2, 3])), + shape=(2, 2), + ) + + adata = anndata_with_sparse_counts(sparse_counts) + source_data = adata.X.data.copy() + + dds = DeseqDataSet(adata=adata, design="~1", quiet=True) + + np.testing.assert_array_equal(dds.X.toarray(), [[1, 0], [0, 2]]) + assert dds.X.nnz == 2 + assert adata.X.nnz == 3 + np.testing.assert_array_equal(adata.X.data, source_data) + + +def test_general_sparse_transcript_length_counts(): + counts, metadata, transcript_lengths = small_tximport_data() + adata = anndata_with_sparse_counts( + dok_array(counts.to_numpy()), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + + dds = DeseqDataSet( + adata=adata, + transcript_lengths=transcript_lengths, + design="~condition", + quiet=True, + ) + dds.fit_size_factors() + + assert dds.X.format == "csr" + np.testing.assert_array_equal(dds.X.toarray(), np.rint(counts.to_numpy())) + assert "normalization_factors" in dds.layers + assert dds.layers["normalization_factors"].shape == dds.shape + + +def test_sparse_anndata_poscounts_size_factors_match_dense(): + counts, metadata, _ = small_tximport_data() + counts = counts.round().astype(int) + for idx in range(counts.shape[1]): + counts.iat[idx, idx] = 0 + adata = ad.AnnData( + X=csc_array(counts.to_numpy()), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + dds = DeseqDataSet(adata=adata, design="~condition", quiet=True) + reference_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + design="~condition", + quiet=True, + ) + + dds.fit_size_factors("poscounts") + reference_dds.fit_size_factors("poscounts") + + np.testing.assert_allclose( + dds.obs["size_factors"], reference_dds.obs["size_factors"] + ) + np.testing.assert_allclose( + dds.layers["normed_counts"], reference_dds.layers["normed_counts"] + ) + assert isinstance(dds.layers["normed_counts"], np.ndarray) + assert isinstance(dds.X, csc_array) + + +def test_sparse_anndata_iterative_size_factors_match_dense(): + counts = load_example_data("raw_counts", "synthetic", debug=False).iloc[:20] + metadata = load_example_data("metadata", "synthetic", debug=False).loc[counts.index] + adata = ad.AnnData( + X=csr_array(counts.to_numpy()), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + dds = DeseqDataSet(adata=adata, design="~1", n_cpus=1, quiet=True) + reference_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + design="~1", + n_cpus=1, + quiet=True, + ) + for dataset in (dds, reference_dds): + dataset.logmeans = np.ones(dataset.n_vars) + dataset.filtered_genes = np.ones(dataset.n_vars, dtype=bool) + dataset._fit_iterate_size_factors(niter=1) + + np.testing.assert_allclose( + dds.obs["size_factors"], reference_dds.obs["size_factors"] + ) + np.testing.assert_allclose( + dds.layers["normed_counts"], reference_dds.layers["normed_counts"] + ) + assert dds.logmeans is None + assert dds.filtered_genes is None + assert isinstance(dds.layers["normed_counts"], np.ndarray) + assert isinstance(dds.X, csr_array) + + +@pytest.mark.parametrize("sparse_constructor", [csr_matrix, csr_array]) +def test_sparse_ratio_falls_back_to_iterative(monkeypatch, sparse_constructor): + counts, metadata, _ = small_tximport_data() + counts = counts.round().astype(int) + for idx in range(counts.shape[1]): + counts.iat[idx, idx] = 0 + adata = ad.AnnData( + X=sparse_constructor(counts.to_numpy()), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + dds = DeseqDataSet(adata=adata, design="~condition", quiet=True) + iterative_called = False + + def fit_iterative(): + nonlocal iterative_called + iterative_called = True + dds.obs["size_factors"] = np.ones(dds.n_obs) + dds.layers["normed_counts"] = dds.X.toarray() + dds.logmeans = None + dds.filtered_genes = None + + monkeypatch.setattr(dds, "_fit_iterate_size_factors", fit_iterative) + + with pytest.warns(UserWarning, match="Switching to iterative mode"): + dds.fit_size_factors("ratio") + + assert iterative_called + + +def test_sparse_array_cooks_outlier(): + counts, metadata, _ = small_tximport_data() + counts = counts.round().astype(int) + adata = ad.AnnData( + X=csr_array(counts.to_numpy()), + obs=metadata.copy(), + var=pd.DataFrame(index=counts.columns), + ) + dds = DeseqDataSet(adata=adata, design="~condition", refit_cooks=False, quiet=True) + dds.layers["cooks"] = np.zeros(dds.shape) + dds.layers["cooks"][-1, 0] = np.inf + np.testing.assert_array_equal(dds.cooks_outlier(), [True, False, False]) + + +def test_moments_dispersions_match_deseq2_with_matrix_factors(): + counts, metadata, transcript_lengths = small_tximport_data() + dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + quiet=True, + ) + dds.fit_size_factors() + + normalization_factors = dds.layers["normalization_factors"] + normed_counts = dds.layers["normed_counts"] + mean_inverse_factor = np.mean(1 / normalization_factors.mean(axis=1)) + means = normed_counts.mean(axis=0) + variances = normed_counts.var(axis=0, ddof=1) + expected = np.nan_to_num((variances - mean_inverse_factor * means) / means**2) + + np.testing.assert_allclose( + fit_moments_dispersions(normed_counts, normalization_factors), + expected, + ) + + +def test_moments_dispersions_align_matrix_factors_after_zero_gene_filtering(): + normed_counts = np.array( + [ + [0.0, 10.0], + [0.0, 14.0], + [0.0, 18.0], + ] + ) + normalization_factors = np.array( + [ + [100.0, 1.0], + [100.0, 1.0], + [100.0, 1.0], + ] + ) + + np.testing.assert_allclose( + fit_moments_dispersions(normed_counts, normalization_factors), + np.array([2.0 / 196.0]), + ) + + +def test_sparse_estimated_counts_round_after_summing_duplicates(): + counts = csr_matrix( + (np.array([0.4, 0.4, 0.4]), np.array([0, 0, 1]), np.array([0, 2, 3])), + shape=(2, 2), + ) + adata = ad.AnnData(X=counts) + adata.obsm["length"] = np.ones((2, 2)) + adata.uns["counts_from_abundance"] = None + + dds = DeseqDataSet(adata=adata, design="~1", quiet=True) + + np.testing.assert_array_equal(dds.X.toarray(), np.array([[1, 0], [0, 0]])) + assert dds.X.nnz == 1 + assert adata.X.nnz == 3 + + +def test_matching_integer_transcript_length_labels_are_accepted(): + counts, metadata, transcript_lengths = small_tximport_data() + integer_samples = pd.Index(range(1, len(counts) + 1)) + integer_genes = pd.Index(range(101, 101 + counts.shape[1])) + counts.index = integer_samples + counts.columns = integer_genes + metadata.index = integer_samples + transcript_lengths.index = integer_samples + transcript_lengths.columns = integer_genes + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ad.ImplicitModificationWarning) + dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + quiet=True, + ) + + np.testing.assert_allclose( + dds.layers["avg_tx_length"], transcript_lengths.to_numpy() + ) + + +def test_scalar_size_factor_refit_clears_stale_matrix_factors(): + counts, metadata, _ = small_tximport_data() + counts = counts.round().astype(int) + dds = DeseqDataSet(counts=counts, metadata=metadata, quiet=True) + dds.layers["normalization_factors"] = np.full(dds.shape, 2.0) + + dds.fit_size_factors() + + assert "normalization_factors" not in dds.layers + np.testing.assert_allclose( + dds.layers["normed_counts"], + dds.X / dds.obs["size_factors"].to_numpy()[:, None], + ) + + +@pytest.mark.parametrize("length_source", ["canonical", "pytximport"]) +def test_implicit_transcript_lengths_clear_scalar_normalization_state( + length_source, +): + counts, metadata, transcript_lengths = small_tximport_data() + counts = counts.round().astype(int) + source = DeseqDataSet( + counts=counts, + metadata=metadata, + design="~condition", + quiet=True, + ) + source.fit_size_factors() + source_size_factors = source.obs["size_factors"].copy() + source_normed_counts = source.layers["normed_counts"].copy() + + if length_source == "canonical": + source.layers["avg_tx_length"] = transcript_lengths.to_numpy() + else: + source.obsm["length"] = transcript_lengths.to_numpy() + source.uns["counts_from_abundance"] = None + + dds = DeseqDataSet(adata=source, design="~condition", quiet=True) + assert "size_factors" not in dds.obs + assert set(dds.layers).isdisjoint({"normalization_factors", "normed_counts"}) + + reference_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + design="~condition", + quiet=True, + ) + dds.fit_genewise_dispersions() + reference_dds.fit_genewise_dispersions() + + np.testing.assert_allclose( + dds.layers["normalization_factors"], + reference_dds.layers["normalization_factors"], + ) + np.testing.assert_allclose( + dds.layers["normed_counts"], reference_dds.layers["normed_counts"] + ) + np.testing.assert_allclose( + dds.var["genewise_dispersions"], reference_dds.var["genewise_dispersions"] + ) + pd.testing.assert_series_equal(source.obs["size_factors"], source_size_factors) + np.testing.assert_allclose(source.layers["normed_counts"], source_normed_counts) + + +def test_inherited_normalization_uses_effective_fit_settings(): + counts, metadata, transcript_lengths = small_tximport_data() + source = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + design="~condition", + quiet=True, + ) + source.fit_size_factors( + fit_type="poscounts", + control_genes=["gene1", "gene2"], + ) + source_factors = source.layers["normalization_factors"].copy() + + mismatched_dds = DeseqDataSet( + adata=source, + design="~condition", + quiet=True, + ) + assert "size_factors" not in mismatched_dds.obs + assert "normalization_factors" not in mismatched_dds.layers + + matching_dds = DeseqDataSet( + adata=source, + design="~condition", + size_factors_fit_type="poscounts", + control_genes=["gene1", "gene2"], + quiet=True, + ) + np.testing.assert_allclose( + matching_dds.layers["normalization_factors"], + source_factors, + ) + np.testing.assert_allclose(source.layers["normalization_factors"], source_factors) + + +def test_explicit_transcript_lengths_clear_inherited_fitted_state(): + counts = load_example_data("raw_counts", "synthetic", debug=False) + metadata = load_example_data("metadata", "synthetic", debug=False) + transcript_lengths = varying_transcript_lengths(counts) + fitted_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + fitted_dds.deseq2() + original_size_factors = fitted_dds.obs["size_factors"].copy() + original_normalization_factors = fitted_dds.layers["normalization_factors"].copy() + original_normed_counts = fitted_dds.layers["normed_counts"].copy() + + mismatched_design_dds = DeseqDataSet( + adata=fitted_dds, + design="~condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + mismatched_refit_dds = DeseqDataSet( + adata=fitted_dds, + design="~0 + condition", + fit_type="mean", + refit_cooks=True, + n_cpus=1, + quiet=True, + ) + for mismatched_dds in (mismatched_design_dds, mismatched_refit_dds): + assert_fit_state_cleared(mismatched_dds) + + restored_dds = DeseqDataSet( + adata=fitted_dds, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + np.testing.assert_array_equal(restored_dds.non_zero_idx, fitted_dds.non_zero_idx) + assert restored_dds.non_zero_genes.equals(fitted_dds.non_zero_genes) + restored_stats = DeseqStats( + restored_dds, contrast=["condition", "B", "A"], quiet=True, n_cpus=1 + ) + restored_stats.summary() + coefficient = restored_stats.LFC.columns[-1] + restored_stats.lfc_shrink(coeff=coefficient, adapt=False) + assert restored_stats.shrunk_LFCs + + replacement_lengths = transcript_lengths.copy() + replacement_lengths.iloc[:, 0] *= np.linspace(1.0, 2.0, len(replacement_lengths)) + dds = DeseqDataSet( + adata=fitted_dds, + transcript_lengths=replacement_lengths, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + + assert_fit_state_cleared(dds) + assert "size_factors" in fitted_dds.obs + assert "genewise_dispersions" in fitted_dds.var + assert "_mu_LFC" in fitted_dds.obsm + assert "LFC" in fitted_dds.varm + assert "cooks" in fitted_dds.layers + assert "mean_disp" in fitted_dds.uns + + reference_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=replacement_lengths, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + dds.fit_dispersion_trend() + reference_dds.fit_dispersion_trend() + dds.fit_LFC() + reference_dds.fit_LFC() + + np.testing.assert_allclose( + dds.layers["normalization_factors"], + reference_dds.layers["normalization_factors"], + ) + np.testing.assert_allclose(dds.var["dispersions"], reference_dds.var["dispersions"]) + np.testing.assert_allclose(dds.varm["LFC"], reference_dds.varm["LFC"]) + np.testing.assert_allclose(fitted_dds.obs["size_factors"], original_size_factors) + np.testing.assert_allclose( + fitted_dds.layers["normalization_factors"], + original_normalization_factors, + ) + np.testing.assert_allclose( + fitted_dds.layers["normed_counts"], original_normed_counts + ) + + +def test_inherited_fit_requires_complete_configuration_provenance(): + class CustomInference(DefaultInference): + pass + + counts = load_example_data("raw_counts", "synthetic", debug=False) + metadata = load_example_data("metadata", "synthetic", debug=False) + transcript_lengths = varying_transcript_lengths(counts) + custom_inference = CustomInference(n_cpus=1) + fitted_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + inference=custom_inference, + n_cpus=1, + quiet=True, + ) + fitted_dds.deseq2() + + picklable_adata = fitted_dds.to_picklable_anndata() + picklable_adata.inference = custom_inference + missing_provenance_dds = DeseqDataSet( + adata=picklable_adata, + design="~0 + condition", + fit_type="parametric", + refit_cooks=True, + inference=custom_inference, + n_cpus=1, + quiet=True, + ) + omitted_custom_inference_dds = DeseqDataSet( + adata=fitted_dds, + design="~0 + condition", + fit_type="mean", + refit_cooks=False, + n_cpus=1, + quiet=True, + ) + + for incompatible_dds in ( + missing_provenance_dds, + omitted_custom_inference_dds, + ): + assert_fit_state_cleared(incompatible_dds) + + assert "LFC" in fitted_dds.varm + assert hasattr(fitted_dds, "non_zero_idx") + + +def test_cooks_outlier_refit_preserves_matrix_factors(): + counts = load_example_data("raw_counts", "synthetic", debug=False) + metadata = load_example_data("metadata", "synthetic", debug=False) + counts.iloc[0, 0] *= 10_000 + transcript_lengths = varying_transcript_lengths(counts) + adata = ad.AnnData( + X=csr_matrix(counts.to_numpy()), + obs=metadata, + var=pd.DataFrame(index=counts.columns), + ) + adata.obsm["length"] = transcript_lengths.to_numpy() + adata.uns["counts_from_abundance"] = None + + dds = DeseqDataSet( + adata=adata, + design="~0 + condition", + fit_type="mean", + refit_cooks=True, + min_replicates=3, + n_cpus=1, + quiet=True, + ) + reference_dds = DeseqDataSet( + counts=counts, + metadata=metadata, + transcript_lengths=transcript_lengths, + design="~0 + condition", + fit_type="mean", + refit_cooks=True, + min_replicates=3, + n_cpus=1, + quiet=True, + ) + dds.deseq2() + reference_dds.deseq2() + + assert dds.var["replaced"].sum() == 1 + assert dds.var["refitted"].sum() == 1 + np.testing.assert_array_equal( + dds.var[["replaced", "refitted"]], + reference_dds.var[["replaced", "refitted"]], + ) + np.testing.assert_allclose( + dds.layers["cooks"], + reference_dds.layers["cooks"], + rtol=1e-7, + atol=1e-10, + equal_nan=True, + ) + np.testing.assert_array_equal(dds.counts_to_refit.X, reference_dds.counts_to_refit.X) + refitted = dds.var["refitted"].to_numpy() + np.testing.assert_allclose( + dds.counts_to_refit.layers["normalization_factors"], + dds.layers["normalization_factors"][:, refitted], + ) + np.testing.assert_allclose( + dds.var.loc[refitted, ["genewise_dispersions", "dispersions"]], + reference_dds.var.loc[refitted, ["genewise_dispersions", "dispersions"]], + ) + np.testing.assert_allclose( + dds.varm["LFC"].loc[dds.var_names[refitted]], + reference_dds.varm["LFC"].loc[reference_dds.var_names[refitted]], + ) + assert isinstance(dds.X, csr_matrix) + + source_lfc = dds.varm["LFC"].copy() + restored_dds = DeseqDataSet( + adata=dds, + design="~0 + condition", + fit_type="mean", + refit_cooks=True, + min_replicates=3, + n_cpus=1, + quiet=True, + ) + assert restored_dds.new_all_zeroes_genes.equals(dds.new_all_zeroes_genes) + assert restored_dds.counts_to_refit is not dds.counts_to_refit + np.testing.assert_array_equal(restored_dds.counts_to_refit.X, dds.counts_to_refit.X) + np.testing.assert_allclose( + restored_dds.counts_to_refit.layers["normalization_factors"], + dds.counts_to_refit.layers["normalization_factors"], + ) + restored_dds.varm["LFC"].iloc[0, 0] += 1 + pd.testing.assert_frame_equal(dds.varm["LFC"], source_lfc) + + +def test_pytximport_backed_anndata_requires_memory(tmp_path): + adata = ad.AnnData(X=np.ones((2, 2))) + adata.obsm["length"] = np.ones((2, 2)) + adata.uns["counts_from_abundance"] = None + path = tmp_path / "backed.h5ad" + adata.write_h5ad(path) + backed = ad.read_h5ad(path, backed="r") + + try: + with pytest.raises(ValueError, match="to_memory"): + DeseqDataSet(adata=backed, quiet=True) + finally: + backed.file.close() + + +def test_pytximport_dataframe_length_labels_are_preserved(): + adata, _, _, transcript_lengths = small_pytximport_adata() + adata.obsm["length"] = transcript_lengths.copy() + + dds = DeseqDataSet(adata=adata, quiet=True) + + assert dds.obs_names.equals(transcript_lengths.index) + assert dds.var_names.equals(transcript_lengths.columns) + np.testing.assert_allclose( + dds.layers["avg_tx_length"], transcript_lengths.to_numpy() + ) + + +def test_pytximport_dataframe_gene_labels_must_match(): + adata, _, _, transcript_lengths = small_pytximport_adata() + adata.obsm["length"] = transcript_lengths.rename(columns={"gene1": "wrong_gene"}) + + with pytest.raises(ValueError, match="same gene columns"): + DeseqDataSet(adata=adata, quiet=True) + + +def test_pytximport_scaled_count_mode_is_rejected(): + adata, _, _, _ = small_pytximport_adata() + adata.uns["counts_from_abundance"] = "length_scaled_tpm" + + with pytest.raises(ValueError, match="unscaled estimated counts"): + DeseqDataSet(adata=adata, quiet=True) + + +@pytest.mark.parametrize("higher_precedence_source", ["explicit", "canonical"]) +def test_scaled_pytximport_mode_rejects_higher_precedence_lengths( + higher_precedence_source, +): + adata, _, _, transcript_lengths = small_pytximport_adata() + adata.uns["counts_from_abundance"] = "length_scaled_tpm" + del adata.obsm["length"] + explicit_lengths = None + if higher_precedence_source == "explicit": + explicit_lengths = transcript_lengths + else: + adata.layers["avg_tx_length"] = transcript_lengths.to_numpy() + + with pytest.raises(ValueError, match="must not be combined"): + DeseqDataSet( + adata=adata, + transcript_lengths=explicit_lengths, + quiet=True, + ) + + +def test_pytximport_lengths_reject_unsynchronized_gene_subsetting(): + adata, _, _, _ = small_pytximport_adata() + subset = adata[:, :2].copy() + + with pytest.raises(ValueError, match="same shape"): + DeseqDataSet(adata=subset, quiet=True) + + +@pytest.mark.parametrize("use_explicit_lengths", [False, True]) +def test_transcript_length_source_precedence(use_explicit_lengths): + adata, _, _, transcript_lengths = small_pytximport_adata() + adata.X = csr_matrix(adata.X) + canonical_lengths = transcript_lengths.to_numpy() + 100.0 + source_lengths = canonical_lengths.copy() + explicit_lengths = transcript_lengths + 200.0 + adata.layers["avg_tx_length"] = canonical_lengths + + dds = DeseqDataSet( + adata=adata, + transcript_lengths=explicit_lengths if use_explicit_lengths else None, + quiet=True, + ) + + expected = explicit_lengths.to_numpy() if use_explicit_lengths else canonical_lengths + np.testing.assert_allclose(dds.layers["avg_tx_length"], expected) + np.testing.assert_allclose(adata.layers["avg_tx_length"], source_lengths) + + +@pytest.mark.parametrize("missing_field", ["counts_from_abundance", "length"]) +def test_incomplete_pytximport_fields_are_not_auto_detected(missing_field): + adata, counts, _, _ = small_pytximport_adata() + adata.X = np.rint(counts.to_numpy()) + if missing_field == "counts_from_abundance": + del adata.uns["counts_from_abundance"] + else: + del adata.obsm["length"] + adata.uns["counts_from_abundance"] = "scaled_tpm" + + dds = DeseqDataSet(adata=adata, quiet=True) + + assert "avg_tx_length" not in dds.layers + + +def test_pytximport_input_anndata_is_not_mutated(): + adata, counts, _, transcript_lengths = small_pytximport_adata() + original_obs = adata.obs.copy() + connectivities = np.eye(adata.n_obs) + adata.uns["neighbors"] = {"connectivities": connectivities} + + dds = DeseqDataSet(adata=adata, quiet=True) + dds.fit_size_factors() + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + picklable_adata = dds.to_picklable_anndata() + + np.testing.assert_array_equal(dds.X, np.rint(counts.to_numpy())) + np.testing.assert_array_equal(adata.X, counts.to_numpy()) + pd.testing.assert_frame_equal(adata.obs, original_obs) + assert dds.var is not adata.var + assert adata.uns["neighbors"]["connectivities"] is connectivities + assert "connectivities" in picklable_adata.obsp + np.testing.assert_array_equal(adata.obsm["length"], transcript_lengths.to_numpy()) + assert "design_matrix" not in adata.obsm + assert not adata.layers diff --git a/tests/test_utils.py b/tests/test_utils.py index 62ff83e2..95ebd5fd 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,9 +3,20 @@ import numpy as np import pytest +from scipy.sparse import bsr_array +from scipy.sparse import bsr_matrix +from scipy.sparse import coo_array +from scipy.sparse import coo_matrix +from scipy.sparse import dia_array +from scipy.sparse import dia_matrix +from scipy.sparse import dok_array +from scipy.sparse import dok_matrix +from scipy.sparse import lil_array +from scipy.sparse import lil_matrix from pydeseq2.utils import load_example_data from pydeseq2.utils import nb_nll +from pydeseq2.utils import test_valid_counts as validate_counts @pytest.mark.parametrize("mu, alpha", [(10, 0.5), (10, 0.1), (3, 0.5), (9, 0.05)]) @@ -53,3 +64,41 @@ def test_rtd_example_data_loading(mocked_function, modality, mocked_dir_flag): dataset="synthetic", debug=False, ) + + +@pytest.mark.parametrize( + "sparse_constructor", + [ + bsr_matrix, + coo_matrix, + dia_matrix, + dok_matrix, + lil_matrix, + bsr_array, + coo_array, + dia_array, + dok_array, + lil_array, + ], +) +def test_valid_sparse_count_formats(sparse_constructor): + validate_counts(sparse_constructor([[1, 0], [0, 2]])) + + +@pytest.mark.parametrize("sparse_constructor", [dia_matrix, dia_array]) +def test_valid_dia_counts_ignore_padding(sparse_constructor): + counts = sparse_constructor( + (np.array([[-1, 2]]), np.array([1])), + shape=(2, 2), + ) + np.testing.assert_array_equal(counts.toarray(), [[0, 2], [0, 0]]) + validate_counts(counts) + + +@pytest.mark.parametrize( + ("value", "message"), + [(np.nan, "NaNs"), (1.5, "integers"), (-1, "non-negative")], +) +def test_invalid_general_sparse_counts(value, message): + with pytest.raises(ValueError, match=message): + validate_counts(dok_matrix([[value]]))