Updated controls to both Irida and PCR.

src/submissions/frontend/visualizations/irida_charts.py

@@ -0,0 +1,153 @@
+"""
+Functions for constructing irida controls graphs using plotly.
+"""
+from datetime import date
+from pprint import pformat
+import plotly
+import plotly.express as px
+import pandas as pd
+from PyQt6.QtWidgets import QWidget
+from . import CustomFigure
+import logging
+from tools import get_unique_values_in_df_column, divide_chunks
+from frontend.widgets.functions import select_save_file
+
+logger = logging.getLogger(f"submissions.{__name__}")
+
+
+class IridaFigure(CustomFigure):
+
+    def __init__(self, df: pd.DataFrame, modes: list, ytitle: str | None = None, parent: QWidget | None = None,
+                 months: int = 6):
+
+        super().__init__(df=df, modes=modes)
+
+        self.construct_chart(df=df, modes=modes)
+        self.generic_figure_markers(modes=modes, ytitle=ytitle, months=months)
+
+    def construct_chart(self, df: pd.DataFrame, modes: list):
+        """
+        Creates a plotly chart for controls from a pandas dataframe
+
+        Args:
+            df (pd.DataFrame): input dataframe of controls
+            modes (list): analysis modes to construct charts for
+            ytitle (str | None, optional): title on the y-axis. Defaults to None.
+
+        Returns:
+            Figure: output stacked bar chart.
+        """
+        # fig = Figure()
+        for ii, mode in enumerate(modes):
+            if "count" in mode:
+                df[mode] = pd.to_numeric(df[mode], errors='coerce')
+                color = "genus"
+                color_discrete_sequence = None
+            elif 'percent' in mode:
+                color = "genus"
+                color_discrete_sequence = None
+            else:
+                color = "target"
+                match get_unique_values_in_df_column(df, 'target'):
+                    case ['Target']:
+                        color_discrete_sequence = ["blue"]
+                    case ['Off-target']:
+                        color_discrete_sequence = ['red']
+                    case _:
+                        color_discrete_sequence = ['blue', 'red']
+            bar = px.bar(df,
+                         x="submitted_date",
+                         y=mode,
+                         color=color,
+                         title=mode,
+                         barmode='stack',
+                         hover_data=["genus", "name", "target", mode],
+                         text="genera",
+                         color_discrete_sequence=color_discrete_sequence
+                         )
+            bar.update_traces(visible=ii == 0)
+            self.add_traces(bar.data)
+        # return generic_figure_markers(modes=modes, ytitle=ytitle)
+
+    def generic_figure_markers(self, modes: list = [], ytitle: str | None = None, months: int = 6):
+        """
+        Adds standard layout to figure.
+
+        Args:
+            fig (Figure): Input figure.
+            modes (list, optional): List of modes included in figure. Defaults to [].
+            ytitle (str, optional): Title for the y-axis. Defaults to None.
+
+        Returns:
+            Figure: Output figure with updated titles, rangeslider, buttons.
+        """
+        if modes:
+            ytitle = modes[0]
+        # Creating visibles list for each mode.
+        self.update_layout(
+            xaxis_title="Submitted Date (* - Date parsed from fastq file creation date)",
+            yaxis_title=ytitle,
+            showlegend=True,
+            barmode='stack',
+            updatemenus=[
+                dict(
+                    type="buttons",
+                    direction="right",
+                    x=0.7,
+                    y=1.2,
+                    showactive=True,
+                    buttons=[button for button in self.make_pyqt_buttons(modes=modes)],
+                )
+            ]
+        )
+
+        self.update_xaxes(
+            rangeslider_visible=True,
+            rangeselector=dict(
+                buttons=[button for button in self.make_plotly_buttons(months=months)]
+            )
+        )
+        assert isinstance(self, CustomFigure)
+
+    def make_plotly_buttons(self, months: int = 6):
+        rng = [1]
+        if months > 2:
+            rng += [iii for iii in range(3, months, 3)]
+        # logger.debug(f"Making buttons for months: {rng}")
+        buttons = [dict(count=iii, label=f"{iii}m", step="month", stepmode="backward") for iii in rng]
+        if months > date.today().month:
+            buttons += [dict(count=1, label="YTD", step="year", stepmode="todate")]
+        buttons += [dict(step="all")]
+        for button in buttons:
+            yield button
+
+    def make_pyqt_buttons(self, modes: list) -> list:
+        """
+        Creates list of buttons with one for each mode to be used in showing/hiding mode traces.
+
+        Args:
+            modes (list): list of modes used by main parser.
+            fig_len (int): number of traces in the figure
+
+        Returns:
+            list: list of buttons.
+        """
+        fig_len = len(self.data)
+        if len(modes) > 1:
+            for ii, mode in enumerate(modes):
+                # NOTE: What I need to do is create a list of bools with the same length as the fig.data
+                mode_vis = [True] * fig_len
+                # NOTE: And break it into {len(modes)} chunks
+                mode_vis = list(divide_chunks(mode_vis, len(modes)))
+                # NOTE: Then, for each chunk, if the chunk index isn't equal to the index of the current mode, set to false
+                for jj, sublist in enumerate(mode_vis):
+                    if jj != ii:
+                        mode_vis[jj] = [not elem for elem in mode_vis[jj]]
+                # NOTE: Finally, flatten list.
+                mode_vis = [item for sublist in mode_vis for item in sublist]
+                # NOTE: Now, yield button to add to list
+                yield dict(label=mode, method="update", args=[
+                    {"visible": mode_vis},
+                    {"yaxis.title.text": mode},
+                ])
+