Making progress through alimanfoo's comments

Author: jonbrenas

malariagen_data/anoph/hap_frq.py

@@ -25,9 +25,6 @@ def __init__(
         # to the superclass constructor.
         super().__init__(**kwargs)
 
-        # Set up cache variables.
-        self._cache_annotator = None
-
     @check_types
     @doc(
         summary="""
@@ -46,37 +43,30 @@ def haplotypes_frequencies(
         region: base_params.region,
         cohorts: base_params.cohorts,
         sample_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         min_cohort_size: base_params.min_cohort_size = 10,
         sample_sets: Optional[base_params.sample_sets] = None,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
     ) -> pd.DataFrame:
         # Access sample metadata.
         df_samples = self.sample_metadata(
-            sample_sets=sample_sets, sample_query=sample_query
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
         )
 
         # Build cohort dictionary, maps cohort labels to boolean indexers.
-        coh_dict = locate_cohorts(cohorts=cohorts, data=df_samples)
-
-        # Remove cohorts below minimum cohort size.
-        coh_dict = {
-            coh: loc_coh
-            for coh, loc_coh in coh_dict.items()
-            if np.count_nonzero(loc_coh) >= min_cohort_size
-        }
-
-        # Early check for no cohorts.
-        if len(coh_dict) == 0:
-            raise ValueError(
-                "No cohorts available for the given sample selection parameters and minimum cohort size."
-            )
+        coh_dict = locate_cohorts(
+            cohorts=cohorts, data=df_samples, min_cohort_size=min_cohort_size
+        )
 
         # Access haplotypes.
         ds_haps = self.haplotypes(
             region=region,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             chunks=chunks,
             inline_array=inline_array,
         )
@@ -152,14 +142,17 @@ def haplotypes_frequencies_advanced(
         period_by: frq_params.period_by,
         sample_sets: Optional[base_params.sample_sets] = None,
         sample_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         min_cohort_size: base_params.min_cohort_size = 10,
         ci_method: Optional[frq_params.ci_method] = frq_params.ci_method_default,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
     ) -> xr.Dataset:
         # Load sample metadata.
         df_samples = self.sample_metadata(
-            sample_sets=sample_sets, sample_query=sample_query
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
         )
 
         # Prepare sample metadata for cohort grouping.
@@ -189,6 +182,7 @@ def haplotypes_frequencies_advanced(
             region=region,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             chunks=chunks,
             inline_array=inline_array,
         )

malariagen_data/anoph/sample_metadata.py

@@ -1343,7 +1343,7 @@ def plot_sample_location_geo(
             return fig
 
 
-def locate_cohorts(*, cohorts, data):
+def locate_cohorts(*, cohorts, data, min_cohort_size):
     # Build cohort dictionary where key=cohort_id, value=loc_coh.
     coh_dict = {}
 
@@ -1374,4 +1374,17 @@ def locate_cohorts(*, cohorts, data):
             loc_coh = data[cohorts] == coh
             coh_dict[coh] = loc_coh.values
 
+    # Remove cohorts below minimum cohort size.
+    coh_dict = {
+        coh: loc_coh
+        for coh, loc_coh in coh_dict.items()
+        if np.count_nonzero(loc_coh) >= min_cohort_size
+    }
+
+    # Early check for no cohorts.
+    if len(coh_dict) == 0:
+        raise ValueError(
+            "No cohorts available for the given sample selection parameters and minimum cohort size."
+        )
+
     return coh_dict

malariagen_data/anoph/snp_frq.py

@@ -149,20 +149,9 @@ def snp_allele_frequencies(
         )
 
         # Build cohort dictionary, maps cohort labels to boolean indexers.
-        coh_dict = locate_cohorts(cohorts=cohorts, data=df_samples)
-
-        # Remove cohorts below minimum cohort size.
-        coh_dict = {
-            coh: loc_coh
-            for coh, loc_coh in coh_dict.items()
-            if np.count_nonzero(loc_coh) >= min_cohort_size
-        }
-
-        # Early check for no cohorts.
-        if len(coh_dict) == 0:
-            raise ValueError(
-                "No cohorts available for the given sample selection parameters and minimum cohort size."
-            )
+        coh_dict = locate_cohorts(
+            cohorts=cohorts, data=df_samples, min_cohort_size=min_cohort_size
+        )
 
         # Access SNP data.
         ds_snp = self.snp_calls(

malariagen_data/anopheles.py

@@ -40,7 +40,7 @@
 from .anoph.cnv_data import AnophelesCnvData
 from .anoph.genome_features import AnophelesGenomeFeaturesData
 from .anoph.genome_sequence import AnophelesGenomeSequenceData
-from .anoph.hap_data import AnophelesHapData, hap_params
+from .anoph.hap_data import hap_params
 from .anoph.hap_frq import AnophelesHapFrequencyAnalysis
 from .anoph.igv import AnophelesIgv
 from .anoph.pca import AnophelesPca
@@ -106,7 +106,7 @@ class AnophelesDataResource(
     AnophelesAimData,
     AnophelesSnpData,
     AnophelesCnvData,
-    AnophelesHapData,
+    # AnophelesHapData,
     AnophelesSampleMetadata,
     AnophelesGenomeFeaturesData,
     AnophelesGenomeSequenceData,
@@ -1152,7 +1152,9 @@ def _gene_cnv_frequencies(
         is_called = cn >= 0
 
         debug("set up cohort dict")
-        coh_dict = locate_cohorts(cohorts=cohorts, data=df_samples)
+        coh_dict = locate_cohorts(
+            cohorts=cohorts, data=df_samples, min_cohort_size=min_cohort_size
+        )
 
         debug("compute cohort frequencies")
         freq_cols = dict()