From 6b25aa3b8a059571b917e33cb598811e01fbd53c Mon Sep 17 00:00:00 2001
From: _ <4256466+leehart@users.noreply.github.com>
Date: Fri, 27 Sep 2024 10:06:01 +0100
Subject: [PATCH] Add sample_query_options to CNV data functions

---
 malariagen_data/anoph/cnv_data.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/malariagen_data/anoph/cnv_data.py b/malariagen_data/anoph/cnv_data.py
index 337206654..ded1d5690 100644
--- a/malariagen_data/anoph/cnv_data.py
+++ b/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,