Merge pull request #19 from aertslab/feature-save_patterns_to_disk

Add functionality to save and load patterns to and from disk.

Author: LukasMahieu

docs/api/io.md

@@ -10,4 +10,6 @@ Functions for reading and writing AnnData objects optimized for seqlet analysis.
 
     save_h5ad
     load_h5ad
+    save_patterns
+    load_patterns
 ```

docs/notebooks/02_analysis_tutorial.ipynb

@@ -306,7 +306,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 3,
    "id": "e174f226",
    "metadata": {},
    "outputs": [
@@ -326,6 +326,17 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "2b6ced23",
+   "metadata": {},
+   "source": [
+    "```{tip}\n",
+    "\n",
+    "You can save patterns using {func}`tfmindi.save_patterns` and load them back using {func}`tfmindi.load_patterns`.\n",
+    "```"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 19,

pyproject.toml

@@ -22,6 +22,7 @@ classifiers = [
 ]
 dependencies = [
   "anndata",
+  "h5py>=3.14.0",
   "igraph",
   "lda",
   "logomaker",

src/tfmindi/__init__.py

@@ -22,7 +22,7 @@
     load_motif_collection,
     load_motif_to_dbd,
 )
-from tfmindi.io import load_h5ad, save_h5ad  # noqa: E402
+from tfmindi.io import load_h5ad, load_patterns, save_h5ad, save_patterns  # noqa: E402
 from tfmindi.types import Pattern, Seqlet  # noqa: E402
 
 __all__ = [
@@ -41,6 +41,8 @@
     "load_motif_to_dbd",
     "save_h5ad",
     "load_h5ad",
+    "save_patterns",
+    "load_patterns",
 ]
 
 __version__ = version("tfmindi")

src/tfmindi/io.py

@@ -3,11 +3,15 @@
 from __future__ import annotations
 
 import pickle
+import warnings
 from pathlib import Path
 
+import h5py  # type: ignore
 import numpy as np
-import pandas as pd
-from anndata import AnnData, read_h5ad
+import pandas as pd  # type: ignore
+from anndata import AnnData, read_h5ad  # type: ignore
+
+from tfmindi.types import _PATTERN_SPEC, _SEQLET_SPEC, Pattern, Seqlet
 
 
 def _sanitize_hdf5_keys(data):
@@ -254,3 +258,97 @@ def _convert_numpy_arrays_to_strings_chunked(df, col, chunk_size=1000):
 
     # Convert to categorical to save memory
     df[col] = pd.Series(converted_values, index=series.index).astype(str).astype("category")
+
+
+def _save_seqlet(seqlet: Seqlet, grp: h5py.Group) -> None:
+    """Save seqlet to h5 group."""
+    grp.attrs["version"] = _SEQLET_SPEC
+    for k, v in seqlet.__dict__.items():
+        if v is None:
+            continue
+        grp[k] = v
+
+
+def _save_pattern(pattern: Pattern, grp: h5py.Group) -> None:
+    """Save pattern to h5 group."""
+    grp.attrs["version"] = _PATTERN_SPEC
+    for k, v in pattern.__dict__.items():
+        if k == "seqlets":
+            continue
+        if v is None:
+            continue
+        grp[k] = v
+    seqlets_grp = grp.create_group("seqlets")
+    for i, seqlet in enumerate(pattern.seqlets):
+        seqlet_grp = seqlets_grp.create_group(f"seqlet_{i}")
+        _save_seqlet(seqlet, seqlet_grp)
+
+
+def _read_seqlet(grp: h5py.Group) -> Seqlet:
+    """Load seqlet from h5 group."""
+    kwargs = {}
+    if grp.attrs["version"] != _SEQLET_SPEC:
+        warnings.warn(
+            f"The version of the seqlet on disk ({grp.attrs['version']}) does not match with the pattern version in TF-MInDi ({_PATTERN_SPEC})! Will try to read anyway.",
+            stacklevel=1,
+        )
+    for k in grp.keys():
+        value = grp[k][()]  # type: ignore
+        if isinstance(value, bytes):
+            value = value.decode("utf-8")
+        kwargs[k] = value
+    return Seqlet(**kwargs)
+
+
+def _load_pattern(grp: h5py.Group) -> Pattern:
+    """Load pattern from h5 group."""
+    kwargs = {}
+    if grp.attrs["version"] != _PATTERN_SPEC:
+        warnings.warn(
+            f"The version of the pattern on disk ({grp.attrs['version']}) does not match with the pattern version in TF-MInDi ({_PATTERN_SPEC})! Will try to read anyway.",
+            stacklevel=1,
+        )
+    for k in grp.keys():
+        if k == "seqlets":
+            continue
+        value = grp[k][()]  # type: ignore
+        if isinstance(value, bytes):
+            value = value.decode("utf-8")
+        kwargs[k] = value
+    seqlets: list[Seqlet] = []
+    # Sorted to make sure that the order of the seqlets is the same as when they were saved.
+    for seqlet_key in sorted(grp["seqlets"].keys(), key=lambda x: int(x.split("_")[1])):  # type: ignore
+        seqlets.append(_read_seqlet(grp["seqlets"][seqlet_key]))  # type: ignore
+    kwargs["seqlets"] = seqlets
+    return Pattern(**kwargs)
+
+
+def save_patterns(patterns: dict[str, Pattern], filename: str | Path) -> None:
+    """Save dict of Patterns to disk.
+
+    Paramaters
+    ----------
+    patterns
+        Dict of patterns.
+    filename
+        output filename.
+    """
+    with h5py.File(filename, "w") as h5_handle:
+        for key, pattern in patterns.items():
+            pattern_grp = h5_handle.create_group(f"pattern_{key}")
+            _save_pattern(pattern, pattern_grp)
+
+
+def load_patterns(filename: str | Path) -> dict[str, Pattern]:
+    """Load patterns from disk.
+
+    Parameters
+    ----------
+    filename
+        input filename.
+    """
+    patterns: dict[str, Pattern] = {}
+    with h5py.File(filename, "r") as h5_handle:
+        for pattern_name in h5_handle.keys():
+            patterns[pattern_name.replace("pattern_", "")] = _load_pattern(h5_handle[pattern_name])
+    return patterns

