Add validate_sample_selection_params to funcs with sample_query, sample_indices. Add sample_query_options to biallelic_snps_to_plink.

Author: leehart

malariagen_data/anoph/distance.py

@@ -115,6 +115,11 @@ def biallelic_diplotype_pairwise_distances(
         # invalidate any previously cached data.
         name = "biallelic_diplotype_pairwise_distances"
 
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         ## Normalize params for consistent hash value.
 
         # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.
@@ -272,6 +277,11 @@ def njt(
         # invalidate any previously cached data.
         name = "njt_v1"
 
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         ## Normalize params for consistent hash value.
 
         # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.
@@ -458,6 +468,11 @@ def plot_njt(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
     ) -> plotly_params.figure:
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         # Only import anjl if needed, as it requires a couple of seconds to compile
         # functions.
         import anjl  # type: ignore

malariagen_data/anoph/pca.py

@@ -82,6 +82,11 @@ def pca(
         # invalidate any previously cached data.
         name = "pca_v5"
 
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         ## Normalize params for consistent hash value.
 
         # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.

malariagen_data/anoph/snp_data.py

@@ -456,7 +456,7 @@ def snp_genotypes(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
     ) -> da.Array:
-        # Additional parameter checks.
+        # Check that either sample_query xor sample_indices are provided.
         base_params.validate_sample_selection_params(
             sample_query=sample_query, sample_indices=sample_indices
         )
@@ -1001,7 +1001,7 @@ def snp_calls(
         max_cohort_size: Optional[base_params.max_cohort_size] = None,
         random_seed: base_params.random_seed = 42,
     ) -> xr.Dataset:
-        # Additional parameter checks.
+        # Check that either sample_query xor sample_indices are provided.
         base_params.validate_sample_selection_params(
             sample_query=sample_query, sample_indices=sample_indices
         )
@@ -1302,6 +1302,11 @@ def snp_allele_counts(
         # to invalidate any previously cached data.
         name = "snp_allele_counts_v2"
 
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         ## Normalize params for consistent hash value.
 
         # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.
@@ -1688,6 +1693,11 @@ def biallelic_snp_calls(
         n_snps: Optional[base_params.n_snps] = None,
         thin_offset: base_params.thin_offset = 0,
     ) -> xr.Dataset:
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         # Perform an allele count.
         ac = self.snp_allele_counts(
             region=region,
@@ -1849,6 +1859,11 @@ def biallelic_diplotypes(
         # invalidate any previously cached data.
         name = "biallelic_diplotypes"
 
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         ## Normalize params for consistent hash value.
 
         # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.

malariagen_data/anoph/to_plink.py

@@ -57,6 +57,7 @@ def biallelic_snps_to_plink(
         thin_offset: base_params.thin_offset = 0,
         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,
         sample_indices: Optional[base_params.sample_indices] = None,
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         min_minor_ac: Optional[
@@ -69,6 +70,11 @@ def biallelic_snps_to_plink(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
     ):
+        # Check that either sample_query xor sample_indices are provided.
+        base_params.validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
         # Define output files
         plink_file_path = f"{output_dir}/{region}.{n_snps}.{min_minor_ac}.{max_missing_an}.{thin_offset}"
 
@@ -84,6 +90,7 @@ def biallelic_snps_to_plink(
             region=region,
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             sample_indices=sample_indices,
             site_mask=site_mask,
             min_minor_ac=min_minor_ac,