gene cnv optimisation (#655)

Author: alimanfoo

malariagen_data/anopheles.py

@@ -886,14 +886,24 @@ def _gene_cnv(
         chunks,
         inline_array,
     ):
-        debug = self._log.debug
-
-        debug("sanity check")
+        # Sanity check.
         assert isinstance(region, Region)
 
-        debug("access HMM data")
+        # Access genes within the region of interest.
+        df_genome_features = self.genome_features(region=region)
+        sample_query_options = sample_query_options or {}
+        df_genes = df_genome_features.query(
+            f"type == '{self._gff_gene_type}'", **sample_query_options
+        )
+
+        # Refine the region for CNV data to ensure coverage of all requested genes.
+        cnv_region = Region(
+            region.contig, df_genes["start"].min(), df_genes["end"].max()
+        )
+
+        # Access HMM data.
         ds_hmm = self.cnv_hmm(
-            region=region.contig,
+            region=cnv_region,
             sample_sets=sample_sets,
             sample_query=sample_query,
             sample_query_options=sample_query_options,
@@ -907,45 +917,38 @@ def _gene_cnv(
         with self._dask_progress(desc="Load CNV HMM data"):
             pos, end, cn = dask.compute(pos, end, cn)
 
-        debug("access genes")
-        df_genome_features = self.genome_features(region=region)
-        sample_query_options = sample_query_options or {}
-        df_genes = df_genome_features.query(
-            f"type == '{self._gff_gene_type}'", **sample_query_options
-        )
-
-        debug("setup intermediates")
+        # Set up intermediates.
         windows = []
         modes = []
         counts = []
 
-        debug("iterate over genes")
+        # Iterate over genes.
         genes_iterator = self._progress(
             df_genes.itertuples(),
             desc="Compute modal gene copy number",
             total=len(df_genes),
         )
         for gene in genes_iterator:
-            # locate windows overlapping the gene
+            # Locate windows overlapping the gene.
             loc_gene_start = bisect_left(end, gene.start)
             loc_gene_stop = bisect_right(pos, gene.end)
             w = loc_gene_stop - loc_gene_start
             windows.append(w)
 
-            # slice out copy number data for the given gene
+            # Slice out copy number data for the given gene.
             cn_gene = cn[loc_gene_start:loc_gene_stop]
 
-            # compute the modes
+            # Compute the modes.
             m, c = _cn_mode(cn_gene, vmax=12)
             modes.append(m)
             counts.append(c)
 
-        debug("combine results")
+        # Combine results.
         windows = np.array(windows)
         modes = np.vstack(modes)
         counts = np.vstack(counts)
 
-        debug("build dataset")
+        # Build dataset.
         ds_out = xr.Dataset(
             coords={
                 "gene_id": (["genes"], df_genes["ID"].values),

notebooks/plot_frequencies_heatmap.ipynb

@@ -381,6 +381,44 @@
    "id": "86c5c594",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "interesting_cyp_genes = [\n",
+    "    \"AGAP002862\",  # Cyp6aa1\n",
+    "    \"AGAP013128\",  # Cyp6aa2\n",
+    "    \"AGAP002865\",  # Cyp6p3\n",
+    "    \"AGAP000818\",  # Cyp9k1\n",
+    "    \"AGAP008212\",  # Cyp6m2\n",
+    "    \"AGAP008218\",  # Cyp6z2    \n",
+    "]\n",
+    "\n",
+    "cyp_cnv_freqs_df = ag3.gene_cnv_frequencies(\n",
+    "    region=interesting_cyp_genes,\n",
+    "    cohorts=\"admin1_year\",\n",
+    "    sample_sets=(\"AG1000G-BF-A\", \"AG1000G-BF-B\", \"AG1000G-BF-C\"),\n",
+    "    sample_query=\"taxon == 'coluzzii'\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d7ad130-30c2-4cd3-8906-a7ada3ccc75f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ag3.plot_frequencies_heatmap(\n",
+    "    df=cyp_cnv_freqs_df,\n",
+    "    color_continuous_scale=\"Blues\",\n",
+    "    title=\"Cyp gene CNV frequencies\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "83aab417-632e-4fd2-8da4-3ffdd6e233f6",
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],

tests/test_af1.py

@@ -6,7 +6,7 @@
 from malariagen_data.util import locate_region, resolve_region
 
 
-def setup_af1(url="simplecache::gs://vo_afun_release/", **kwargs):
+def setup_af1(url="simplecache::gs://vo_afun_release_master_us_central1/", **kwargs):
     kwargs.setdefault("check_location", False)
     kwargs.setdefault("show_progress", False)
     if url is None:

tests/test_ag3.py

@@ -15,7 +15,7 @@
 contigs = "2R", "2L", "3R", "3L", "X"
 
 
-def setup_ag3(url="simplecache::gs://vo_agam_release/", **kwargs):
+def setup_ag3(url="simplecache::gs://vo_agam_release_master_us_central1/", **kwargs):
     kwargs.setdefault("check_location", False)
     kwargs.setdefault("show_progress", False)
     if url is None: