Merge branch 'master' into GH1223-seed-random

Author: jonbrenas

malariagen_data/anoph/base_params.py

@@ -10,6 +10,7 @@
     single_contig_param_type,
     single_region_param_type,
     chunks_param_type,
+    Region,
 )
 
 contig: TypeAlias = Annotated[
@@ -47,6 +48,8 @@
     """,
 ]
 
+regions_tuple: TypeAlias = Tuple[Region, ...]
+
 release: TypeAlias = Annotated[
     Union[str, Sequence[str]],
     "Release version identifier.",
@@ -65,6 +68,8 @@
     """,
 ]
 
+sample_sets_tuple: TypeAlias = Tuple[sample_set, ...]
+
 sample_query: TypeAlias = Annotated[
     str,
     """
@@ -94,6 +99,8 @@
     """,
 ]
 
+sample_indices_tuple: TypeAlias = Tuple[int, ...]
+
 sample: TypeAlias = Annotated[
     Union[str, int],
     "Sample identifier or index within sample set.",

malariagen_data/anoph/snp_data.py

@@ -67,12 +67,15 @@ def __init__(
         self._default_site_mask = default_site_mask
 
         # Set up caches.
-        # TODO review type annotations here, maybe can tighten
-        self._cache_snp_sites = None
-        self._cache_snp_genotypes: Dict = dict()
-        self._cache_site_filters: Dict = dict()
-        self._cache_site_annotations = None
-        self._cache_locate_site_class: Dict = dict()
+        self._cache_snp_sites: Optional[zarr.hierarchy.Group] = None
+        self._cache_snp_genotypes: Dict[
+            base_params.sample_set, zarr.hierarchy.Group
+        ] = dict()
+        self._cache_site_filters: Dict[
+            base_params.site_mask, zarr.hierarchy.Group
+        ] = dict()
+        self._cache_site_annotations: Optional[zarr.hierarchy.Group] = None
+        self._cache_locate_site_class: Dict[Tuple[Any, ...], np.ndarray] = dict()
 
         # Create the SNP-calls cache as a per-instance lru_cache wrapping the
         # bound method.  Storing it on the instance (rather than using a
@@ -214,7 +217,7 @@ def _site_filters_for_contig(
         field: base_params.field,
         inline_array: base_params.inline_array,
         chunks: base_params.chunks,
-    ):
+    ) -> da.Array:
         if contig in self.virtual_contigs:
             contigs = self.virtual_contigs[contig]
             arrs = [
@@ -245,7 +248,7 @@ def _site_filters_for_region(
         field: base_params.field,
         inline_array: base_params.inline_array,
         chunks: base_params.chunks,
-    ):
+    ) -> da.Array:
         d = self._site_filters_for_contig(
             contig=region.contig,
             mask=mask,
@@ -579,7 +582,7 @@ def _snp_variants_for_contig(
         contig: base_params.contig,
         inline_array: base_params.inline_array,
         chunks: base_params.chunks,
-    ):
+    ) -> xr.Dataset:
         if contig in self.virtual_contigs:
             contigs = self.virtual_contigs[contig]
             datasets = []
@@ -652,7 +655,7 @@ def snp_variants(
         site_mask: Optional[base_params.site_mask] = None,
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
-    ):
+    ) -> xr.Dataset:
         # Normalise parameters.
         regions: List[Region] = _parse_multi_region(self, region)
         del region
@@ -789,7 +792,7 @@ def _locate_site_class(
         site_class: base_params.site_class,
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
-    ):
+    ) -> np.ndarray:
         # Cache these data in memory to avoid repeated computation.
         cache_key = (region, site_mask, site_class)
 
