Solved the MRO issue

Author: jonbrenas

malariagen_data/anoph/fst.py

@@ -82,10 +82,11 @@ def _fst_gwss(
             ).compute()
 
         with self._spinner(desc="Compute Fst"):
-            fst = allel.moving_hudson_fst(ac1, ac2, size=window_size)
-            # Sometimes Fst can be very slightly below zero, clip for simplicity.
-            fst = np.clip(fst, a_min=clip_min, a_max=1)
-            x = allel.moving_statistic(pos, statistic=np.mean, size=window_size)
+            with np.errstate(divide="ignore", invalid="ignore"):
+                fst = allel.moving_hudson_fst(ac1, ac2, size=window_size)
+                # Sometimes Fst can be very slightly below zero, clip for simplicity.
+                fst = np.clip(fst, a_min=clip_min, a_max=1)
+                x = allel.moving_statistic(pos, statistic=np.mean, size=window_size)
 
         results = dict(x=x, fst=fst)
 

malariagen_data/anoph/gplt_params.py

@@ -1,7 +1,7 @@
 """Parameters for genome plotting functions. N.B., genome plots are always
 plotted with bokeh."""
 
-from typing import Literal, Mapping, Optional, Union, Sequence
+from typing import Literal, Mapping, Optional, Union, Final, Sequence
 
 import bokeh.models
 from typing_extensions import Annotated, TypeAlias
@@ -112,4 +112,11 @@
     "Passed through to bokeh line() function.",
 ]
 
+contig_colors: TypeAlias = Annotated[
+    list[str],
+    "A sequence of colors.",
+]
+
+contig_colors_default: Final[contig_colors] = list(bokeh.palettes.d3["Category20b"][5])
+
 colors: TypeAlias = Annotated[Sequence[str], "List of colors."]

malariagen_data/anoph/h12.py

