Merge pull request #628 from malariagen/GH410_add_sample_query_options

Add sample_query_options

Author: leehart

malariagen_data/anoph/aim_data.py

@@ -121,6 +121,7 @@ def aim_calls(
         aims: aim_params.aims,
         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,
     ) -> xr.Dataset:
         self._require_aim_analysis()
 
@@ -144,7 +145,8 @@ def aim_calls(
         # Handle sample query.
         if sample_query is not None:
             df_samples = self.sample_metadata(sample_sets=sample_sets_prepped)
-            loc_samples = df_samples.eval(sample_query).values
+            sample_query_options = sample_query_options or {}
+            loc_samples = df_samples.eval(sample_query, **sample_query_options).values
             if np.count_nonzero(loc_samples) == 0:
                 raise ValueError(f"No samples found for query {sample_query!r}")
             ds = ds.isel(samples=loc_samples)
@@ -170,6 +172,7 @@ def plot_aim_heatmap(
         aims: aim_params.aims,
         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,
         sort: bool = True,
         row_height: int = 4,
         xgap: float = 0,
@@ -183,6 +186,7 @@ def plot_aim_heatmap(
             aims=aims,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
         ).compute()
         samples = ds["sample_id"].values
         variant_contig = ds["variant_contig"].values

malariagen_data/anoph/base_params.py

@@ -73,6 +73,14 @@
     """,
 ]
 
+sample_query_options: TypeAlias = Annotated[
+    dict,
+    """
+    A dictionary of arguments that will be passed through to pandas query() or
+    eval(), e.g. parser, engine, local_dict, global_dict, resolvers.
+    """,
+]
+
 sample_indices: TypeAlias = Annotated[
     List[int],
     """

malariagen_data/anoph/cnv_data.py

@@ -177,6 +177,7 @@ def cnv_hmm(
         region: base_params.regions,
         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,
         max_coverage_variance: cnv_params.max_coverage_variance = cnv_params.max_coverage_variance_default,
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
@@ -241,7 +242,10 @@ def cnv_hmm(
                 )
 
                 debug("apply the query")
-                loc_query_samples = df_samples_cnv.eval(sample_query).values
+                sample_query_options = sample_query_options or {}
+                loc_query_samples = df_samples_cnv.eval(
+                    sample_query, **sample_query_options
+                ).values
                 if np.count_nonzero(loc_query_samples) == 0:
                     raise ValueError(f"No samples found for query {sample_query!r}")
 
@@ -536,6 +540,7 @@ def cnv_discordant_read_calls(
         contig: base_params.contigs,
         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,
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
     ) -> xr.Dataset:
@@ -588,7 +593,10 @@ def cnv_discordant_read_calls(
             )
 
             debug("apply the query")
-            loc_query_samples = df_samples_cnv.eval(sample_query).values
+            sample_query_options = sample_query_options or {}
+            loc_query_samples = df_samples_cnv.eval(
+                sample_query, **sample_query_options
+            ).values
             if np.count_nonzero(loc_query_samples) == 0:
                 raise ValueError(f"No samples found for query {sample_query!r}")
 
@@ -801,6 +809,7 @@ def plot_cnv_hmm_heatmap_track(
         region: base_params.region,
         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,
         max_coverage_variance: cnv_params.max_coverage_variance = cnv_params.max_coverage_variance_default,
         sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
         width: gplt_params.width = gplt_params.width_default,
@@ -823,6 +832,7 @@ def plot_cnv_hmm_heatmap_track(
             region=region_prepped,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             max_coverage_variance=max_coverage_variance,
         )
 
@@ -942,6 +952,7 @@ def plot_cnv_hmm_heatmap(
         region: base_params.region,
         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,
         max_coverage_variance: cnv_params.max_coverage_variance = cnv_params.max_coverage_variance_default,
         sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
         width: gplt_params.width = gplt_params.width_default,
@@ -960,6 +971,7 @@ def plot_cnv_hmm_heatmap(
             region=region,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             max_coverage_variance=max_coverage_variance,
             sizing_mode=sizing_mode,
             width=width,

malariagen_data/anoph/dipclust.py

@@ -50,6 +50,7 @@ def plot_diplotype_clustering(
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         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,
         cohort_size: Optional[base_params.cohort_size] = None,
         random_seed: base_params.random_seed = 42,
         color: plotly_params.color = None,
@@ -89,14 +90,17 @@ def plot_diplotype_clustering(
 
         # 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,
         )
 
         dist, gt_samples, n_snps_used = self.diplotype_pairwise_distances(
             region=region,
             site_mask=site_mask,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             cohort_size=cohort_size,
             distance_metric=distance_metric,
             random_seed=random_seed,
@@ -196,6 +200,7 @@ def diplotype_pairwise_distances(
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         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,
         site_class: Optional[base_params.site_class] = None,
         cohort_size: Optional[base_params.cohort_size] = None,
         distance_metric: dipclust_params.distance_metric = dipclust_params.distance_metric_default,
@@ -215,6 +220,7 @@ def diplotype_pairwise_distances(
             site_mask=site_mask,
             sample_sets=sample_sets_prepped,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             site_class=site_class,
             cohort_size=cohort_size,
             distance_metric=distance_metric,
@@ -245,6 +251,7 @@ def _diplotype_pairwise_distances(
         site_mask,
         sample_sets,
         sample_query,
+        sample_query_options,
         site_class,
         cohort_size,
         distance_metric,
@@ -261,6 +268,7 @@ def _diplotype_pairwise_distances(
         ds_snps = self.snp_calls(
             region=region,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             site_class=site_class,
@@ -310,6 +318,7 @@ def _dipclust_het_bar_trace(
         dendro_sample_id_order: np.ndarray,
         sample_sets: Optional[base_params.sample_sets],
         sample_query: Optional[base_params.sample_query],
+        sample_query_options: Optional[base_params.sample_query_options],
         site_mask: base_params.site_mask,
         cohort_size: Optional[base_params.cohort_size],
         random_seed: base_params.random_seed,
@@ -320,6 +329,7 @@ def _dipclust_het_bar_trace(
         ds_snps = self.snp_calls(
             region=region,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             cohort_size=cohort_size,
             site_mask=site_mask,
@@ -375,6 +385,7 @@ def _dipclust_cnv_bar_trace(
         dendro_sample_id_order: np.ndarray,
         sample_sets: Optional[base_params.sample_sets],
         sample_query: Optional[base_params.sample_query],
+        sample_query_options: Optional[base_params.sample_query_options],
         max_coverage_variance: Optional[cnv_params.max_coverage_variance],
         colorscale: Optional[plotly_params.color_continuous_scale],
         chunks: base_params.chunks = base_params.native_chunks,
@@ -389,6 +400,7 @@ def _dipclust_cnv_bar_trace(
                 region=cnv_region,
                 sample_sets=sample_sets,
                 sample_query=sample_query,
+                sample_query_options=sample_query_options,
                 max_coverage_variance=max_coverage_variance,
                 chunks=chunks,
                 inline_array=inline_array,
@@ -438,6 +450,7 @@ def _dipclust_snp_trace(
         snp_query: Optional[base_params.snp_query] = AA_CHANGE_QUERY,
         sample_sets: Optional[base_params.sample_sets],
         sample_query: Optional[base_params.sample_query],
+        sample_query_options: Optional[base_params.sample_query_options],
         site_mask: Optional[base_params.site_mask],
         dendro_sample_id_order: np.ndarray,
         snp_filter_min_maf: float,
@@ -450,6 +463,7 @@ def _dipclust_snp_trace(
             transcript=transcript,
             snp_query=snp_query,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             chunks=chunks,
@@ -555,6 +569,7 @@ def plot_diplotype_clustering_advanced(
         site_mask: Optional[base_params.site_mask] = None,
         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,
         random_seed: base_params.random_seed = 42,
         cohort_size: Optional[base_params.cohort_size] = None,
         color: plotly_params.color = None,
@@ -594,6 +609,7 @@ def plot_diplotype_clustering_advanced(
             region=region,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             site_mask=site_mask,
             count_sort=count_sort,
             distance_metric=distance_metric,
@@ -635,6 +651,7 @@ def plot_diplotype_clustering_advanced(
                 dendro_sample_id_order=dendro_sample_id_order,
                 sample_sets=sample_sets,
                 sample_query=sample_query,
+                sample_query_options=sample_query_options,
                 cohort_size=cohort_size,
                 site_mask=site_mask,
                 color_continuous_scale=heterozygosity_colorscale,
@@ -651,6 +668,7 @@ def plot_diplotype_clustering_advanced(
                 dendro_sample_id_order=dendro_sample_id_order,
                 sample_sets=sample_sets,
                 sample_query=sample_query,
+                sample_query_options=sample_query_options,
                 max_coverage_variance=cnv_max_coverage_variance,
                 colorscale=cnv_colorscale,
                 chunks=chunks,
@@ -667,6 +685,7 @@ def plot_diplotype_clustering_advanced(
                 transcript=snp_transcript,
                 sample_sets=sample_sets,
                 sample_query=sample_query,
+                sample_query_options=sample_query_options,
                 snp_query=snp_query,
                 site_mask=site_mask,
                 dendro_sample_id_order=dendro_sample_id_order,
@@ -693,7 +712,7 @@ def plot_diplotype_clustering_advanced(
             height=height,
             row_heights=subplot_heights,
             sample_sets=sample_sets,
-            sample_query=sample_query,
+            sample_query=sample_query,  # Only uses query for title.
             region=region,
             n_snps=n_snps_cluster,
         )

malariagen_data/anoph/fst.py

@@ -415,6 +415,7 @@ def pairwise_average_fst(
         cohorts: base_params.cohorts,
         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,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
         min_cohort_size: Optional[
             base_params.min_cohort_size
@@ -432,6 +433,7 @@ def pairwise_average_fst(
             cohorts,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             cohort_size=cohort_size,
             min_cohort_size=min_cohort_size,
         )