@@ -1082,11 +1085,11 @@ def snp_calls(
         del site_mask
 
         # Convert lists to tuples to avoid CacheMiss "TypeError: unhashable type: 'list'".
-        prepared_regions_tuple: Tuple[Region, ...] = tuple(prepared_regions)
-        prepared_sample_sets_tuple: Optional[Tuple[str, ...]] = (
+        prepared_regions_tuple: base_params.regions_tuple = tuple(prepared_regions)
+        prepared_sample_sets_tuple: Optional[base_params.sample_sets_tuple] = (
             tuple(prepared_sample_sets) if prepared_sample_sets is not None else None
         )
-        prepared_sample_indices_tuple: Optional[Tuple[int, ...]] = (
+        prepared_sample_indices_tuple: Optional[base_params.sample_indices_tuple] = (
             tuple(prepared_sample_indices)
             if prepared_sample_indices is not None
             else None
@@ -1110,18 +1113,19 @@ def snp_calls(
     def _raw_snp_calls(
         self,
         *,
-        regions: Tuple[Region, ...],
-        sample_sets,
-        site_mask,
-        site_class,
-        inline_array,
-        chunks,
-    ):
+        regions: base_params.regions_tuple,
+        sample_sets: Optional[base_params.sample_sets_tuple],
+        site_mask: Optional[base_params.site_mask],
+        site_class: Optional[base_params.site_class],
+        inline_array: base_params.inline_array,
+        chunks: base_params.chunks,
+    ) -> xr.Dataset:
         # Access SNP calls and concatenate multiple sample sets and/or regions.
         with self._spinner("Access SNP calls"):
             lx = []
             for r in regions:
                 ly = []
+                assert sample_sets is not None
                 for s in sample_sets:
                     y = self._snp_calls_for_contig(
                         contig=r.contig,
@@ -1179,18 +1183,18 @@ def _raw_snp_calls(
     def _snp_calls(
         self,
         *,
-        regions: Tuple[Region, ...],
-        sample_sets,
-        sample_indices,
-        site_mask,
-        site_class,
-        cohort_size,
-        min_cohort_size,
-        max_cohort_size,
-        random_seed,
-        inline_array,
-        chunks,
-    ):
+        regions: base_params.regions_tuple,
+        sample_sets: Optional[base_params.sample_sets_tuple],
+        sample_indices: Optional[base_params.sample_indices_tuple],
+        site_mask: Optional[base_params.site_mask],
+        site_class: Optional[base_params.site_class],
+        cohort_size: Optional[base_params.cohort_size],
+        min_cohort_size: Optional[base_params.min_cohort_size],
+        max_cohort_size: Optional[base_params.max_cohort_size],
+        random_seed: base_params.random_seed,
+        inline_array: base_params.inline_array,
+        chunks: base_params.chunks,
+    ) -> xr.Dataset:
         ## Get SNP calls and concatenate multiple sample sets and/or regions.
 
         # Note: sample_sets should be "prepared" before being passed to this private function.
@@ -1305,23 +1309,25 @@ def _results_cache_add_analysis_params(self, params: dict):
     def _snp_allele_counts(
         self,
         *,
-        region,
-        sample_sets,
-        sample_indices,
-        site_mask,
-        site_class,
-        cohort_size,
-        min_cohort_size,
-        max_cohort_size,
-        random_seed,
-        inline_array,
-        chunks,
-    ):
+        region: Union[dict, List[dict]],
+        sample_sets: Optional[base_params.sample_sets_tuple],
+        sample_indices: Optional[base_params.sample_indices_tuple],
+        site_mask: Optional[base_params.site_mask],
+        site_class: Optional[base_params.site_class],
+        cohort_size: Optional[base_params.cohort_size],
+        min_cohort_size: Optional[base_params.min_cohort_size],
+        max_cohort_size: Optional[base_params.max_cohort_size],
+        random_seed: base_params.random_seed,
+        inline_array: base_params.inline_array,
+        chunks: base_params.chunks,
+    ) -> Dict[str, np.ndarray]:
         # Access SNP calls.
+        # N.B., snp_calls is a public API with @_check_types, which expects
+        # List[int] for sample_indices, not a tuple. Convert back here.
         ds_snps = self.snp_calls(
             region=region,
-            sample_sets=sample_sets,
-            sample_indices=sample_indices,
+            sample_sets=list(sample_sets) if sample_sets is not None else None,
+            sample_indices=list(sample_indices) if sample_indices is not None else None,
             site_mask=site_mask,
             site_class=site_class,
             cohort_size=cohort_size,
@@ -1402,6 +1408,15 @@ def snp_allele_counts(
             sample_query_options=sample_query_options,
             sample_indices=sample_indices,
         )
+        # Convert lists to tuples to avoid CacheMiss "TypeError: unhashable type: 'list'".
+        sample_sets_prepped_tuple: Optional[base_params.sample_sets_tuple] = (
+            tuple(sample_sets_prepped) if sample_sets_prepped is not None else None
+        )
+        sample_indices_prepped_tuple: Optional[base_params.sample_indices_tuple] = (
+            tuple(sample_indices_prepped)
+            if sample_indices_prepped is not None
+            else None
+        )
         del sample_sets
         del sample_query
         del sample_query_options
@@ -1412,8 +1427,8 @@ def snp_allele_counts(
         del site_mask
         params = dict(
             region=region_prepped,
-            sample_sets=sample_sets_prepped,
-            sample_indices=sample_indices_prepped,
+            sample_sets=sample_sets_prepped_tuple,
+            sample_indices=sample_indices_prepped_tuple,
             site_mask=site_mask_prepped,
             site_class=site_class,
             cohort_size=cohort_size,
@@ -1427,7 +1442,17 @@ def snp_allele_counts(
 
         except CacheMiss:
             results = self._snp_allele_counts(
-                **params, inline_array=inline_array, chunks=chunks
+                inline_array=inline_array,
+                chunks=chunks,
+                region=region_prepped,
+                sample_sets=sample_sets_prepped_tuple,
+                sample_indices=sample_indices_prepped_tuple,
+                site_mask=site_mask_prepped,
+                site_class=site_class,
+                cohort_size=cohort_size,
+                min_cohort_size=min_cohort_size,
+                max_cohort_size=max_cohort_size,
+                random_seed=random_seed,
             )
             self.results_cache_set(name=name, params=params, results=results)
 
@@ -1964,6 +1989,16 @@ def biallelic_diplotypes(
         prepared_region = self._prep_region_cache_param(region=region)
         prepared_site_mask = self._prep_optional_site_mask_param(site_mask=site_mask)
 
+        # Convert lists to tuples to avoid CacheMiss "TypeError: unhashable type: 'list'".
+        prepared_sample_sets_tuple: Optional[base_params.sample_sets_tuple] = (
+            tuple(prepared_sample_sets) if prepared_sample_sets is not None else None
+        )
+        prepared_sample_indices_tuple: Optional[base_params.sample_indices_tuple] = (
+            tuple(prepared_sample_indices)
+            if prepared_sample_indices is not None
+            else None
+        )
+
         # Delete original parameters to prevent accidental use.
         del sample_sets
         del sample_query
@@ -1976,8 +2011,8 @@ def biallelic_diplotypes(
             region=prepared_region,
             n_snps=n_snps,
             thin_offset=thin_offset,
-            sample_sets=prepared_sample_sets,
-            sample_indices=prepared_sample_indices,
+            sample_sets=prepared_sample_sets_tuple,
+            sample_indices=prepared_sample_indices_tuple,
             site_mask=prepared_site_mask,
             site_class=site_class,
             cohort_size=cohort_size,
@@ -1994,7 +2029,21 @@ def biallelic_diplotypes(
 
         except CacheMiss:
             results = self._biallelic_diplotypes(
-                inline_array=inline_array, chunks=chunks, **params
+                inline_array=inline_array,
+                chunks=chunks,
+                region=prepared_region,
+                sample_sets=prepared_sample_sets_tuple,
+                sample_indices=prepared_sample_indices_tuple,
+                site_mask=prepared_site_mask,
+                site_class=site_class,
+                cohort_size=cohort_size,
+                min_cohort_size=min_cohort_size,
+                max_cohort_size=max_cohort_size,
+                random_seed=random_seed,
+                n_snps=n_snps,
+                thin_offset=thin_offset,
+                min_minor_ac=min_minor_ac,
+                max_missing_an=max_missing_an,
             )
             self.results_cache_set(name=name, params=params, results=results)
 
@@ -2007,29 +2056,31 @@ def biallelic_diplotypes(
     def _biallelic_diplotypes(
         self,
         *,
-        region,
-        sample_sets,
-        sample_indices,
-        site_mask,
-        site_class,
-        cohort_size,
-        min_cohort_size,
-        max_cohort_size,
-        random_seed,
-        max_missing_an,
-        min_minor_ac,
-        n_snps,
-        thin_offset,
-        inline_array,
-        chunks,
-    ):
+        region: Union[dict, List[dict]],
+        sample_sets: Optional[base_params.sample_sets_tuple],
+        sample_indices: Optional[base_params.sample_indices_tuple],
+        site_mask: Optional[base_params.site_mask],
+        site_class: Optional[base_params.site_class],
+        cohort_size: Optional[base_params.cohort_size],
+        min_cohort_size: Optional[base_params.min_cohort_size],
+        max_cohort_size: Optional[base_params.max_cohort_size],
+        random_seed: base_params.random_seed,
+        max_missing_an: Optional[base_params.max_missing_an],
+        min_minor_ac: Optional[base_params.min_minor_ac],
+        n_snps: Optional[base_params.n_snps],
+        thin_offset: base_params.thin_offset,
+        inline_array: base_params.inline_array,
+        chunks: base_params.chunks,
+    ) -> Dict[str, np.ndarray]:
         # Note: this function uses sample_indices and should not expect a sample_query.
 
         # Access biallelic SNPs.
+        # N.B., biallelic_snp_calls is a public API with @_check_types, which
+        # expects List[int] for sample_indices, not a tuple. Convert back here.
         ds = self.biallelic_snp_calls(
             region=region,
-            sample_sets=sample_sets,
-            sample_indices=sample_indices,
+            sample_sets=list(sample_sets) if sample_sets is not None else None,
+            sample_indices=list(sample_indices) if sample_indices is not None else None,
             site_mask=site_mask,
             site_class=site_class,
             cohort_size=cohort_size,