fix: address review feedback - use cache for dataset returns, add missing functions

Author: 31puneet

malariagen_data/anoph/base_params.py

@@ -326,3 +326,12 @@ def _validate_sample_selection_params(
     to select SNPs to be included
     """,
 ]
+
+return_dataset: TypeAlias = Annotated[
+    bool,
+    """
+    If True, return an xarray Dataset containing computed results as
+    additional data variables. If False (default), return the legacy
+    format (numpy array or tuple) for backward compatibility.
+    """,
+]

malariagen_data/anoph/distance.py

@@ -1,5 +1,5 @@
 # Standard library imports.
-from typing import Optional, Tuple
+from typing import Any, Optional, Tuple
 import math
 
 # Third-party library imports.
@@ -86,7 +86,12 @@ def __init__(self, **kwargs):
         summary="""
             Compute pairwise distances between samples using biallelic SNP genotypes.
         """,
-        returns=("dist", "samples", "n_snps_used"),
+        returns="""
+            If `return_dataset` is False (default), return a tuple
+            `(dist, samples, n_snps_used)`. If `return_dataset` is True,
+            return an xarray Dataset with `dist`, `sample_id`, and
+            `n_snps_used` as variables/coordinates.
+        """,
     )
     def biallelic_diplotype_pairwise_distances(
         self,
@@ -108,22 +113,14 @@ def biallelic_diplotype_pairwise_distances(
         random_seed: base_params.random_seed = 42,
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
-    ) -> Tuple[
-        distance_params.dist, distance_params.samples, distance_params.n_snps_used
-    ]:
-        # Change this name if you ever change the behaviour of this function, to
-        # invalidate any previously cached data.
+        return_dataset: base_params.return_dataset = False,
+    ) -> Any:
         name = "biallelic_diplotype_pairwise_distances"
 
-        # Check that either sample_query xor sample_indices are provided.
         base_params._validate_sample_selection_params(
             sample_query=sample_query, sample_indices=sample_indices
         )
 
-        ## Normalize params for consistent hash value.
-
-        # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.
-        # So `sample_query` and `sample_query_options` should not be used beyond this point. (`sample_indices` should be used instead.)
         (
             sample_sets_prepped,
             sample_indices_prepped,
@@ -158,7 +155,6 @@ def biallelic_diplotype_pairwise_distances(
             max_missing_an=max_missing_an,
         )
 
-        # Try to retrieve results from the cache.
         try:
             results = self.results_cache_get(name=name, params=params)
 
@@ -168,10 +164,24 @@ def biallelic_diplotype_pairwise_distances(
             )
             self.results_cache_set(name=name, params=params, results=results)
 
-        # Unpack results.
         dist: np.ndarray = results["dist"]
         samples: np.ndarray = results["samples"]
-        n_snps_used: int = int(results["n_snps"][()])  # ensure scalar
+        n_snps_used: int = int(results["n_snps"][()])
+
+        if return_dataset:
+            import xarray as xr
+            from scipy.spatial.distance import squareform
+            dist_square = squareform(dist)
+            ds = xr.Dataset(
+                data_vars={
+                    "dist": (("sample_x", "sample_y"), dist_square),
+                },
+                coords={
+                    "sample_id": ("sample_x", samples),
+                },
+                attrs={"n_snps_used": n_snps_used},
+            )
+            return ds
 
         return dist, samples, n_snps_used
 

malariagen_data/anoph/hapclust.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Optional, Tuple
+from typing import Any, Optional, Tuple
 
 import allel  # type: ignore
 import numpy as np
@@ -204,11 +204,12 @@ def plot_haplotype_clustering(
         summary="""
             Compute pairwise distances between haplotypes.
         """,
-        returns=dict(
-            dist="Pairwise distance.",
-            phased_samples="Sample identifiers for haplotypes.",
-            n_snps="Number of SNPs used.",
-        ),
+        returns="""
+            If `return_dataset` is False (default), return a tuple
+            `(dist, phased_samples, n_snps)`. If `return_dataset` is True,
+            return an xarray Dataset with `dist`, `sample_id`, and
+            `n_snps` as variables/attributes.
+        """,
     )
     def haplotype_pairwise_distances(
         self,
@@ -222,12 +223,10 @@ def haplotype_pairwise_distances(
         random_seed: base_params.random_seed = 42,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
-    ) -> Tuple[np.ndarray, np.ndarray, int]:
-        # Change this name if you ever change the behaviour of this function, to
-        # invalidate any previously cached data.
+        return_dataset: base_params.return_dataset = False,
+    ) -> Any:
         name = "haplotype_pairwise_distances"
 
-        # Normalize params for consistent hash value.
         sample_sets_prepped = self._prep_sample_sets_param(sample_sets=sample_sets)
         del sample_sets
         sample_query_prepped = self._prep_sample_query_param(sample_query=sample_query)
@@ -245,7 +244,6 @@ def haplotype_pairwise_distances(
             random_seed=random_seed,
         )
 
-        # Try to retrieve results from the cache.
         try:
             results = self.results_cache_get(name=name, params=params)
 
@@ -255,10 +253,24 @@ def haplotype_pairwise_distances(
             )
             self.results_cache_set(name=name, params=params, results=results)
 
-        # Unpack results")
         dist: np.ndarray = results["dist"]
         phased_samples: np.ndarray = results["phased_samples"]
-        n_snps: int = int(results["n_snps"][()])  # ensure scalar
+        n_snps: int = int(results["n_snps"][()])
+
+        if return_dataset:
+            import xarray as xr
+            from scipy.spatial.distance import squareform
+            dist_square = squareform(dist)
+            ds = xr.Dataset(
+                data_vars={
+                    "dist": (("sample_x", "sample_y"), dist_square),
+                },
+                coords={
+                    "sample_id": ("sample_x", phased_samples),
+                },
+                attrs={"n_snps": n_snps},
+            )
+            return ds
 
         return dist, phased_samples, n_snps