Merge branch 'master' into 583-GWSS-colors

Author: jonbrenas

docs/source/Af1.rst

@@ -164,6 +164,8 @@ Genome-wide selection scans
     plot_h12_calibration
     h12_gwss
     plot_h12_gwss
+    plot_h12_gwss_multi_panel
+    plot_h12_gwss_multi_overlay
     h1x_gwss
     plot_h1x_gwss
     g123_calibration

docs/source/Ag3.rst

@@ -174,6 +174,8 @@ Genome-wide selection scans
     plot_h12_calibration
     h12_gwss
     plot_h12_gwss
+    plot_h12_gwss_multi_panel
+    plot_h12_gwss_multi_overlay
     h1x_gwss
     plot_h1x_gwss
     g123_calibration

malariagen_data/anoph/distance.py

@@ -367,8 +367,13 @@ def _njt(
         # neighbour-joining iterations.
         progress_options = dict(desc="Construct neighbour-joining tree", leave=False)
 
-        # Decide which algorithm to use and run the neighbour-joining.
-        if algorithm == "rapid":
+        # Decide which algorithm to use and run the neighbour-joining. The "dynamic"
+        # algorithm is fastest.
+        if algorithm == "dynamic":
+            Z = anjl.dynamic_nj(
+                D=D, progress=self._progress, progress_options=progress_options
+            )
+        elif algorithm == "rapid":
             Z = anjl.rapid_nj(
                 D=D, progress=self._progress, progress_options=progress_options
             )

malariagen_data/anoph/distance_params.py

@@ -14,8 +14,8 @@
 default_distance_metric: distance_metric = "cityblock"
 
 nj_algorithm: TypeAlias = Annotated[
-    Literal["rapid", "canonical"],
-    "Neighbour-joining algorithm to use.",
+    Literal["dynamic", "rapid", "canonical"],
+    "Neighbour-joining algorithm to use. The 'dynamic' algorithm is fastest.",
 ]
 
-default_nj_algorithm: nj_algorithm = "rapid"
+default_nj_algorithm: nj_algorithm = "dynamic"

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

@@ -83,6 +83,14 @@
     "A bokeh figure (only returned if show=False).",
 ]
 
