Merge branch 'master' into GH407_allow_gene_labels

Author: leehart

malariagen_data/ag3.py

@@ -378,6 +378,7 @@ def karyotype(
         inversion: inversion_param,
         sample_sets: Optional[base_params.sample_sets] = None,
         sample_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
     ) -> pd.DataFrame:
         # load tag snp data
         df_tagsnps = self.load_inversion_tags(inversion=inversion)
@@ -390,7 +391,10 @@ def karyotype(
         region = f"{contig}:{start}-{end}"
 
         ds_snps = self.snp_calls(
-            region=region, sample_sets=sample_sets, sample_query=sample_query
+            region=region,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
         )
 
         with self._spinner("Inferring karyotype from tag SNPs"):

malariagen_data/anoph/fst.py

@@ -28,11 +28,13 @@ def __init__(
 
     def _fst_gwss(
         self,
+        *,
         contig,
         window_size,
         sample_sets,
         cohort1_query,
         cohort2_query,
+        sample_query_options,
         site_mask,
         cohort_size,
         min_cohort_size,
@@ -46,6 +48,7 @@ def _fst_gwss(
         ac1 = self.snp_allele_counts(
             region=contig,
             sample_query=cohort1_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             cohort_size=cohort_size,
@@ -58,6 +61,7 @@ def _fst_gwss(
         ac2 = self.snp_allele_counts(
             region=contig,
             sample_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             cohort_size=cohort_size,
@@ -78,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)
 
@@ -104,6 +109,7 @@ def fst_gwss(
         window_size: fst_params.window_size,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
@@ -127,6 +133,7 @@ def fst_gwss(
             window_size=window_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=self._prep_sample_sets_param(sample_sets=sample_sets),
             site_mask=self._prep_optional_site_mask_param(site_mask=site_mask),
             cohort_size=cohort_size,
@@ -161,6 +168,7 @@ def plot_fst_gwss_track(
         window_size: fst_params.window_size,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
@@ -189,6 +197,7 @@ def plot_fst_gwss_track(
             max_cohort_size=max_cohort_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             random_seed=random_seed,
@@ -265,6 +274,7 @@ def plot_fst_gwss(
         window_size: fst_params.window_size,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
@@ -290,6 +300,7 @@ def plot_fst_gwss(
             window_size=window_size,
             cohort1_query=cohort1_query,
             cohort2_query=cohort2_query,
+            sample_query_options=sample_query_options,
             sample_sets=sample_sets,
             site_mask=site_mask,
             cohort_size=cohort_size,
@@ -348,6 +359,7 @@ def average_fst(
         region: base_params.region,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
+        sample_query_options: Optional[base_params.sample_query] = None,
         sample_sets: Optional[base_params.sample_sets] = None,
         cohort_size: Optional[base_params.cohort_size] = fst_params.cohort_size_default,
         min_cohort_size: Optional[
@@ -366,6 +378,7 @@ def average_fst(
             region=region,
             sample_sets=sample_sets,
             sample_query=cohort1_query,
+            sample_query_options=sample_query_options,
             cohort_size=cohort_size,
             site_mask=site_mask,
             site_class=site_class,
@@ -377,6 +390,7 @@ def average_fst(
             region=region,
             sample_sets=sample_sets,
             sample_query=cohort2_query,
+            sample_query_options=sample_query_options,
             cohort_size=cohort_size,
             site_mask=site_mask,
             site_class=site_class,

malariagen_data/anoph/genome_features.py

@@ -149,10 +149,11 @@ def genome_features(
                         df_part = df_part.query(f"end >= {r.start}")
                     parts.append(df_part)
                 df = pd.concat(parts, axis=0)
-                return df.reset_index(drop=True).copy()
+                return df.sort_values(["contig", "start"]).reset_index(drop=True).copy()
 
             return (
                 self._genome_features(attributes=attributes_normed)
+                .sort_values(["contig", "start"])
                 .reset_index(drop=True)
                 .copy()
             )

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,6 +112,13 @@
     "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."]
 
 gene_labels: TypeAlias = Annotated[

malariagen_data/anoph/h12.py

@@ -266,8 +266,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
 
@@ -277,6 +285,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(
@@ -297,10 +306,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,
@@ -327,8 +336,9 @@ def h12_gwss(
 
         x = results["x"]
         h12 = results["h12"]
+        contigs = results["contigs"]
 
-        return x, h12
+        return x, h12, contigs
 
     @check_types
     @doc(
@@ -354,14 +364,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,
@@ -412,15 +423,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"
@@ -457,6 +467,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,
@@ -479,6 +490,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,
@@ -575,7 +587,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:
@@ -610,7 +622,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,