Add Ag inversion karyotyping support (#608)

* add karyotyping support

* not biallelic snp calls

* add new function to check exact alleles

* add new function to check exact alleles

* pre-commit

* pre-commit

* add tag snp file

* Update malariagen_data/ag3.py

Co-authored-by: Alistair Miles <alimanfoo@googlemail.com>

* alimanfoo comments

* add karyo params

* add manifest.in

* pre-cmmit hooks

* move karyo_params to ag3.py

* move typing import to top

* Update malariagen_data/ag3.py

Co-authored-by: Alistair Miles <alimanfoo@googlemail.com>

* Update malariagen_data/ag3.py

Co-authored-by: Alistair Miles <alimanfoo@googlemail.com>

* use doc generator

* pre-commit

* fix for tests - new col name

* Update malariagen_data/ag3.py

Co-authored-by: Alistair Miles <alimanfoo@googlemail.com>

* fix 2 for tests

* Update tests/test_ag3.py

* Update tests/test_ag3.py

* add notebook

* Update malariagen_data/ag3.py

Co-authored-by: Alistair Miles <alimanfoo@googlemail.com>

* update notebook

* fix dtypes

* pre-commit

---------

Co-authored-by: Alistair Miles <alimanfoo@googlemail.com>
Co-authored-by: Alistair Miles <am23@sanger.ac.uk>

Author: 3 people

MANIFEST.in

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

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.
 
@@ -341,3 +355,74 @@ 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
+        )
+        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()
+
+        with self._spinner("Inferring karyotype from tag SNPs"):
+            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,
+                },
+            )
+            kt_dtype = CategoricalDtype(categories=[0, 1, 2], ordered=True)
+            df[f"karyotype_{inversion}"] = df[f"karyotype_{inversion}"].astype(kt_dtype)
+
+        return df