+def_figure: TypeAlias = Annotated[
+    # Use quite a broad type here to accommodate both single-panel figures
+    # created via bokeh.plotting and multi-panel figures created via
+    # bokeh.layouts.
+    bokeh.model.Model,
+    "A bokeh figure.",
+]
+
 output_backend: TypeAlias = Annotated[
     Literal["canvas", "webgl", "svg"],
     """

malariagen_data/anoph/h12.py

@@ -526,6 +526,307 @@ def plot_h12_gwss(
         else:
             return fig
 
+    @check_types
+    @doc(
+        summary="Plot h12 GWSS data track with multiple traces overlaid.",
+    )
+    def plot_h12_gwss_multi_overlay_track(
+        self,
+        contig: base_params.contig,
+        cohorts: base_params.cohorts,
+        window_size: h12_params.multi_window_size,
+        cohort_size: Optional[base_params.cohort_size] = h12_params.cohort_size_default,
+        sample_query: Optional[base_params.sample_query] = None,
+        analysis: hap_params.analysis = base_params.DEFAULT,
+        min_cohort_size: Optional[
+            base_params.min_cohort_size
+        ] = h12_params.min_cohort_size_default,
+        max_cohort_size: Optional[
+            base_params.max_cohort_size
+        ] = h12_params.max_cohort_size_default,
+        title: Optional[gplt_params.title] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        colors: gplt_params.colors = bokeh.palettes.d3["Category10"][10],
+        random_seed: base_params.random_seed = 42,
+        sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
+        width: gplt_params.width = gplt_params.width_default,
+        height: gplt_params.height = 200,
+        show: gplt_params.show = True,
+        x_range: Optional[gplt_params.x_range] = None,
+        output_backend: gplt_params.output_backend = gplt_params.output_backend_default,
+    ) -> gplt_params.figure:
+        cohort_queries = self._setup_cohort_queries(
+            cohorts=cohorts,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
+            cohort_size=cohort_size,
+            min_cohort_size=None,
+        )
+
+        if isinstance(window_size, int):
+            window_size = {k: window_size for k in cohort_queries.keys()}
+        elif isinstance(window_size, Mapping):
+            if set(window_size.keys()) != set(cohort_queries.keys()):
+                raise ValueError("Cohorts and window_sizes should have the same keys.")
+
+        # Compute H12.
+        res = {}
+        for cohort_label, cohort_query in cohort_queries.items():
+            res[cohort_label] = self.h12_gwss(
+                contig=contig,
+                analysis=analysis,
+                window_size=window_size[cohort_label],
+                cohort_size=cohort_size,
+                min_cohort_size=min_cohort_size,
+                max_cohort_size=max_cohort_size,
+                sample_query=cohort_query,
+                sample_sets=sample_sets,
+                random_seed=random_seed,
+            )
+
+        # Determine X axis range.
+        x, _ = res[list(cohort_queries.keys())[0]]
+        x_min = x[0]
+        x_max = x[-1]
+        if x_range is None:
+            x_range = bokeh.models.Range1d(x_min, x_max, bounds="auto")
+
+        # Create a figure.
+        xwheel_zoom = bokeh.models.WheelZoomTool(
+            dimensions="width", maintain_focus=False
+        )
+
+        fig = bokeh.plotting.figure(
+            title=title,
+            tools=[
+                "xpan",
+                "xzoom_in",
+                "xzoom_out",
+                xwheel_zoom,
+                "reset",
+                "save",
+                "crosshair",
+            ],
+            active_inspect=None,
+            active_scroll=xwheel_zoom,
+            active_drag="xpan",
+            sizing_mode=sizing_mode,
+            width=width,
+            height=height,
+            toolbar_location="above",
+            x_range=x_range,
+            y_range=(0, 1),
+            output_backend=output_backend,
+        )
+
+        # Plot H12.
+        for i, (cohort_label, (x, h12)) in enumerate(res.items()):
+            fig.scatter(
+                x=x,
+                y=h12,
+                marker="circle",
+                size=3,
+                line_width=1,
+                line_color=colors[i % len(colors)],
+                fill_color=None,
+                legend_label=cohort_label,
+            )
+
+        # Tidy up the plot.
+        fig.yaxis.axis_label = "H12"
+        fig.yaxis.ticker = [0, 1]
+        self._bokeh_style_genome_xaxis(fig, contig)
+
+        if show:  # pragma: no cover
+            bokeh.plotting.show(fig)
+            return None
+        else:
+            return fig
+
+    @check_types
+    @doc(
+        summary="Plot h12 GWSS data with multiple traces overlaid.",
+    )
+    def plot_h12_gwss_multi_overlay(
+        self,
+        contig: base_params.contig,
+        cohorts: base_params.cohorts,
+        window_size: h12_params.multi_window_size,
+        cohort_size: Optional[base_params.cohort_size] = h12_params.cohort_size_default,
+        sample_query: Optional[base_params.sample_query] = None,
+        analysis: hap_params.analysis = base_params.DEFAULT,
+        min_cohort_size: Optional[
+            base_params.min_cohort_size
+        ] = h12_params.min_cohort_size_default,
+        max_cohort_size: Optional[
+            base_params.max_cohort_size
+        ] = h12_params.max_cohort_size_default,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        colors: gplt_params.colors = bokeh.palettes.d3["Category10"][10],
+        random_seed: base_params.random_seed = 42,
+        title: Optional[gplt_params.title] = None,
+        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,
+        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,
+    ) -> gplt_params.figure:
+        # Plot GWSS track.
+        fig1 = self.plot_h12_gwss_multi_overlay_track(
+            contig=contig,
+            sample_query=sample_query,
+            cohorts=cohorts,
+            cohort_size=cohort_size,
+            window_size=window_size,
+            analysis=analysis,
+            min_cohort_size=min_cohort_size,
+            max_cohort_size=max_cohort_size,
+            sample_query_options=sample_query_options,
+            sample_sets=sample_sets,
+            colors=colors,
+            random_seed=random_seed,
+            title=title,
+            sizing_mode=sizing_mode,
+            width=width,
+            height=track_height,
+            show=False,
+            output_backend=output_backend,
+        )
+
+        fig1.xaxis.visible = False
+        fig1.legend.location = "top_right"
+        fig1.legend.click_policy = "hide"
+
+        # Plot genes.
+        fig2 = self.plot_genes(
+            region=contig,
+            sizing_mode=sizing_mode,
+            width=width,
+            height=genes_height,
+            x_range=fig1.x_range,
+            show=False,
+            output_backend=output_backend,
+        )
+
+        # Combine plots into a single figure.
+        fig = bokeh.layouts.gridplot(
+            [fig1, fig2],
+            ncols=1,
+            toolbar_location="above",
+            merge_tools=True,
+            sizing_mode=sizing_mode,
+            toolbar_options=dict(active_inspect=None),
+        )
+
+        if show:  # pragma: no cover
+            bokeh.plotting.show(fig)
+            return None
+        else:
+            return fig
+
+    @check_types
+    @doc(
+        summary="Plot h12 GWSS data with multiple tracks.",
+    )
+    def plot_h12_gwss_multi_panel(
+        self,
+        contig: base_params.contig,
+        cohorts: base_params.cohorts,
+        window_size: h12_params.multi_window_size,
+        cohort_size: Optional[base_params.cohort_size] = h12_params.cohort_size_default,
+        sample_query: Optional[base_params.sample_query] = None,
+        analysis: hap_params.analysis = base_params.DEFAULT,
+        min_cohort_size: Optional[
+            base_params.min_cohort_size
+        ] = h12_params.min_cohort_size_default,
+        max_cohort_size: Optional[
+            base_params.max_cohort_size
+        ] = h12_params.max_cohort_size_default,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        random_seed: base_params.random_seed = 42,
+        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,
+        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,
+    ) -> gplt_params.figure:
+        cohort_queries = self._setup_cohort_queries(
+            cohorts=cohorts,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
+            cohort_size=cohort_size,
+            min_cohort_size=None,
+        )
+
+        if isinstance(window_size, int):
+            window_size = {k: window_size for k in cohort_queries.keys()}
+        elif isinstance(window_size, Mapping):
+            if set(window_size.keys()) != set(cohort_queries.keys()):
+                raise ValueError("Cohorts and window_sizes should have the same keys.")
+
+        # Plot GWSS track.
+        figs: list[gplt_params.def_figure] = []
+        for i, (cohort_label, cohort_query) in enumerate(cohort_queries.items()):
+            params = dict(
+                contig=contig,
+                analysis=analysis,
+                window_size=window_size[cohort_label],
+                sample_sets=sample_sets,
+                sample_query=cohort_query,
+                cohort_size=cohort_size,
+                min_cohort_size=min_cohort_size,
+                max_cohort_size=max_cohort_size,
+                random_seed=random_seed,
+                title=cohort_label,  # Deal with a choice of titles later
+                sizing_mode=sizing_mode,
+                width=width,
+                height=track_height,
+                show=False,
+                output_backend=output_backend,
+            )
+            if i > 0:
+                track = self.plot_h12_gwss_track(x_range=figs[0].x_range, **params)
+            else:
+                track = self.plot_h12_gwss_track(**params)
+            track.xaxis.visible = False
+            figs.append(track)
+
+        # Plot genes.
+        fig2 = self.plot_genes(
+            region=contig,
+            sizing_mode=sizing_mode,
+            width=width,
+            height=genes_height,
+            x_range=figs[0].x_range,
+            show=False,
+            output_backend=output_backend,
+        )
+
+        figs.append(fig2)
+
+        # Combine plots into a single figure.
+        fig = bokeh.layouts.gridplot(
+            figs,
+            ncols=1,
+            toolbar_location="above",
+            merge_tools=True,
+            sizing_mode=sizing_mode,
+            toolbar_options=dict(active_inspect=None),
+        )
+
+        if show:  # pragma: no cover
+            bokeh.plotting.show(fig)
+            return None
+        else:
+            return fig
+
 
 def haplotype_frequencies(h):
     """Compute haplotype frequencies, returning a dictionary that maps