Implement plot_samples_bar() (#413)

* implement plot_samples_bar

* exclude lines from coverage

* exclude some more lines from coverage

malariagen_data/ag3.py

@@ -1893,7 +1893,7 @@ def plot_cnv_hmm_coverage_track(
         self._bokeh_style_genome_xaxis(fig, region_prepped.contig)
         fig.add_layout(fig.legend[0], "right")
 
-        if show:
+        if show:  # pragma: no cover
             bkplt.show(fig)
             return None
         else:
@@ -1986,7 +1986,7 @@ def plot_cnv_hmm_coverage(
             sizing_mode=sizing_mode,
         )
 
-        if show:
+        if show:  # pragma: no cover
             bkplt.show(fig)
             return None
         else:
@@ -2156,7 +2156,7 @@ def plot_cnv_hmm_heatmap_track(
         )
         fig.add_layout(color_bar, "right")
 
-        if show:
+        if show:  # pragma: no cover
             bkplt.show(fig)
             return None
         else:
@@ -2244,7 +2244,7 @@ def plot_cnv_hmm_heatmap(
             sizing_mode=sizing_mode,
         )
 
-        if show:
+        if show:  # pragma: no cover
             bkplt.show(fig)
             return None
         else:

malariagen_data/anoph/aim_data.py

@@ -307,7 +307,7 @@ def plot_aim_heatmap(
             height=max(300, row_height * len(samples) + 100),
         )
 
-        if show:
+        if show:  # pragma: no cover
             fig.show(renderer=renderer)
             return None
         else:

malariagen_data/anoph/genome_features.py

@@ -277,7 +277,7 @@ def plot_transcript(
         fig.yaxis.ticker = []
         self._bokeh_style_genome_xaxis(fig, parent.contig)
 
-        if show:
+        if show:  # pragma: no cover
             bokeh.plotting.show(fig)
             return None
         else:
@@ -383,7 +383,7 @@ def plot_genes(
         fig.yaxis.axis_label = "Genes"
         self._bokeh_style_genome_xaxis(fig, contig)
 
-        if show:
+        if show:  # pragma: no cover
             bokeh.plotting.show(fig)
             return None
         else:

malariagen_data/anoph/sample_metadata.py

@@ -4,10 +4,11 @@
 import ipyleaflet
 import numpy as np
 import pandas as pd
+import plotly.express as px
 from numpydoc_decorator import doc
 
 from ..util import check_types
-from . import base_params, map_params
+from . import base_params, map_params, plotly_params
 from .base import AnophelesBase
 
 
@@ -712,3 +713,85 @@ def _results_cache_add_analysis_params(self, params: dict):
         super()._results_cache_add_analysis_params(params)
         params["cohorts_analysis"] = self._cohorts_analysis
         params["aim_analysis"] = self._aim_analysis
+
+    @doc(
+        summary="""
+            Plot a bar chart showing the number of samples available, grouped by
+            some variable such as country or year.
+        """,
+        parameters=dict(
+            x="Name of sample metadata column to plot on the X axis.",
+            color="Name of the sample metadata column to color bars by.",
+            sort="If True, sort the bars in size order.",
+            kwargs="Passed through to px.bar().",
+        ),
+    )
+    def plot_samples_bar(
+        self,
+        x: str,
+        color: Optional[str] = None,
+        sort: bool = True,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        sample_query: Optional[base_params.sample_query] = None,
+        template: plotly_params.template = "plotly_white",
+        width: plotly_params.width = 800,
+        height: plotly_params.height = 600,
+        show: plotly_params.show = True,
+        renderer: plotly_params.renderer = None,
+        **kwargs,
+    ) -> plotly_params.figure:
+        # Load sample metadata.
+        df_samples = self.sample_metadata(
+            sample_sets=sample_sets, sample_query=sample_query
+        )
+
+        # Special handling for plotting by year.
+        if x == "year":
+            # Remove samples with missing year.
+            df_samples = df_samples.query("year > 0")
+
+        # Construct a long-form dataframe to plot.
+        if color:
+            grouper: Union[str, List[str]] = [x, color]
+        else:
+            grouper = x
+        df_plot = df_samples.groupby(grouper).agg({"sample_id": "count"}).reset_index()
+
+        # Deal with request to sort by bar size.
+        if sort:
+            df_sort = (
+                df_samples.groupby(x)
+                .agg({"sample_id": "count"})
+                .reset_index()
+                .sort_values("sample_id")
+            )
+            x_order = df_sort[x].values
+            category_orders = kwargs.get("category_orders", dict())
+            category_orders.setdefault(x, x_order)
+            kwargs["category_orders"] = category_orders
+
+        # Make the plot.
+        fig = px.bar(
+            df_plot,
+            x=x,
+            y="sample_id",
+            color=color,
+            template=template,
+            width=width,
+            height=height,
+            **kwargs,
+        )
+
+        # Visual styling.
+        fig.update_layout(
+            xaxis_title=x.capitalize(),
+            yaxis_title="No. samples",
+        )
+        if color:
+            fig.update_layout(legend_title=color.capitalize())
+
+        if show:  # pragma: no cover
+            fig.show(renderer=renderer)
+            return None
+        else:
+            return fig

malariagen_data/anoph/snp_data.py

@@ -1068,7 +1068,7 @@ def plot_snps(
             sizing_mode=sizing_mode,
         )
 
-        if show:
+        if show:  # pragma: no cover
             bokeh.plotting.show(fig)
             return None
         else:
@@ -1250,7 +1250,7 @@ def plot_snps_track(
         fig.xaxis.minor_tick_line_color = None
         fig.xaxis[0].formatter = bokeh.models.NumeralTickFormatter(format="0,0")
 
-        if show:
+        if show:  # pragma: no cover
             bokeh.plotting.show(fig)
             return None
         else: