Addressing comments

Author: jonbrenas

malariagen_data/anoph/karyotype.py

@@ -4,15 +4,36 @@
 import allel  # type: ignore
 
 from numpydoc_decorator import doc
-from ..util import check_types, _karyotype_tags_n_alt
+from ..util import check_types
 from . import base_params
 from typing import Optional
 
 from .snp_data import AnophelesSnpData
 from .karyotype_params import inversion_param
 
 
-class AnophelesKaryotypeData(AnophelesSnpData):
+def _karyotype_tags_n_alt(gt, alts, inversion_alts):
+    # could be Numba'd for speed but was already quick (not many inversion tag snps)
+    n_sites = gt.shape[0]
+    n_samples = gt.shape[1]
+
+    # create empty array
+    inv_n_alt = np.empty((n_sites, n_samples), dtype=np.int8)
+
+    # for every site
+    for i in range(n_sites):
+        # find the index of the correct tag snp allele
+        tagsnp_index = np.where(alts[i] == inversion_alts[i])[0]
+
+        for j in range(n_samples):
+            # count alleles which == tag snp allele and store
+            n_tag_alleles = np.sum(gt[i, j] == tagsnp_index[0])
+            inv_n_alt[i, j] = n_tag_alleles
+
+    return inv_n_alt
+
+
+class AnophelesKaryotypeAnalysis(AnophelesSnpData):
     def __init__(
         self,
         inversion_tag_path: Optional[str] = None,
@@ -23,13 +44,11 @@ def __init__(
         # to the superclass constructor.
         super().__init__(**kwargs)
 
-        # If provided, this analysis version will override the
-        # default value provided in the release configuration.
         self._inversion_tag_path = inversion_tag_path
 
     @check_types
     @doc(
-        summary="Load tag SNPs for a given inversion in Ag.",
+        summary="Load tag SNPs for a given inversion.",
     )
     def load_inversion_tags(self, inversion: inversion_param) -> pd.DataFrame:
         # needs to be modified depending on where we are hosting

malariagen_data/anoph/karyotype_params.py

@@ -1,10 +1,9 @@
 """Parameter definitions for karyotype analysis functions."""
 
-from typing import Literal
 
 from typing_extensions import Annotated, TypeAlias
 
 inversion_param: TypeAlias = Annotated[
-    Literal["2La", "2Rb", "2Rc_gam", "2Rc_col", "2Rd", "2Rj"],
+    str,
     "Name of inversion to infer karyotype for.",
 ]

malariagen_data/anopheles.py

@@ -29,7 +29,7 @@
     plotly_params,
     xpehh_params,
 )
-from .anoph.karyotype import AnophelesKaryotypeData
+from .anoph.karyotype import AnophelesKaryotypeAnalysis
 from .anoph.aim_data import AnophelesAimData
 from .anoph.base import AnophelesBase
 from .anoph.cnv_data import AnophelesCnvData
@@ -95,7 +95,7 @@ class AnophelesDataResource(
     AnophelesPca,
     PlinkConverter,
     AnophelesIgv,
-    AnophelesKaryotypeData,
+    AnophelesKaryotypeAnalysis,
     AnophelesAimData,
     AnophelesHapData,
     AnophelesSnpData,

malariagen_data/util.py

@@ -1591,27 +1591,6 @@ def distributed_client():
     return client
 
 
-def _karyotype_tags_n_alt(gt, alts, inversion_alts):
-    # could be Numba'd for speed but was already quick (not many inversion tag snps)
-    n_sites = gt.shape[0]
-    n_samples = gt.shape[1]
-
-    # create empty array
-    inv_n_alt = np.empty((n_sites, n_samples), dtype=np.int8)
-
-    # for every site
-    for i in range(n_sites):
-        # find the index of the correct tag snp allele
-        tagsnp_index = np.where(alts[i] == inversion_alts[i])[0]
-
-        for j in range(n_samples):
-            # count alleles which == tag snp allele and store
-            n_tag_alleles = np.sum(gt[i, j] == tagsnp_index[0])
-            inv_n_alt[i, j] = n_tag_alleles
-
-    return inv_n_alt
-
-
 def prep_samples_for_cohort_grouping(*, df_samples, area_by, period_by):
     # Take a copy, as we will modify the dataframe.
     df_samples = df_samples.copy()

notebooks/karyotype.ipynb

@@ -359,14 +359,6 @@
    "source": [
     "ag3.plot_pca_coords(pca_df_2rc_col, color=\"karyotype_2Rc_col\", symbol=\"country\", width=600, height=500)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "d6fb7237-bb4e-490e-9ae0-33a52d4fa650",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -391,7 +383,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.10.11"
   }
  },
  "nbformat": 4,

tests/integration/test_af1.py

@@ -84,15 +84,9 @@ def test_locate_region(region_raw):
 def test_karyotyping(inversion):
     af1 = setup_af1()
 
-    if inversion == "X_x":
-        with pytest.raises(TypeError):
-            af1.karyotype(
-                inversion=inversion, sample_sets="AG1000G-GH", sample_query=None
-            )
-    else:
-        with pytest.raises(FileNotFoundError):
-            af1.karyotype(
-                inversion=inversion,
-                sample_sets="1229-VO-GH-DADZIE-VMF00095",
-                sample_query=None,
-            )
+    with pytest.raises(FileNotFoundError):
+        af1.karyotype(
+            inversion=inversion,
+            sample_sets="1229-VO-GH-DADZIE-VMF00095",
+            sample_query=None,
+        )

tests/integration/test_ag3.py

@@ -167,7 +167,7 @@ def test_karyotyping(inversion):
     ag3 = setup_ag3(cohorts_analysis="20230516")
 
     if inversion == "X_x":
-        with pytest.raises(TypeError):
+        with pytest.raises(ValueError):
             ag3.karyotype(
                 inversion=inversion, sample_sets="AG1000G-GH", sample_query=None
             )