Merge pull request malariagen#962 from Tanisha127/fix/250-consistent-aim-coloring

feat: add consistent aim_species colors in plot_pca_coords

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/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/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,