diff --git a/src/pydeseq2/dds.py b/src/pydeseq2/dds.py index 614bf12d..cc40458b 100644 --- a/src/pydeseq2/dds.py +++ b/src/pydeseq2/dds.py @@ -447,7 +447,8 @@ def vst_transform(self, counts: np.ndarray | None = None) -> np.ndarray: if "vst_trend_coeffs" not in self.uns: raise RuntimeError("Fit the dispersion curve prior to applying VST.") - a0, a1 = self.uns["vst_trend_coeffs"] + a0 = self.uns["vst_trend_coeffs"]["a0"] + a1 = self.uns["vst_trend_coeffs"]["a1"] return np.log2( ( 1 @@ -1218,10 +1219,12 @@ def _fit_parametric_dispersion_trend(self, vst: bool = False): inplace=True, ) + # Store the coefficients as a plain dict so the object can be written to h5ad + trend_coeffs = {"a0": float(coeffs[0]), "a1": float(coeffs[1])} if vst: - self.uns["vst_trend_coeffs"] = pd.Series(coeffs, index=["a0", "a1"]) + self.uns["vst_trend_coeffs"] = trend_coeffs else: - self.uns["trend_coeffs"] = pd.Series(coeffs, index=["a0", "a1"]) + self.uns["trend_coeffs"] = trend_coeffs self.var["fitted_dispersions"] = np.full(self.n_vars, np.nan) self.uns["disp_function_type"] = "parametric" diff --git a/src/pydeseq2/dispersions.py b/src/pydeseq2/dispersions.py index 9a69da73..4ee78f9b 100644 --- a/src/pydeseq2/dispersions.py +++ b/src/pydeseq2/dispersions.py @@ -17,7 +17,7 @@ def dispersion_trend( normed_mean: float | np.ndarray, - coeffs: pd.Series | np.ndarray, + coeffs: dict[str, float] | pd.Series | np.ndarray, ) -> float | np.ndarray: r"""Return dispersion trend from normalized counts. @@ -34,7 +34,7 @@ def dispersion_trend( ------- Dispersion trend :math:`a_1/ \mu + a_0`. """ - if isinstance(coeffs, pd.Series): + if isinstance(coeffs, (dict, pd.Series)): return coeffs["a0"] + coeffs["a1"] / normed_mean else: return coeffs[0] + coeffs[1] / normed_mean diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py index 23484664..4daf97d3 100644 --- a/tests/test_pydeseq2.py +++ b/tests/test_pydeseq2.py @@ -893,6 +893,22 @@ def test_vst_transform_no_fit(train_counts, train_metadata, test_counts): train_dds.vst_transform(test_counts.to_numpy()) +def test_write_h5ad(counts_df, metadata, tmp_path): + """Test that a fitted DeseqDataSet can be written to and read from h5ad.""" + dds = DeseqDataSet(counts=counts_df, metadata=metadata, design="~condition") + dds.deseq2() + dds.vst(use_design=False) + + path = tmp_path / "dds.h5ad" + dds.write_h5ad(path) + adata = ad.read_h5ad(path) + + for key in ["trend_coeffs", "vst_trend_coeffs"]: + assert adata.uns[key].keys() == {"a0", "a1"} + for coeff in ["a0", "a1"]: + assert adata.uns[key][coeff] == pytest.approx(dds.uns[key][coeff]) + + def assert_res_almost_equal(py_res, r_res, tol=0.02): # check that the same p-values are NaN assert (py_res.pvalue.isna() == r_res.pvalue.isna()).all()