Merge branch 'master' into 583-GWSS-colors

Author: jonbrenas

MANIFEST.in

@@ -0,0 +1 @@
+include malariagen_data/resources/*

docs/source/Af1.rst

@@ -53,6 +53,8 @@ Sample metadata access
     count_samples
     plot_samples_bar
     plot_samples_interactive_map
+    plot_sample_location_mapbox
+    plot_sample_location_geo
     wgs_data_catalog
     cohorts
 
@@ -132,6 +134,7 @@ Genetic distance and neighbour-joining trees (NJT)
     :toctree: generated/
 
     plot_njt
+    njt
     biallelic_diplotype_pairwise_distances
 
 Heterozygosity analysis

docs/source/Ag3.rst

@@ -54,6 +54,8 @@ Sample metadata access
     lookup_sample
     plot_samples_bar
     plot_samples_interactive_map
+    plot_sample_location_mapbox
+    plot_sample_location_geo
     wgs_data_catalog
     cohorts
 
@@ -142,6 +144,7 @@ Genetic distance and neighbour-joining trees (NJT)
     :toctree: generated/
 
     plot_njt
+    njt
     biallelic_diplotype_pairwise_distances
 
 Heterozygosity analysis

malariagen_data/af1.py

@@ -40,7 +40,7 @@ class Af1(AnophelesDataResource):
     debug : bool, optional
         Set to True to enable debug level logging.
     show_progress : bool, optional
-        If True, show a progress bar during longer-running computations.
+        If True, show a progress bar during longer-running computations. The default can be overridden using an environmental variable named MGEN_SHOW_PROGRESS.
     check_location : bool, optional
         If True, use ipinfo to check the location of the client system.
     **kwargs
@@ -82,7 +82,7 @@ def __init__(
         results_cache=None,
         log=sys.stdout,
         debug=False,
-        show_progress=True,
+        show_progress=None,
         check_location=True,
         cohorts_analysis=None,
         site_filters_analysis=None,

malariagen_data/ag3.py

@@ -2,11 +2,19 @@
 
 import dask
 import pandas as pd  # type: ignore
+from pandas import CategoricalDtype
+import numpy as np  # type: ignore
+import allel  # type: ignore
 import plotly.express as px  # type: ignore
 
 import malariagen_data
 from .anopheles import AnophelesDataResource
 
+from numpydoc_decorator import doc
+from .util import check_types, _karyotype_tags_n_alt
+from .anoph import base_params
+from typing import Optional, Literal, Annotated, TypeAlias
+
 # silence dask performance warnings
 dask.config.set(**{"array.slicing.split_native_chunks": False})  # type: ignore
 
@@ -75,6 +83,12 @@ def _setup_aim_palettes():
 }
 
 
+inversion_param: TypeAlias = Annotated[
+    Literal["2La", "2Rb", "2Rc_gam", "2Rc_col", "2Rd", "2Rj"],
+    "Name of inversion to infer karyotype for.",
+]
+
+
 class Ag3(AnophelesDataResource):
     """Provides access to data from Ag3.x releases.
 
@@ -98,7 +112,7 @@ class Ag3(AnophelesDataResource):
     debug : bool, optional
         Set to True to enable debug level logging.
     show_progress : bool, optional
-        If True, show a progress bar during longer-running computations.
+        If True, show a progress bar during longer-running computations. The default can be overridden using an environmental variable named MGEN_SHOW_PROGRESS.
     check_location : bool, optional
         If True, use ipinfo to check the location of the client system.
     **kwargs
@@ -140,7 +154,7 @@ def __init__(
         results_cache=None,
         log=sys.stdout,
         debug=False,
-        show_progress=True,
+        show_progress=None,
         check_location=True,
         cohorts_analysis=None,
         aim_analysis=None,
@@ -341,3 +355,78 @@ def _results_cache_add_analysis_params(self, params):
         super()._results_cache_add_analysis_params(params)
         # override parent class to add AIM analysis
         params["aim_analysis"] = self._aim_analysis
+
+    @check_types
+    @doc(
+        summary="Load tag SNPs for a given inversion in Ag.",
+    )
+    def load_inversion_tags(self, inversion: inversion_param) -> pd.DataFrame:
+        # needs to be modified depending on where we are hosting
+        import importlib.resources
+        from . import resources
+
+        with importlib.resources.path(resources, "karyotype_tag_snps.csv") as path:
+            df_tag_snps = pd.read_csv(path, sep=",")
+        return df_tag_snps.query(f"inversion == '{inversion}'").reset_index()
+
+    @check_types
+    @doc(
+        summary="Infer karyotype from tag SNPs for a given inversion in Ag.",
+    )
+    def karyotype(
+        self,
+        inversion: inversion_param,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        sample_query: Optional[base_params.sample_query] = None,
+    ) -> pd.DataFrame:
+        # load tag snp data
+        df_tagsnps = self.load_inversion_tags(inversion=inversion)
+        inversion_pos = df_tagsnps["position"]
+        inversion_alts = df_tagsnps["alt_allele"]
+        contig = inversion[0:2]
+
+        # get snp calls for inversion region
+        start, end = np.min(inversion_pos), np.max(inversion_pos)
+        region = f"{contig}:{start}-{end}"
+
+        ds_snps = self.snp_calls(
+            region=region, sample_sets=sample_sets, sample_query=sample_query
+        )
+
+        with self._spinner("Inferring karyotype from tag SNPs"):
+            # access variables we need
+            geno = allel.GenotypeDaskArray(ds_snps["call_genotype"].data)
+            pos = allel.SortedIndex(ds_snps["variant_position"].values)
+            samples = ds_snps["sample_id"].values
+            alts = ds_snps["variant_allele"].values.astype(str)
+
+            # subset to position of inversion tags
+            mask = pos.locate_intersection(inversion_pos)[0]
+            alts = alts[mask]
+            geno = geno.compress(mask, axis=0).compute()
+
+            # infer karyotype
+            gn_alt = _karyotype_tags_n_alt(
+                gt=geno, alts=alts, inversion_alts=inversion_alts
+            )
+            is_called = geno.is_called()
+
+            # calculate mean genotype for each sample whilst masking missing calls
+            av_gts = np.mean(np.ma.MaskedArray(gn_alt, mask=~is_called), axis=0)
+            total_sites = np.sum(is_called, axis=0)
+
+            df = pd.DataFrame(
+                {
+                    "sample_id": samples,
+                    "inversion": inversion,
+                    f"karyotype_{inversion}_mean": av_gts,
+                    # round the genotypes then convert to int
+                    f"karyotype_{inversion}": av_gts.round().astype(int),
+                    "total_tag_snps": total_sites,
+                },
+            )
+            # Allow filling missing values with "<NA>" visible placeholder.
+            kt_dtype = CategoricalDtype(categories=[0, 1, 2, "<NA>"], ordered=True)
+            df[f"karyotype_{inversion}"] = df[f"karyotype_{inversion}"].astype(kt_dtype)
+
+        return df

malariagen_data/anoph/base.py

@@ -1,3 +1,5 @@
+import os
+
 import json
 from contextlib import nullcontext
 from datetime import date
@@ -54,12 +56,24 @@ def __init__(
         bokeh_output_notebook: bool = False,
         log: Optional[Union[str, IO]] = None,
         debug: bool = False,
-        show_progress: bool = False,
+        show_progress: Optional[bool] = None,
         check_location: bool = False,
         storage_options: Optional[Mapping] = None,
         results_cache: Optional[str] = None,
         tqdm_class=None,
     ):
+        # If show_progress has not been specified, then determine the default.
+        if show_progress is None:
+            # Get the env var, if it exists.
+            show_progress_env = os.getenv("MGEN_SHOW_PROGRESS")
+
+            # If the env var does not exist, then use the class default.
+            # Otherwise, convert the env var value to a boolean and use that.
+            if show_progress_env is None:
+                show_progress = True
+            else:
+                show_progress = show_progress_env.lower() in ("true", "1", "yes", "on")
+
         self._config_path = config_path
         self._pre = pre
         self._gcs_default_url = gcs_default_url

malariagen_data/anoph/dipclust_params.py

@@ -1,6 +1,6 @@
 """Parameters for diplotype clustering functions."""
 
-from .diplotype_distance_params import distance_metric
+from .distance_params import distance_metric
 from .clustering_params import linkage_method