WIP: dev support for surveillance_use_only, unrestricted_use_only params

Author: leehart

malariagen_data/anoph/aim_data.py

@@ -171,29 +171,12 @@ def aim_calls(
             # If there are no sample query options, then default to an empty dict.
             sample_query_options = sample_query_options or {}
 
-            # Determine which samples match the sample query.
-            loc_samples = df_samples.eval(prepared_sample_query, **sample_query_options)
-
-            # Raise an error if no samples match the sample query.
-            if not loc_samples.any():
-                raise ValueError(
-                    f"No samples found for query {prepared_sample_query!r}"
-                )
-
-            # Get the relevant sample ids from the sample metadata DataFrame, using the boolean mask.
-            relevant_sample_ids = df_samples.loc[loc_samples, "sample_id"].values
-
-            # Get all the sample ids from the unfiltered AIM calls Dataset.
-            ds_sample_ids = ds.coords["sample_id"].values
-
-            # Get the indices of samples in the AIM calls Dataset that match the relevant sample ids.
-            # Note: we use `[0]` to get the first element of the tuple returned by `np.where`.
-            relevant_sample_indices = np.where(
-                np.isin(ds_sample_ids, relevant_sample_ids)
-            )[0]
-
-            # Select only the relevant samples from the AIM calls Dataset.
-            ds = ds.isel(samples=relevant_sample_indices)
+            ds = self._filter_sample_dataset(
+                ds=ds,
+                df_samples=df_samples,
+                sample_query=prepared_sample_query,
+                sample_query_options=sample_query_options,
+            )
 
         return ds
 

malariagen_data/anoph/base.py

@@ -28,6 +28,7 @@
 from tqdm.auto import tqdm as tqdm_auto  # type: ignore
 from tqdm.dask import TqdmCallback  # type: ignore
 from yaspin import yaspin  # type: ignore
+import xarray as xr
 
 from ..util import (
     CacheMiss,
@@ -933,6 +934,45 @@ def _prep_sample_query_param(
 
         return prepped_sample_query
 
+    def _filter_sample_dataset(
+        self,
+        *,
+        ds: xr.Dataset,
+        df_samples: pd.DataFrame,
+        sample_query: str,
+        sample_query_options: dict,
+    ) -> xr.Dataset:
+        """Filters the given Dataset using the given DataFrame and query."""
+
+        # Note: "prepare" the params before calling this function.
+
+        # Determine which samples match the sample query.
+        if sample_query != "":
+            loc_samples = df_samples.eval(sample_query, **sample_query_options)
+        else:
+            loc_samples = pd.Series(True, index=df_samples.index)
+
+        # Raise an error if no samples match the sample query.
+        if not loc_samples.any():
+            raise ValueError(f"No samples found for query {sample_query!r}")
+
+        # Get the relevant sample ids from the sample metadata DataFrame, using the boolean mask.
+        relevant_sample_ids = df_samples.loc[loc_samples, "sample_id"].values
+
+        # Get all the sample ids from the unfiltered Dataset.
+        ds_sample_ids = ds.coords["sample_id"].values
+
+        # Get the indices of samples in the Dataset that match the relevant sample ids.
+        # Note: we use `[0]` to get the first element of the tuple returned by `np.where`.
+        relevant_sample_indices = np.where(np.isin(ds_sample_ids, relevant_sample_ids))[
+            0
+        ]
+
+        # Select only the relevant samples from the Dataset.
+        ds = ds.isel(samples=relevant_sample_indices)
+
+        return ds
+
     def _results_cache_add_analysis_params(self, params: dict):
         # Expect sub-classes will override to add any analysis parameters.
         pass

malariagen_data/anoph/cnv_data.py

@@ -256,32 +256,13 @@ def cnv_hmm(
                 # If there are no sample query options, then default to an empty dict.
                 sample_query_options = sample_query_options or {}
 
-                # Determine which samples match the sample query.
-                loc_samples = df_samples.eval(
-                    prepared_sample_query, **sample_query_options
+                ds = self._filter_sample_dataset(
+                    ds=ds,
+                    df_samples=df_samples,
+                    sample_query=prepared_sample_query,
+                    sample_query_options=sample_query_options,
                 )
 
-                # Raise an error if no samples match the sample query.
-                if not loc_samples.any():
-                    raise ValueError(
-                        f"No samples found for query {prepared_sample_query!r}"
-                    )
-
-                # Get the relevant sample ids from the sample metadata DataFrame, using the boolean mask.
-                relevant_sample_ids = df_samples.loc[loc_samples, "sample_id"].values
-
-                # Get all the sample ids from the unfiltered CNV HMM Dataset.
-                ds_sample_ids = ds.coords["sample_id"].values
-
-                # Get the indices of samples in the CNV HMM Dataset that match the relevant sample ids.
-                # Note: we use `[0]` to get the first element of the tuple returned by `np.where`.
-                relevant_sample_indices = np.where(
-                    np.isin(ds_sample_ids, relevant_sample_ids)
-                )[0]
-
-                # Select only the relevant samples from the CNV HMM Dataset.
-                ds = ds.isel(samples=relevant_sample_indices)
-
             debug("handle coverage variance filter")
             if max_coverage_variance is not None:
                 cov_var = ds["sample_coverage_variance"].values
@@ -476,30 +457,15 @@ def cnv_coverage_calls(
         # Get the relevant sample metadata.
         df_samples = self.sample_metadata(sample_sets=prepared_sample_set)
 
-        # Determine which samples match the sample query.
-        if prepared_sample_query != "":
-            loc_samples = df_samples.eval(prepared_sample_query)
-        else:
-            loc_samples = pd.Series(True, index=df_samples.index)
-
-        # Raise an error if no samples match the sample query.
-        if not loc_samples.any():
-            raise ValueError(f"No samples found for query {prepared_sample_query!r}")
-
-        # Get the relevant sample ids from the sample metadata DataFrame, using the boolean mask.
-        relevant_sample_ids = df_samples.loc[loc_samples, "sample_id"].values
+        # If there is no sample query, then default to an empty str.
+        prepared_sample_query = prepared_sample_query or ""
 
-        # Get all the sample ids from the unfiltered CNV coverage calls Dataset.
-        ds_sample_ids = ds.coords["sample_id"].values
-
-        # Get the indices of samples in the CNV coverage calls Dataset that match the relevant sample ids.
-        # Note: we use `[0]` to get the first element of the tuple returned by `np.where`.
-        relevant_sample_indices = np.where(np.isin(ds_sample_ids, relevant_sample_ids))[
-            0
-        ]
-
-        # Select only the relevant samples from the CNV coverage calls Dataset.
-        ds = ds.isel(samples=relevant_sample_indices)
+        ds = self._filter_sample_dataset(
+            ds=ds,
+            df_samples=df_samples,
+            sample_query=prepared_sample_query,
+            sample_query_options={},
+        )
 
         return ds
 
@@ -690,29 +656,12 @@ def cnv_discordant_read_calls(
             # If there are no sample query options, then default to an empty dict.
             sample_query_options = sample_query_options or {}
 
-            # Determine which samples match the sample query.
-            loc_samples = df_samples.eval(prepared_sample_query, **sample_query_options)
-
-            # Raise an error if no samples match the sample query.
-            if not loc_samples.any():
-                raise ValueError(
-                    f"No samples found for query {prepared_sample_query!r}"
-                )
-
-            # Get the relevant sample ids from the sample metadata DataFrame, using the boolean mask.
-            relevant_sample_ids = df_samples.loc[loc_samples, "sample_id"].values
-
-            # Get all the sample ids from the unfiltered CNV discordant reads Dataset.
-            ds_sample_ids = ds.coords["sample_id"].values
-
-            # Get the indices of samples in the CNV discordant reads Dataset that match the relevant sample ids.
-            # Note: we use `[0]` to get the first element of the tuple returned by `np.where`.
-            relevant_sample_indices = np.where(
-                np.isin(ds_sample_ids, relevant_sample_ids)
-            )[0]
-
-            # Select only the relevant samples from the CNV discordant reads Dataset.
-            ds = ds.isel(samples=relevant_sample_indices)
+            ds = self._filter_sample_dataset(
+                ds=ds,
+                df_samples=df_samples,
+                sample_query=prepared_sample_query,
+                sample_query_options=sample_query_options,
+            )
 
         return ds
 

malariagen_data/anoph/distance.py

@@ -115,7 +115,10 @@ def biallelic_diplotype_pairwise_distances(
         # invalidate any previously cached data.
         name = "biallelic_diplotype_pairwise_distances"
 
-        # Normalize params for consistent hash value.
+        ## 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,
@@ -269,7 +272,10 @@ def njt(
         # invalidate any previously cached data.
         name = "njt_v1"
 
-        # Normalize params for consistent hash value.
+        ## 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,

malariagen_data/anoph/pca.py

@@ -80,27 +80,39 @@ def pca(
     ) -> Tuple[pca_params.df_pca, pca_params.evr]:
         # Change this name if you ever change the behaviour of this function, to
         # invalidate any previously cached data.
-        name = "pca_v4"
+        name = "pca_v5"
 
-        # Normalize params for consistent hash value.
+        ## 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,
+            prepared_sample_sets,
+            prepared_sample_indices,
         ) = self._prep_sample_selection_cache_params(
             sample_sets=sample_sets,
             sample_query=sample_query,
             sample_query_options=sample_query_options,
             sample_indices=sample_indices,
         )
-        region_prepped = self._prep_region_cache_param(region=region)
-        site_mask_prepped = self._prep_optional_site_mask_param(site_mask=site_mask)
+        prepared_region = self._prep_region_cache_param(region=region)
+        prepared_site_mask = self._prep_optional_site_mask_param(site_mask=site_mask)
+
+        # Delete original parameters to prevent accidental use.
+        del sample_sets
+        del sample_indices
+        del sample_query
+        del sample_query_options
+        del region
+        del site_mask
+
         params = dict(
-            region=region_prepped,
+            region=prepared_region,
             n_snps=n_snps,
             thin_offset=thin_offset,
-            sample_sets=sample_sets_prepped,
-            sample_indices=sample_indices_prepped,
-            site_mask=site_mask_prepped,
+            sample_sets=prepared_sample_sets,
+            sample_indices=prepared_sample_indices,
+            site_mask=prepared_site_mask,
             site_class=site_class,
             min_minor_ac=min_minor_ac,
             max_missing_an=max_missing_an,
@@ -127,22 +139,18 @@ def pca(
         samples = results["samples"]
         loc_keep_fit = results["loc_keep_fit"]
 
-        # Load sample metadata.
-        df_samples = self.sample_metadata(
-            sample_sets=sample_sets,
-        )
+        # Create a new DataFrame containing the PCA coords data.
+        df_pca = pd.DataFrame(coords, index=samples)
 
-        # Ensure aligned with genotype data.
-        df_samples = df_samples.set_index("sample_id").loc[samples].reset_index()
+        # Name the DataFrame's columns PC1, PC2, etc.
+        df_pca.columns = pd.Index([f"PC{i+1}" for i in range(coords.shape[1])])
 
-        # Combine coords and sample metadata.
-        df_coords = pd.DataFrame(
-            {f"PC{i + 1}": coords[:, i] for i in range(coords.shape[1])}
-        )
-        df_pca = df_samples.join(df_coords, how="inner")
-        # Add a column for which samples were included in fitting.
+        # Add a column to indicate which samples were included in fitting.
         df_pca["pca_fit"] = loc_keep_fit
 
+        # Name the index.
+        df_pca.index.name = "sample_id"
+
         return df_pca, evr
 
     def _pca(

malariagen_data/anoph/sample_metadata.py

@@ -1054,20 +1054,26 @@ def _prep_sample_selection_cache_params(
         sample_indices: Optional[base_params.sample_indices],
     ) -> Tuple[List[str], Optional[List[int]]]:
         # Normalise sample sets.
-        sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
-        sample_query = self._prep_sample_query_param(sample_query=sample_query)
+        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
+        prepared_sample_query = self._prep_sample_query_param(sample_query=sample_query)
 
-        if sample_query is not None:
+        # Delete original parameters to prevent accidental use.
+        del sample_sets
+        del sample_query
+
+        if prepared_sample_query is not None:
             # Resolve query to a list of integers for more cache hits - we
             # do this because there are different ways to write the same pandas
             # query, and so it's better to evaluate the query and use a list of
             # integer indices instead.
-            df_samples = self.sample_metadata(sample_sets=sample_sets)
+            df_samples = self.sample_metadata(sample_sets=prepared_sample_sets)
             sample_query_options = sample_query_options or {}
-            loc_samples = df_samples.eval(sample_query, **sample_query_options).values
+            loc_samples = df_samples.eval(
+                prepared_sample_query, **sample_query_options
+            ).values
             sample_indices = np.nonzero(loc_samples)[0].tolist()
 
-        return sample_sets, sample_indices
+        return prepared_sample_sets, sample_indices
 
     def _results_cache_add_analysis_params(self, params: dict):
         super()._results_cache_add_analysis_params(params)