Merge pull request #645 from malariagen/GH635_add_sample_query_options

Add sample_query_options to more functions

Author: leehart

malariagen_data/ag3.py

@@ -378,6 +378,7 @@ def karyotype(
         inversion: inversion_param,
         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,
     ) -> pd.DataFrame:
         # load tag snp data
         df_tagsnps = self.load_inversion_tags(inversion=inversion)
@@ -390,7 +391,10 @@ def karyotype(
         region = f"{contig}:{start}-{end}"
 
         ds_snps = self.snp_calls(
-            region=region, sample_sets=sample_sets, sample_query=sample_query
+            region=region,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
         )
 
         with self._spinner("Inferring karyotype from tag SNPs"):

malariagen_data/anoph/fst.py

@@ -28,11 +28,13 @@ def __init__(
 
     def _fst_gwss(
         self,
+        *,
         contig,
         window_size,
         sample_sets,
         cohort1_query,
         cohort2_query,
+        sample_query_options,
         site_mask,
         cohort_size,
         min_cohort_size,
@@ -46,6 +48,7 @@ def _fst_gwss(
         ac1 = self.snp_allele_counts(
             region=contig,
             sample_query=cohort1_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             cohort_size=cohort_size,
@@ -58,6 +61,7 @@ def _fst_gwss(
         ac2 = self.snp_allele_counts(
             region=contig,
             sample_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             cohort_size=cohort_size,
@@ -104,6 +108,7 @@ def fst_gwss(
         window_size: fst_params.window_size,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
@@ -127,6 +132,7 @@ def fst_gwss(
             window_size=window_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=self._prep_sample_sets_param(sample_sets=sample_sets),
             site_mask=self._prep_optional_site_mask_param(site_mask=site_mask),
             cohort_size=cohort_size,
@@ -161,6 +167,7 @@ def plot_fst_gwss_track(
         window_size: fst_params.window_size,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
@@ -189,6 +196,7 @@ def plot_fst_gwss_track(
             max_cohort_size=max_cohort_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             random_seed=random_seed,
@@ -265,6 +273,7 @@ def plot_fst_gwss(
         window_size: fst_params.window_size,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
@@ -290,6 +299,7 @@ def plot_fst_gwss(
             window_size=window_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             cohort_size=cohort_size,
@@ -348,6 +358,7 @@ def average_fst(
         region: base_params.region,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
         min_cohort_size: Optional[
@@ -366,6 +377,7 @@ def average_fst(
             region=region,
             sample_sets=sample_sets,
             sample_query=cohort1_query,
+            sample_query_options=sample_query_options,
             cohort_size=cohort_size,
             site_mask=site_mask,
             site_class=site_class,
@@ -377,6 +389,7 @@ def average_fst(
             region=region,
             sample_sets=sample_sets,
             sample_query=cohort2_query,
+            sample_query_options=sample_query_options,
             cohort_size=cohort_size,
             site_mask=site_mask,
             site_class=site_class,

malariagen_data/anoph/h1x.py

@@ -26,12 +26,14 @@ def __init__(
 
     def _h1x_gwss(
         self,
+        *,
         contig,
         analysis,
         window_size,
         sample_sets,
         cohort1_query,
         cohort2_query,
+        sample_query_options,
         cohort_size,
         min_cohort_size,
         max_cohort_size,
@@ -44,6 +46,7 @@ def _h1x_gwss(
             region=contig,
             analysis=analysis,
             sample_query=cohort1_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             cohort_size=cohort_size,
             min_cohort_size=min_cohort_size,
@@ -56,6 +59,7 @@ def _h1x_gwss(
             region=contig,
             analysis=analysis,
             sample_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             cohort_size=cohort_size,
             min_cohort_size=min_cohort_size,
@@ -102,6 +106,7 @@ def h1x_gwss(
         window_size: h12_params.window_size,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         analysis: hap_params.analysis = base_params.DEFAULT,
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort_size: Optional[base_params.cohort_size] = h12_params.cohort_size_default,
@@ -128,6 +133,7 @@ def h1x_gwss(
             # are different in the haplotype data.
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=self._prep_sample_sets_param(sample_sets=sample_sets),
             cohort_size=cohort_size,
             min_cohort_size=min_cohort_size,
@@ -160,6 +166,7 @@ def plot_h1x_gwss_track(
         window_size: h12_params.window_size,
         cohort1_query: base_params.cohort1_query,
         cohort2_query: base_params.cohort2_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         analysis: hap_params.analysis = base_params.DEFAULT,
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort_size: Optional[base_params.cohort_size] = h12_params.cohort_size_default,
@@ -190,6 +197,7 @@ def plot_h1x_gwss_track(
             max_cohort_size=max_cohort_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             random_seed=random_seed,
             chunks=chunks,
@@ -266,6 +274,7 @@ def plot_h1x_gwss(
         window_size: h12_params.window_size,
         cohort1_query: base_params.cohort1_query,
         cohort2_query: base_params.cohort2_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         analysis: hap_params.analysis = base_params.DEFAULT,
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort_size: Optional[base_params.cohort_size] = h12_params.cohort_size_default,
@@ -293,6 +302,7 @@ def plot_h1x_gwss(
             window_size=window_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             cohort_size=cohort_size,
             min_cohort_size=min_cohort_size,

malariagen_data/anopheles.py

@@ -2253,6 +2253,7 @@ def plot_xpehh_gwss(
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort1_query: Optional[base_params.sample_query] = None,
         cohort2_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         window_size: xpehh_params.window_size = xpehh_params.window_size_default,
         percentiles: xpehh_params.percentiles = xpehh_params.percentiles_default,
         filter_min_maf: xpehh_params.filter_min_maf = xpehh_params.filter_min_maf_default,
@@ -2287,6 +2288,7 @@ def plot_xpehh_gwss(
             sample_sets=sample_sets,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             window_size=window_size,
             percentiles=percentiles,
             palette=palette,
@@ -2457,6 +2459,7 @@ def xpehh_gwss(
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort1_query: Optional[base_params.sample_query] = None,
         cohort2_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         window_size: xpehh_params.window_size = xpehh_params.window_size_default,
         percentiles: xpehh_params.percentiles = xpehh_params.percentiles_default,
         filter_min_maf: xpehh_params.filter_min_maf = xpehh_params.filter_min_maf_default,
@@ -2498,6 +2501,7 @@ def xpehh_gwss(
             # are different in the haplotype data.
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             min_cohort_size=min_cohort_size,
             max_cohort_size=max_cohort_size,
             random_seed=random_seed,
@@ -2525,6 +2529,7 @@ def _xpehh_gwss(
         sample_sets,
         cohort1_query,
         cohort2_query,
+        sample_query_options,
         window_size,
         percentiles,
         filter_min_maf,
@@ -2544,6 +2549,7 @@ def _xpehh_gwss(
             region=contig,
             analysis=analysis,
             sample_query=cohort1_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             min_cohort_size=min_cohort_size,
             max_cohort_size=max_cohort_size,
@@ -2556,6 +2562,7 @@ def _xpehh_gwss(
             region=contig,
             analysis=analysis,
             sample_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             min_cohort_size=min_cohort_size,
             max_cohort_size=max_cohort_size,
@@ -2628,6 +2635,7 @@ def plot_xpehh_gwss_track(
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort1_query: Optional[base_params.sample_query] = None,
         cohort2_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         window_size: xpehh_params.window_size = xpehh_params.window_size_default,
         percentiles: xpehh_params.percentiles = xpehh_params.percentiles_default,
         filter_min_maf: xpehh_params.filter_min_maf = xpehh_params.filter_min_maf_default,
@@ -2672,6 +2680,7 @@ def plot_xpehh_gwss_track(
             max_cohort_size=max_cohort_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             random_seed=random_seed,
             chunks=chunks,