Store dispersion trend coefficients as a dict so the dataset can be written to h5ad

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"

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

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()