Merge branch 'master' into fix/mjn-duplicate-nodes

Author: jonbrenas

malariagen_data/ag3.py

@@ -74,6 +74,14 @@ def _setup_aim_palettes():
     "gcx4": TAXON_PALETTE[10],
     "unassigned": "black",
 }
+# Colors for aim_species column, matching the AIM palettes.
+AIM_SPECIES_COLORS = {
+    "gambiae": AIM_PALETTES["gamb_vs_colu"][1],
+    "coluzzii": AIM_PALETTES["gamb_vs_colu"][3],
+    "arabiensis": AIM_PALETTES["gambcolu_vs_arab"][3],
+    "gambcolu": AIM_PALETTES["gambcolu_vs_arab"][1],
+    "unassigned": "black",
+}
 
 # Note: These column names will be treated as case-insensitive,
 # because these column names and the column names from the CSV
@@ -197,6 +205,7 @@ def __init__(
             storage_options=storage_options,
             tqdm_class=tqdm_class,
             taxon_colors=TAXON_COLORS,
+            aim_species_colors=AIM_SPECIES_COLORS,
             virtual_contigs=VIRTUAL_CONTIGS,
             gene_names=GENE_NAMES,
             inversion_tag_path=INVERSION_TAG_PATH,

malariagen_data/anoph/aim_data.py

@@ -208,6 +208,7 @@ def plot_aim_heatmap(
         show: plotly_params.show = True,
         renderer: plotly_params.renderer = None,
     ) -> plotly_params.figure:
+        aims = self._prep_aims_param(aims=aims)
         # Load AIM calls.
         ds = self.aim_calls(
             aims=aims,

malariagen_data/anoph/sample_metadata.py

@@ -33,6 +33,7 @@ def __init__(
         aim_analysis: Optional[str] = None,
         aim_metadata_dtype: Optional[Mapping[str, Any]] = None,
         taxon_colors: Optional[Mapping[str, str]] = None,
+        aim_species_colors: Optional[Mapping[str, str]] = None,
         **kwargs,
     ):
         # N.B., this class is designed to work cooperatively, and
@@ -73,6 +74,8 @@ def __init__(
         # Set up taxon colors.
         self._taxon_colors = taxon_colors
 
+        self._aim_species_colors = aim_species_colors
+
         # Set up extra metadata.
         self._extra_metadata: List = []
 
@@ -1304,6 +1307,11 @@ def _setup_sample_colors_plotly(
             # Special case, default taxon colors and order.
             color_discrete_map = self._taxon_colors
 
+        # Special handling for aim_species colors.
+        if color == "aim_species" and color_discrete_map is None:
+            # Special case, default aim_species colors and order.
+            color_discrete_map = self._aim_species_colors
+
         if isinstance(color, str):
             if "cohort_" + color in data.columns:
                 # Convenience to allow things like "admin1_year" instead of "cohort_admin1_year".

malariagen_data/anoph/snp_data.py

@@ -1,3 +1,4 @@
+import warnings
 from functools import lru_cache
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -1253,6 +1254,12 @@ def _snp_calls(
         if max_cohort_size is not None:
             n_samples = ds.sizes["samples"]
             if n_samples > max_cohort_size:
+                warnings.warn(
+                    f"Cohort downsampled from {n_samples} to {max_cohort_size} "
+                    "samples. Set max_cohort_size=None to disable downsampling.",
+                    UserWarning,
+                    stacklevel=2,
+                )
                 rng = np.random.default_rng(seed=random_seed)
                 loc_downsample = rng.choice(
                     n_samples, size=max_cohort_size, replace=False

malariagen_data/anopheles.py

@@ -133,12 +133,13 @@ def __init__(
         gff_default_attributes: Tuple[str, ...],
         tqdm_class,
         storage_options: Mapping,
-        taxon_colors: Optional[Mapping[str, str]],
-        virtual_contigs: Optional[Mapping[str, Sequence[str]]],
-        gene_names: Optional[Mapping[str, str]],
-        inversion_tag_path: Optional[str],
-        unrestricted_use_only: Optional[bool],
-        surveillance_use_only: Optional[bool],
+        taxon_colors: Optional[Mapping[str, str]] = None,
+        aim_species_colors: Optional[Mapping[str, str]] = None,
+        virtual_contigs: Optional[Mapping[str, Sequence[str]]] = None,
+        gene_names: Optional[Mapping[str, str]] = None,
+        inversion_tag_path: Optional[str] = None,
+        unrestricted_use_only: Optional[bool] = None,
+        surveillance_use_only: Optional[bool] = None,
     ):
         super().__init__(
             url=url,
@@ -171,6 +172,7 @@ def __init__(
             results_cache=results_cache,
             tqdm_class=tqdm_class,
             taxon_colors=taxon_colors,
+            aim_species_colors=aim_species_colors,
             virtual_contigs=virtual_contigs,
             gene_names=gene_names,
             inversion_tag_path=inversion_tag_path,

tests/anoph/test_dipclust.py

@@ -11,6 +11,9 @@ def random_transcripts_contig(*, api, contig, n):
     df_gff = api.genome_features(attributes=["ID", "Parent"])
     df_transcripts = df_gff.query(f"type == 'mRNA' and contig == '{contig}'")
     transcript_ids = df_transcripts["ID"].dropna().to_list()
+    n = min(n, len(transcript_ids))
+    if n == 0:
+        pytest.skip(f"No mRNA transcripts found for contig '{contig}'")
     transcripts = random.sample(transcript_ids, n)
     return transcripts