src/tfmindi/types.py

@@ -10,6 +10,12 @@
 _BASE_TO_BIN = {"A": 0, "C": 1, "G": 2, "T": 3}
 _BIN_TO_BASE = {0: "A", 1: "C", 2: "G", 3: "T"}
 
+# Change these version numbers when breaking changes are introduced in Pattern and/or Seqlet.
+# That way incompatibilities can be detected when serializing to or from disk.
+# This version number is saved on disk along with the pattern and seqlet data.
+_PATTERN_SPEC = "1.0"
+_SEQLET_SPEC = "1.0"
+
 
 @dataclass
 class Seqlet:

tests/test_io.py

@@ -204,3 +204,37 @@ def test_save_h5ad_preserves_original(sample_adata_with_arrays):
         assert isinstance(sample_adata_with_arrays.var["motif_pwm"].iloc[0], np.ndarray)
         np.testing.assert_array_equal(sample_adata_with_arrays.obs["seqlet_matrix"].iloc[0], original_obs_array)
         np.testing.assert_array_equal(sample_adata_with_arrays.var["motif_pwm"].iloc[0], original_var_array)
+
+
+def test_save_and_load_pattern(sample_patterns):
+    """Test saving and loading of patterns."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        filepath = Path(tmp_dir) / "test_pattern.hdf5"
+
+        tm.save_patterns(sample_patterns, filepath)
+
+        loaded_patterns = tm.load_patterns(filepath)
+
+        assert set(sample_patterns.keys()) == set(loaded_patterns.keys())
+
+        for k in loaded_patterns.keys():
+            pattern_orig = sample_patterns[k]
+            pattern_loaded = loaded_patterns[k]
+            assert set(pattern_orig.__dict__.keys()) == set(pattern_loaded.__dict__.keys())
+            for attr in pattern_loaded.__dict__.keys():
+                if attr == "seqlets":
+                    continue
+                if isinstance(pattern_orig.__dict__[attr], np.ndarray):
+                    np.testing.assert_array_equal(pattern_orig.__dict__[attr], pattern_loaded.__dict__[attr])
+                else:
+                    assert pattern_orig.__dict__[attr] == pattern_loaded.__dict__[attr]
+            seqlets_orig = pattern_orig.seqlets
+            seqlets_loaded = pattern_loaded.seqlets
+            assert len(seqlets_orig) == len(seqlets_loaded)
+            for s_orig, s_loaded in zip(seqlets_orig, seqlets_loaded, strict=True):
+                assert set(s_orig.__dict__.keys()) == set(s_loaded.__dict__.keys())
+                for attr in s_loaded.__dict__.keys():
+                    if isinstance(s_orig.__dict__[attr], np.ndarray):
+                        np.testing.assert_array_equal(s_orig.__dict__[attr], s_loaded.__dict__[attr])
+                    else:
+                        assert s_orig.__dict__[attr] == s_loaded.__dict__[attr]