@@ -265,8 +265,16 @@ def _h12_gwss(
             # Compute window midpoints.
             pos = ds_haps["variant_position"].values
             x = allel.moving_statistic(pos, statistic=np.mean, size=window_size)
+            contigs = np.asarray(
+                allel.moving_statistic(
+                    ds_haps["variant_contig"].values,
+                    statistic=np.median,
+                    size=window_size,
+                ),
+                dtype=int,
+            )
 
-        results = dict(x=x, h12=h12)
+        results = dict(x=x, h12=h12, contigs=contigs)
 
         return results
 
@@ -276,6 +284,7 @@ def _h12_gwss(
         returns=dict(
             x="An array containing the window centre point genomic positions.",
             h12="An array with h12 statistic values for each window.",
+            contigs="An array with the contig for each window. The median is chosen for windows overlapping a change of contig.",
         ),
     )
     def h12_gwss(
@@ -296,10 +305,10 @@ def h12_gwss(
         random_seed: base_params.random_seed = 42,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         # Change this name if you ever change the behaviour of this function, to
         # invalidate any previously cached data.
-        name = "h12_gwss_v1"
+        name = "h12_gwss_v2"
 
         params = dict(
             contig=contig,
@@ -326,8 +335,9 @@ def h12_gwss(
 
         x = results["x"]
         h12 = results["h12"]
+        contigs = results["contigs"]
 
-        return x, h12
+        return x, h12, contigs
 
     @check_types
     @doc(
@@ -353,14 +363,15 @@ def plot_h12_gwss_track(
         sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
         width: gplt_params.width = gplt_params.width_default,
         height: gplt_params.height = 200,
+        contig_colors: gplt_params.contig_colors = gplt_params.contig_colors_default,
         show: gplt_params.show = True,
         x_range: Optional[gplt_params.x_range] = None,
         output_backend: gplt_params.output_backend = gplt_params.output_backend_default,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
     ) -> gplt_params.figure:
         # Compute H12.
-        x, h12 = self.h12_gwss(
+        x, h12, contigs = self.h12_gwss(
             contig=contig,
             analysis=analysis,
             window_size=window_size,
@@ -411,15 +422,14 @@ def plot_h12_gwss_track(
         )
 
         # Plot H12.
-        fig.scatter(
-            x=x,
-            y=h12,
-            marker="circle",
-            size=3,
-            line_width=1,
-            line_color="black",
-            fill_color=None,
-        )
+        for s in set(contigs):
+            idxs = contigs == s
+            fig.scatter(
+                x=x[idxs],
+                y=h12[idxs],
+                marker="circle",
+                color=contig_colors[s % len(contig_colors)],
+            )
 
         # Tidy up the plot.
         fig.yaxis.axis_label = "H12"
@@ -456,6 +466,7 @@ def plot_h12_gwss(
         sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
         width: gplt_params.width = gplt_params.width_default,
         track_height: gplt_params.track_height = 170,
+        contig_colors: gplt_params.contig_colors = gplt_params.contig_colors_default,
         genes_height: gplt_params.genes_height = gplt_params.genes_height_default,
         show: gplt_params.show = True,
         output_backend: gplt_params.output_backend = gplt_params.output_backend_default,
@@ -478,6 +489,7 @@ def plot_h12_gwss(
             sizing_mode=sizing_mode,
             width=width,
             height=track_height,
+            contig_colors=contig_colors,
             show=False,
             output_backend=output_backend,
             chunks=chunks,
@@ -574,7 +586,7 @@ def plot_h12_gwss_multi_overlay_track(
             )
 
         # Determine X axis range.
-        x, _ = res[list(cohort_queries.keys())[0]]
+        x, _, _ = res[list(cohort_queries.keys())[0]]
         x_min = x[0]
         x_max = x[-1]
         if x_range is None:
@@ -609,7 +621,7 @@ def plot_h12_gwss_multi_overlay_track(
         )
 
         # Plot H12.
-        for i, (cohort_label, (x, h12)) in enumerate(res.items()):
+        for i, (cohort_label, (x, h12, contig)) in enumerate(res.items()):
             fig.scatter(
                 x=x,
                 y=h12,

malariagen_data/anoph/h1x.py

@@ -83,8 +83,16 @@ def _h1x_gwss(
             # Compute window midpoints.
             pos = ds1["variant_position"].values
             x = allel.moving_statistic(pos, statistic=np.mean, size=window_size)
+            contigs = np.asarray(
+                allel.moving_statistic(
+                    ds1["variant_contig"].values,
+                    statistic=np.median,
+                    size=window_size,
+                ),
+                dtype=int,
+            )
 
-        results = dict(x=x, h1x=h1x)
+        results = dict(x=x, h1x=h1x, contigs=contigs)
 
         return results
 
@@ -97,6 +105,7 @@ def _h1x_gwss(
         returns=dict(
             x="An array containing the window centre point genomic positions.",
             h1x="An array with H1X statistic values for each window.",
+            contigs="An array with the contig for each window. The median is chosen for windows overlapping a change of contig.",
         ),
     )
     def h1x_gwss(
@@ -118,10 +127,10 @@ def h1x_gwss(
         random_seed: base_params.random_seed = 42,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         # Change this name if you ever change the behaviour of this function, to
         # invalidate any previously cached data.
-        name = "h1x_gwss_v1"
+        name = "h1x_gwss_v2"
 
         params = dict(
             contig=contig,
@@ -149,8 +158,9 @@ def h1x_gwss(
 
         x = results["x"]
         h1x = results["h1x"]
+        contigs = results["contigs"]
 
-        return x, h1x
+        return x, h1x, contigs
 
     @check_types
     @doc(
@@ -180,14 +190,15 @@ def plot_h1x_gwss_track(
         sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
         width: gplt_params.width = gplt_params.width_default,
         height: gplt_params.height = 200,
+        contig_colors: gplt_params.contig_colors = gplt_params.contig_colors_default,
         show: gplt_params.show = True,
         x_range: Optional[gplt_params.x_range] = None,
         output_backend: gplt_params.output_backend = gplt_params.output_backend_default,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
     ) -> gplt_params.figure:
         # Compute H1X.
-        x, h1x = self.h1x_gwss(
+        x, h1x, contigs = self.h1x_gwss(
             contig=contig,
             analysis=analysis,
             window_size=window_size,
@@ -239,15 +250,14 @@ def plot_h1x_gwss_track(
         )
 
         # Plot H1X.
-        fig.scatter(
-            x=x,
-            y=h1x,
-            marker="circle",
-            size=3,
-            line_width=1,
-            line_color="black",
-            fill_color=None,
-        )
+        for s in set(contigs):
+            idxs = contigs == s
+            fig.scatter(
+                x=x[idxs],
+                y=h1x[idxs],
+                marker="circle",
+                color=contig_colors[s % len(contig_colors)],
+            )
 
         # Tidy up the plot.
         fig.yaxis.axis_label = "H1X"
@@ -288,6 +298,7 @@ def plot_h1x_gwss(
         sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
         width: gplt_params.width = gplt_params.width_default,
         track_height: gplt_params.track_height = 190,
+        contig_colors: gplt_params.contig_colors = gplt_params.contig_colors_default,
         genes_height: gplt_params.genes_height = gplt_params.genes_height_default,
         show: gplt_params.show = True,
         output_backend: gplt_params.output_backend = gplt_params.output_backend_default,
@@ -311,6 +322,7 @@ def plot_h1x_gwss(
             sizing_mode=sizing_mode,
             width=width,
             height=track_height,
+            contig_colors=contig_colors,
             show=False,
             output_backend=output_backend,
             chunks=chunks,

malariagen_data/anopheles.py

@@ -888,14 +888,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,
@@ -909,45 +919,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": []
   }
  ],