pg-sui 0.2.0__py3-none-any.whl → 1.6.14.dev9__py3-none-any.whl

pgsui/utils/plotting.py

@@ -1,949 +1,1116 @@
-import os
-import sys
-from itertools import cycle
+import logging
+import warnings
 from pathlib import Path
+from typing import Dict, List, Literal, Optional, Sequence, cast
 
+import matplotlib as mpl
+
+# Use Agg backend for headless plotting
+mpl.use("Agg")
+
+import matplotlib.pyplot as plt
 import numpy as np
+import optuna
 import pandas as pd
-import matplotlib.pyplot as plt
 import seaborn as sns
-import plotly.express as px
-
-from sklearn.decomposition import PCA
-from sklearn.preprocessing import StandardScaler
-from sklearn_genetic.utils import logbook_to_pandas
-from sklearn.metrics import ConfusionMatrixDisplay
-
-try:
-    from . import misc
-except (ModuleNotFoundError, ValueError, ImportError):
-    from utils import misc
+import torch
+from optuna.exceptions import ExperimentalWarning
+from sklearn.metrics import (
+    ConfusionMatrixDisplay,
+    auc,
+    average_precision_score,
+    confusion_matrix,
+    precision_recall_curve,
+    roc_curve,
+)
+from sklearn.preprocessing import label_binarize
+from snpio import SNPioMultiQC
+from snpio.utils.logging import LoggerManager
+
+from pgsui.utils import misc
+from pgsui.utils.logging_utils import configure_logger
+
+# Quiet Matplotlib/fontTools INFO logging when saving PDF/SVG
+for name in (
+    "fontTools",
+    "fontTools.subset",
+    "fontTools.ttLib",
+    "matplotlib.font_manager",
+):
+    lg = logging.getLogger(name)
+    lg.setLevel(logging.WARNING)
+    lg.propagate = False
 
 
 class Plotting:
-    """Functions for plotting imputer scoring and results."""
-
-    @staticmethod
-    def plot_grid_search(cv_results, nn_method, prefix):
-        """Plot cv_results\_ from a grid search for each parameter.
-
-        Saves a figure to disk.
+    """Class for plotting imputer scoring and results.
+
+    This class is used to plot the performance metrics of imputation models. It can plot ROC and Precision-Recall curves, model history, and the distribution of genotypes in the dataset.
+
+    Example:
+        >>> from pgsui import Plotting
+        >>> plotter = Plotting(model_name="ImputeVAE", prefix="pgsui_test", plot_format="png")
+        >>> plotter.plot_metrics(metrics, num_classes)
+        >>> plotter.plot_history(history)
+        >>> plotter.plot_confusion_matrix(y_true_1d, y_pred_1d)
+        >>> plotter.plot_tuning(study, model_name, optimize_dir, target_name="Objective Value")
+        >>> plotter.plot_gt_distribution(df)
+
+    Attributes:
+        model_name (str): Name of the model.
+        prefix (str): Prefix for the output directory.
+        plot_format (Literal["pdf", "png", "jpeg", "jpg", "svg"]): Format for the plots ('pdf', 'png', 'jpeg', 'jpg', 'svg').
+        plot_fontsize (int): Font size for the plots.
+        plot_dpi (int): Dots per inch for the plots.
+        title_fontsize (int): Font size for the plot titles.
+        show_plots (bool): Whether to display the plots inline or during execution.
+        output_dir (Path): Directory where plots will be saved.
+        logger (logging.Logger): Logger instance for logging messages.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        *,
+        prefix: str = "pgsui",
+        plot_format: Literal["pdf", "png", "jpeg", "jpg", "svg"] = "pdf",
+        plot_fontsize: int = 18,
+        plot_dpi: int = 300,
+        title_fontsize: int = 20,
+        despine: bool = True,
+        show_plots: bool = False,
+        verbose: int = 0,
+        debug: bool = False,
+        multiqc: bool = False,
+        multiqc_section: Optional[str] = None,
+    ) -> None:
+        """Initialize the Plotting object.
+
+        This class is used to plot the performance metrics of imputation models. It can plot ROC and Precision-Recall curves, model history, and the distribution of genotypes in the dataset.
 
         Args:
-            cv_results (np.ndarray): the cv_results\_ attribute from a trained grid search object.
-
-            nn_method (str): Neural network algorithm name.
-
-            prefix (str): Prefix to use for saving the plot to file.
+            model_name (str): Name of the model.
+            prefix (str, optional): Prefix for the output directory. Defaults to 'pgsui'.
+            plot_format (Literal["pdf", "png", "jpeg", "jpg"]): Format for the plots ('pdf', 'png', 'jpeg', 'jpg'). Defaults to 'pdf'.
+            plot_fontsize (int): Font size for the plots. Defaults to 18.
+            plot_dpi (int): Dots per inch for the plots. Defaults to 300.
+            title_fontsize (int): Font size for the plot titles. Defaults to 20.
+            despine (bool): Whether to remove the top and right spines from the plots. Defaults to True.
+            show_plots (bool): Whether to display the plots. Defaults to False.
+            verbose (int): Verbosity level for logging. Defaults to 0.
+            debug (bool): Whether to enable debug mode. Defaults to False.
+            multiqc (bool): Whether to queue plots for a MultiQC HTML report. Defaults to False.
+            multiqc_section (Optional[str]): Section name to use in MultiQC. Defaults to 'PG-SUI (<model_name>)'.
         """
-        ## Results from grid search
-        results = pd.DataFrame(cv_results)
-        means_test = [col for col in results if col.startswith("mean_test_")]
-        filter_col = [col for col in results if col.startswith("param_")]
-        params_df = results[filter_col].astype(str)
-        for i, col in enumerate(means_test):
-            params_df[col] = results[means_test[i]]
-
-        # Get number of needed subplot rows.
-        tot = len(filter_col)
-        cols = 4
-        rows = int(np.ceil(tot / cols))
-
-        fig = plt.figure(1, figsize=(20, 10))
-        fig.tight_layout(pad=3.0)
-
-        # Set font properties.
-        font = {"size": 12}
-        plt.rc("font", **font)
-
-        for i, p in enumerate(filter_col, start=1):
-            ax = fig.add_subplot(rows, cols, i)
-
-            # Plot each metric.
-            for col in means_test:
-                # Get maximum score for each parameter setting.
-                df_plot = params_df.groupby(p)[col].agg("max")
-
-                # Convert to float if not supposed to be string.
-                try:
-                    df_plot.index = df_plot.index.astype(float)
-                except TypeError:
-                    pass
-
-                # Sort by index (numerically if possible).
-                df_plot = df_plot.sort_index()
-
-                # Remove prefix from score name.
-                col_new_name = col[len("mean_test_") :]
-
-                ax.plot(
-                    df_plot.index.astype(str),
-                    df_plot.values,
-                    "-o",
-                    label=col_new_name,
-                )
+        logman = LoggerManager(
+            name=__name__, prefix=prefix, verbose=bool(verbose), debug=bool(debug)
+        )
+        self.logger = configure_logger(
+            logman.get_logger(), verbose=bool(verbose), debug=bool(debug)
+        )
 
-                ax.legend(loc="best")
+        self.model_name = model_name
+        self.prefix = prefix
+        self.plot_format = plot_format
+        self.plot_fontsize = plot_fontsize
+        self.plot_dpi = plot_dpi
+        self.title_fontsize = title_fontsize
+        self.show_plots = show_plots
 
-            param_new_name = p[len("param_") :]
-            ax.set_xlabel(param_new_name.lower())
-            ax.set_ylabel("Max Score")
-            ax.set_ylim([0, 1])
+        # MultiQC configuration
+        self.use_multiqc: bool = bool(multiqc)
 
-        fig.savefig(
-            os.path.join(
-                f"{prefix}_output",
-                "plots",
-                "Unsupervised",
-                nn_method,
-                "gridsearch_metrics.pdf",
-            ),
-            bbox_inches="tight",
-            facecolor="white",
+        self.multiqc_section: str = (
+            multiqc_section if multiqc_section is not None else f"PG-SUI ({model_name})"
         )
 
-    @staticmethod
-    def plot_metrics(metrics, num_classes, prefix, nn_method):
-        """Plot AUC-ROC and Precision-Recall performance metrics for neural network classifier.
+        if self.plot_format.startswith("."):
+            self.plot_format = self.plot_format.lstrip(".")
+
+        self.param_dict = {
+            "axes.labelsize": self.plot_fontsize,
+            "axes.titlesize": self.title_fontsize,
+            "axes.spines.top": despine,
+            "axes.spines.right": despine,
+            "xtick.labelsize": self.plot_fontsize,
+            "ytick.labelsize": self.plot_fontsize,
+            "legend.fontsize": self.plot_fontsize,
+            "legend.facecolor": "white",
+            "figure.titlesize": self.title_fontsize,
+            "figure.dpi": self.plot_dpi,
+            "figure.facecolor": "white",
+            "axes.linewidth": 2.0,
+            "lines.linewidth": 2.0,
+            "font.size": self.plot_fontsize,
+            "savefig.bbox": "tight",
+            "savefig.facecolor": "white",
+            "savefig.dpi": self.plot_dpi,
+        }
 
-        Saves plot to PDF file on disk.
+        mpl.rcParams.update(self.param_dict)
 
-        Args:
-            metrics (Dict[str, Any]): Per-class, micro, and macro-averaged metrics including accuracy, ROC-AUC, and Precision-Recall with Average Precision scores.
+        unsuper = {"ImputeVAE", "ImputeNLPCA", "ImputeAutoencoder", "ImputeUBP"}
 
-            num_classes (int): Number of classes evaluated.
+        det = {
+            "ImputeRefAllele",
+            "ImputeMostFrequent",
+            "ImputeMostFrequentPerPop",
+            "ImputePhylo",
+        }
+
+        sup = {"ImputeRandomForest", "ImputeHistGradientBoosting"}
 
-            prefix (str): Prefix to use for output plot.
+        if model_name in unsuper:
+            plot_dir = "Unsupervised"
+        elif model_name in det:
+            plot_dir = "Deterministic"
+        elif model_name in sup:
+            plot_dir = "Supervised"
+        else:
+            msg = f"model_name '{model_name}' not recognized."
+            self.logger.error(msg)
+            raise ValueError(msg)
+
+        self.output_dir = Path(f"{self.prefix}_output", plot_dir)
+        self.output_dir = self.output_dir / "plots" / model_name
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+    # --------------------------------------------------------------------- #
+    # Core plotting methods                                                #
+    # --------------------------------------------------------------------- #
+    def plot_tuning(
+        self,
+        study: optuna.study.Study,
+        model_name: str,
+        optimize_dir: Path,
+        target_name: str = "Objective Value",
+    ) -> None:
+        """Plot the optimization history of a study.
+
+        This method plots the optimization history of a study. The plot is saved to disk as a ``<plot_format>`` file.
 
-            nn_method (str): Neural network algorithm being used.
+        Args:
+            study (optuna.study.Study): Optuna study object.
+            model_name (str): Name of the model.
+            optimize_dir (Path): Directory to save the optimization plots.
+            target_name (str, optional): Name of the target value. Defaults to 'Objective Value'.
         """
-        # Set font properties.
-        font = {"size": 12}
-        plt.rc("font", **font)
-
-        fn = os.path.join(
-            f"{prefix}_output",
-            "plots",
-            "Unsupervised",
-            nn_method,
-            f"auc_pr_curves.pdf",
-        )
-        fig = plt.figure(figsize=(20, 10))
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning)
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            warnings.filterwarnings("ignore", category=ExperimentalWarning)
 
-        acc = round(metrics["accuracy"] * 100, 2)
-        ham = round(metrics["hamming"], 2)
+            target_name = target_name.title()
 
-        fig.suptitle(
-            f"Performance Metrics\nAccuracy: {acc}\nHamming Loss: {ham}"
-        )
-        axs = fig.subplots(nrows=1, ncols=2)
-        plt.subplots_adjust(hspace=0.5)
-
-        # Line weight
-        lw = 2
-
-        roc_auc = metrics["roc_auc"]
-        pr_ap = metrics["precision_recall"]
-
-        metric_list = [roc_auc, pr_ap]
-
-        for metric, ax in zip(metric_list, axs):
-            if "fpr_micro" in metric:
-                prefix1 = "fpr"
-                prefix2 = "tpr"
-                lab1 = "ROC"
-                lab2 = "AUC"
-                xlab = "False Positive Rate"
-                ylab = "True Positive Rate"
-                title = "Receiver Operating Characteristic (ROC)"
-                baseline = [0, 1]
-
-            elif "recall_micro" in metric:
-                prefix1 = "recall"
-                prefix2 = "precision"
-                lab1 = "Precision-Recall"
-                lab2 = "AP"
-                xlab = "Recall"
-                ylab = "Precision"
-                title = "Precision-Recall"
-                baseline = [metric["baseline"], metric["baseline"]]
-
-                # Plot iso-f1 curves.
-                f_scores = np.linspace(0.2, 0.8, num=4)
-                for i, f_score in enumerate(f_scores):
-                    x = np.linspace(0.01, 1)
-                    y = f_score * x / (2 * x - f_score)
-                    ax.plot(
-                        x[y >= 0],
-                        y[y >= 0],
-                        color="gray",
-                        alpha=0.2,
-                        linewidth=lw,
-                        label="Iso-F1 Curves" if i == 0 else "",
-                    )
-                    ax.annotate(f"F1={f_score:0.1f}", xy=(0.9, y[45] + 0.02))
-
-            # Plot ROC curves.
-            ax.plot(
-                metric[f"{prefix1}_micro"],
-                metric[f"{prefix2}_micro"],
-                label=f"Micro-averaged {lab1} Curve ({lab2} = {metric['micro']:.2f})",
-                color="deeppink",
-                linestyle=":",
-                linewidth=4,
+            ax = optuna.visualization.matplotlib.plot_optimization_history(
+                study, target_name=target_name
             )
-
-            ax.plot(
-                metric[f"{prefix1}_macro"],
-                metric[f"{prefix2}_macro"],
-                label=f"Macro-averaged {lab1} Curve ({lab2} = {metric['macro']:.2f})",
-                color="navy",
-                linestyle=":",
-                linewidth=4,
+            ax.set_title(f"{model_name} Optimization History")
+            ax.set_xlabel("Trial")
+            ax.set_ylabel(target_name)
+            ax.legend(
+                loc="best",
+                shadow=True,
+                fancybox=True,
+                fontsize=mpl.rcParamsDefault["legend.fontsize"],
             )
 
-            colors = cycle(["aqua", "darkorange", "cornflowerblue"])
-            for i, color in zip(range(num_classes), colors):
-                if f"{prefix1}_{i}" in metric:
-                    ax.plot(
-                        metric[f"{prefix1}_{i}"],
-                        metric[f"{prefix2}_{i}"],
-                        color=color,
-                        lw=lw,
-                        label=f"{lab1} Curve of class {i} ({lab2} = {metric[i]:.2f})",
-                    )
-
-            if "fpr_micro" in metric:
-                # Make center baseline
-                ax.plot(
-                    baseline,
-                    baseline,
-                    "k--",
-                    linewidth=lw,
-                    label="No Classification Skill",
-                )
-            else:
-                ax.plot(
-                    [0, 1],
-                    baseline,
-                    "k--",
-                    linewidth=lw,
-                    label="No Classification Skill",
-                )
-
-            ax.set_xlim(0.0, 1.0)
-            ax.set_ylim(0.0, 1.05)
-            ax.set_xlabel(f"{xlab}")
-            ax.set_ylabel(f"{ylab}")
-            ax.set_title(f"{title}")
-            ax.legend(loc="best")
-
-        fig.savefig(fn, bbox_inches="tight", facecolor="white")
-        plt.close()
-        plt.clf()
-        plt.cla()
-
-    @staticmethod
-    def plot_search_space(
-        estimator,
-        height=2,
-        s=25,
-        features=None,
-    ):
-        """Make density and contour plots for showing search space during grid search.
-
-        Modified from sklearn-genetic-opt function to implement exception handling.
-
-        Args:
-            estimator (sklearn estimator object): A fitted estimator from :class:`~sklearn_genetic.GASearchCV`.
-
-            height (float, optional): Height of each facet. Defaults to 2.
-
-            s (float, optional): Size of the markers in scatter plot. Defaults to 5.
-
-            features (list, optional): Subset of features to plot, if ``None`` it plots all the features by default. Defaults to None.
+            od = optimize_dir
+            fn = od / f"optuna_optimization_history.{self.plot_format}"
 
-        Returns:
-            g (seaborn.PairGrid): Pair plot of the used hyperparameters during the search.
-        """
-        sns.set_style("white")
-
-        df = logbook_to_pandas(estimator.logbook)
-        if features:
-            _stats = df[features]
-        else:
-            variables = [*estimator.space.parameters, "score"]
-            _stats = df[variables]
+            if not fn.parent.exists():
+                fn.parent.mkdir(parents=True, exist_ok=True)
 
-        g = sns.PairGrid(_stats, diag_sharey=False, height=height)
-
-        g = g.map_upper(sns.scatterplot, s=s, color="r", alpha=0.2)
+            plt.savefig(fn)
+            plt.close()
 
-        try:
-            g = g.map_lower(
-                sns.kdeplot,
-                shade=True,
-                cmap=sns.color_palette("ch:s=.25,rot=-.25", as_cmap=True),
+            ax = optuna.visualization.matplotlib.plot_edf(
+                study, target_name=target_name
             )
-        except np.linalg.LinAlgError as err:
-            if "singular matrix" in str(err).lower():
-                g = g.map_lower(sns.scatterplot, s=s, color="b", alpha=1.0)
-            else:
-                raise
-
-        try:
-            g = g.map_diag(
-                sns.kdeplot,
-                shade=True,
-                palette="crest",
-                alpha=0.2,
-                color="red",
+            ax.set_title(f"{model_name} Empirical Distribution Function (EDF)")
+            ax.set_xlabel(target_name)
+            ax.set_ylabel(f"{model_name} Cumulative Probability")
+            ax.legend(
+                loc="best",
+                shadow=True,
+                fancybox=True,
+                fontsize=mpl.rcParamsDefault["legend.fontsize"],
             )
-        except np.linalg.LinAlgError as err:
-            if "singular matrix" in str(err).lower():
-                g = g.map_diag(sns.histplot, color="red", alpha=1.0, kde=False)
-
-        return g
-
-    @staticmethod
-    def visualize_missingness(
-        genotype_data,
-        df,
-        zoom=True,
-        prefix="imputer",
-        horizontal_space=0.6,
-        vertical_space=0.6,
-        bar_color="gray",
-        heatmap_palette="magma",
-        plot_format="pdf",
-        dpi=300,
-    ):
-        """Make multiple plots to visualize missing data.
-
-        Args:
-            genotype_data (GenotypeData): Initialized GentoypeData object.
 
-            df (pandas.DataFrame): DataFrame with snps to visualize.
+            plt.savefig(fn.with_stem("optuna_edf_plot"))
+            plt.close()
 
-            zoom (bool, optional): If True, zooms in to the missing proportion range on some of the plots. If False, the plot range is fixed at [0, 1]. Defaults to True.
+            ax = optuna.visualization.matplotlib.plot_param_importances(
+                study, target_name=target_name
+            )
+            ax.set_xlabel("Parameter Importance")
+            ax.set_ylabel("Parameter")
+            ax.legend(loc="best", shadow=True, fancybox=True)
 
-            prefix (str, optional): Prefix for output directory and files. Plots and files will be written to a directory called <prefix>_reports. The report directory will be created if it does not already exist. If prefix is None, then the reports directory will not have a prefix. Defaults to 'imputer'.
+            plt.savefig(fn.with_stem("optuna_param_importances_plot"))
+            plt.close()
 
-            horizontal_space (float, optional): Set width spacing between subplots. If your plot are overlapping horizontally, increase horizontal_space. If your plots are too far apart, decrease it. Defaults to 0.6.
+            ax = optuna.visualization.matplotlib.plot_timeline(study)
+            ax.set_title(f"{model_name} Timeline Plot")
+            ax.set_xlabel("Datetime")
+            ax.set_ylabel("Trial")
+            plt.savefig(fn.with_stem("optuna_timeline_plot"))
+            plt.close()
 
-            vertical_space (float, optioanl): Set height spacing between subplots. If your plots are overlapping vertically, increase vertical_space. If your plots are too far apart, decrease it. Defaults to 0.6.
+            # Reset the style from Optuna's plotting.
+            sns.set_style("white", rc=self.param_dict)
+            mpl.rcParams.update(self.param_dict)
 
-            bar_color (str, optional): Color of the bars on the non-stacked barplots. Can be any color supported by matplotlib. See matplotlib.pyplot.colors documentation. Defaults to 'gray'.
+        # ---- MultiQC: Optuna tuning line graph + best-params table --------
+        if self._multiqc_enabled():
+            try:
+                self._queue_multiqc_tuning(
+                    study=study, model_name=model_name, target_name=target_name
+                )
+            except Exception as exc:  # pragma: no cover - defensive
+                self.logger.warning(f"Failed to queue MultiQC tuning plots: {exc}")
 
-            heatmap_palette (str, optional): Palette to use for heatmap plot. Can be any palette supported by seaborn. See seaborn documentation. Defaults to 'magma'.
+    def plot_metrics(
+        self,
+        y_true: np.ndarray,
+        y_pred_proba: np.ndarray,
+        metrics: Dict[str, float],
+        label_names: Optional[Sequence[str]] = None,
+        prefix: str = "",
+    ) -> None:
+        """Plot multi-class ROC-AUC and Precision-Recall curves.
 
-            plot_format (str, optional): Format to save plots. Can be any of the following: "pdf", "png", "svg", "ps", "eps". Defaults to "pdf".
+        This method plots the multi-class ROC-AUC and Precision-Recall curves. The plot is saved to disk as a ``<plot_format>`` file.
 
-            dpi (int): The resolution in dots per inch. Defaults to 300.
+        Args:
+            y_true (np.ndarray): 1D array of true integer labels in [0, n_classes-1].
+            y_pred_proba (np.ndarray): (n_samples, n_classes) array of predicted probabilities.
+            metrics (Dict[str, float]): Dict of summary metrics to annotate the figure.
+            label_names (Optional[Sequence[str]]): Optional sequence of class names (length must equal n_classes).
+                If provided, legends will use these names instead of 'Class i'.
+            prefix (str): Optional prefix for the output filename.
 
-        Returns:
-            pandas.DataFrame: Per-locus missing data proportions.
-            pandas.DataFrame: Per-individual missing data proportions.
-            pandas.DataFrame: Per-population + per-locus missing data proportions.
-            pandas.DataFrame: Per-population missing data proportions.
-            pandas.DataFrame: Per-individual and per-population missing data proportions.
+        Raises:
+            ValueError: If model_name is not recognized (legacy guard).
         """
+        num_classes = y_pred_proba.shape[1]
 
-        loc, ind, poploc, poptotal, indpop = genotype_data.calc_missing(df)
-
-        ncol = 3
-        nrow = 1 if genotype_data.pops is None else 2
-
-        fig, axes = plt.subplots(nrow, ncol, figsize=(8, 11))
-        plt.subplots_adjust(wspace=horizontal_space, hspace=vertical_space)
-        fig.suptitle("Missingness Report")
-
-        ax = axes[0, 0]
-
-        ax.set_title("Per-Individual")
-        ax.barh(genotype_data.samples, ind, color=bar_color, height=1.0)
-        if not zoom:
-            ax.set_xlim([0, 1])
-        ax.set_ylabel("Sample")
-        ax.set_xlabel("Missing Prop.")
-        ax.tick_params(
-            axis="y",
-            which="both",
-            left=False,
-            right=False,
-            labelleft=False,
-        )
-
-        ax = axes[0, 1]
+        # Validate/normalize label names
+        if label_names is not None and len(label_names) != num_classes:
+            self.logger.warning(
+                f"plot_metrics: len(label_names)={len(label_names)} "
+                f"!= n_classes={num_classes}. Ignoring label_names."
+            )
+            label_names = None
+        if label_names is None:
+            label_names = [f"Class {i}" for i in range(num_classes)]
+
+        # Binarize y_true for one-vs-rest curves
+        y_true_bin = np.asarray(label_binarize(y_true, classes=np.arange(num_classes)))
+
+        # Containers
+        fpr, tpr, roc_auc_vals = {}, {}, {}
+        precision, recall, average_precision_vals = {}, {}, {}
+
+        # Per-class ROC & PR
+        for i in range(num_classes):
+            fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_pred_proba[:, i])
+            roc_auc_vals[i] = auc(fpr[i], tpr[i])
+            precision[i], recall[i], _ = precision_recall_curve(
+                y_true_bin[:, i], y_pred_proba[:, i]
+            )
+            average_precision_vals[i] = average_precision_score(
+                y_true_bin[:, i], y_pred_proba[:, i]
+            )
 
-        ax.set_title("Per-Locus")
-        ax.barh(
-            range(genotype_data.num_snps), loc, color=bar_color, height=1.0
+        # Micro-averages
+        fpr["micro"], tpr["micro"], _ = roc_curve(
+            y_true_bin.ravel(), y_pred_proba.ravel()
         )
-        if not zoom:
-            ax.set_xlim([0, 1])
-        ax.set_ylabel("Locus")
-        ax.set_xlabel("Missing Prop.")
-        ax.tick_params(
-            axis="y",
-            which="both",
-            left=False,
-            right=False,
-            labelleft=False,
+        roc_auc_vals["micro"] = auc(fpr["micro"], tpr["micro"])
+        precision["micro"], recall["micro"], _ = precision_recall_curve(
+            y_true_bin.ravel(), y_pred_proba.ravel()
+        )
+        average_precision_vals["micro"] = average_precision_score(
+            y_true_bin, y_pred_proba, average="micro"
         )
 
-        id_vars = ["SampleID"]
-        if poptotal is not None:
-            ax = axes[0, 2]
-
-            ax.set_title("Per-Population Total")
-            ax.barh(poptotal.index, poptotal, color=bar_color, height=1.0)
-            if not zoom:
-                ax.set_xlim([0, 1])
-            ax.set_xlabel("Missing Prop.")
-            ax.set_ylabel("Population")
-
-            ax = axes[1, 0]
-
-            ax.set_title("Per-Population + Per-Locus")
-            npops = len(poploc.columns)
+        # Macro-average ROC
+        all_fpr = np.unique(np.concatenate([fpr[i] for i in range(num_classes)]))
+        mean_tpr = np.zeros_like(all_fpr)
+        for i in range(num_classes):
+            mean_tpr += np.interp(all_fpr, fpr[i], tpr[i])
+        mean_tpr /= num_classes
+        fpr["macro"], tpr["macro"] = all_fpr, mean_tpr
+        roc_auc_vals["macro"] = auc(fpr["macro"], tpr["macro"])
+
+        # Macro-average PR
+        all_recall = np.unique(np.concatenate([recall[i] for i in range(num_classes)]))
+        mean_precision = np.zeros_like(all_recall)
+        for i in range(num_classes):
+            # recall[i] increases, but precision[i] is given over decreasing thresholds
+            mean_precision += np.interp(all_recall, recall[i][::-1], precision[i][::-1])
+        mean_precision /= num_classes
+        average_precision_vals["macro"] = average_precision_score(
+            y_true_bin, y_pred_proba, average="macro"
+        )
 
-            vmax = None if zoom else 1.0
+        # Plot
+        fig, axes = plt.subplots(1, 2, figsize=(14, 6))
 
-            sns.heatmap(
-                poploc,
-                vmin=0.0,
-                vmax=vmax,
-                cmap=sns.color_palette(heatmap_palette, as_cmap=True),
-                yticklabels=False,
-                cbar_kws={"label": "Missing Prop."},
-                ax=ax,
+        # ROC
+        axes[0].plot(
+            fpr["micro"],
+            tpr["micro"],
+            label=f"Micro-average ROC (AUC = {roc_auc_vals['micro']:.2f})",
+            linestyle=":",
+            linewidth=4,
+        )
+        axes[0].plot(
+            fpr["macro"],
+            tpr["macro"],
+            label=f"Macro-average ROC (AUC = {roc_auc_vals['macro']:.2f})",
+            linestyle="--",
+            linewidth=4,
+        )
+        for i in range(num_classes):
+            axes[0].plot(
+                fpr[i],
+                tpr[i],
+                label=f"{label_names[i]} ROC (AUC = {roc_auc_vals[i]:.2f})",
             )
-            ax.set_xlabel("Population")
-            ax.set_ylabel("Locus")
-
-            id_vars.append("Population")
-
-        melt_df = indpop.isna()
-        melt_df["SampleID"] = genotype_data.samples
-        indpop["SampleID"] = genotype_data.samples
-
-        if poptotal is not None:
-            melt_df["Population"] = genotype_data.pops
-            indpop["Population"] = genotype_data.pops
-
-        melt_df = melt_df.melt(value_name="Missing", id_vars=id_vars)
-        melt_df.sort_values(by=id_vars[::-1], inplace=True)
-        melt_df["Missing"].replace(False, "Present", inplace=True)
-        melt_df["Missing"].replace(True, "Missing", inplace=True)
-
-        ax = axes[0, 2] if poptotal is None else axes[1, 1]
+        axes[0].plot([0, 1], [0, 1], linestyle="--", color="black", label="Random")
+        axes[0].set_xlabel("False Positive Rate")
+        axes[0].set_ylabel("True Positive Rate")
+        axes[0].set_title("Multi-class ROC-AUC Curve")
+        axes[0].legend(
+            loc="upper center",
+            bbox_to_anchor=(0.5, -0.15),
+            fancybox=True,
+            shadow=True,
+            ncol=2,
+        )
 
-        ax.set_title("Per-Individual")
-        g = sns.histplot(
-            data=melt_df,
-            y="variable",
-            hue="Missing",
-            multiple="fill",
-            ax=ax,
+        # PR
+        axes[1].plot(
+            recall["micro"],
+            precision["micro"],
+            label=f"Micro-average PR (AP = {average_precision_vals['micro']:.2f})",
+            linestyle=":",
+            linewidth=4,
         )
-        ax.tick_params(
-            axis="y",
-            which="both",
-            left=False,
-            right=False,
-            labelleft=False,
+        axes[1].plot(
+            all_recall,
+            mean_precision,
+            label=f"Macro-average PR (AP = {average_precision_vals['macro']:.2f})",
+            linestyle="--",
+            linewidth=4,
         )
-        g.get_legend().set_title(None)
-
-        if poptotal is not None:
-            ax = axes[1, 2]
-
-            ax.set_title("Per-Population")
-            g = sns.histplot(
-                data=melt_df,
-                y="Population",
-                hue="Missing",
-                multiple="fill",
-                ax=ax,
+        for i in range(num_classes):
+            axes[1].plot(
+                recall[i],
+                precision[i],
+                label=f"{label_names[i]} PR (AP = {average_precision_vals[i]:.2f})",
             )
-            g.get_legend().set_title(None)
-
-        fig.savefig(
-            os.path.join(
-                f"{prefix}_output", "plots", f"missingness.{plot_format}"
-            ),
-            bbox_inches="tight",
-            facecolor="white",
+        axes[1].plot([0, 1], [1, 0], linestyle="--", color="black", label="Random")
+        axes[1].set_xlabel("Recall")
+        axes[1].set_ylabel("Precision")
+        axes[1].set_title("Multi-class Precision-Recall Curve")
+        axes[1].legend(
+            loc="upper center",
+            bbox_to_anchor=(0.5, -0.15),
+            fancybox=True,
+            shadow=True,
+            ncol=2,
         )
-        plt.cla()
-        plt.clf()
-        plt.close()
 
-        return loc, ind, poploc, poptotal, indpop
+        # Title & save
+        fig.suptitle("\n".join([f"{k}: {v:.2f}" for k, v in metrics.items()]), y=1.35)
 
-    @staticmethod
-    def run_and_plot_pca(
-        original_genotype_data,
-        imputer_object,
-        prefix="imputer",
-        n_components=3,
-        center=True,
-        scale=False,
-        n_axes=2,
-        point_size=15,
-        font_size=15,
-        plot_format="pdf",
-        bottom_margin=0,
-        top_margin=0,
-        left_margin=0,
-        right_margin=0,
-        width=1088,
-        height=700,
-    ):
-        """Runs PCA and makes scatterplot with colors showing missingness.
-
-        Genotypes are plotted as separate shapes per population and colored according to missingness per individual.
+        prefix_for_name = f"{prefix}_" if prefix != "" else ""
+        out_name = (
+            f"{self.model_name}_{prefix_for_name}roc_pr_curves.{self.plot_format}"
+        )
+        fig.savefig(self.output_dir / out_name, bbox_inches="tight")
+        if self.show_plots:
+            plt.show()
+        plt.close(fig)
+
+        # ---- MultiQC: metrics table + per-class AUC/AP heatmap ------------
+        if self._multiqc_enabled():
+            try:
+                self._queue_multiqc_metrics(
+                    metrics=metrics,
+                    roc_auc=roc_auc_vals,
+                    average_precision=average_precision_vals,
+                    label_names=label_names,
+                    panel_prefix=prefix,
+                )
+            except Exception as exc:  # pragma: no cover - defensive
+                self.logger.warning(f"Failed to queue MultiQC metrics plots: {exc}")
+
+            try:
+                self._queue_multiqc_roc_curves(
+                    fpr=fpr,
+                    tpr=tpr,
+                    label_names=label_names,
+                    panel_prefix=prefix,
+                )
+                self._queue_multiqc_pr_curves(
+                    precision=precision,
+                    recall=recall,
+                    label_names=label_names,
+                    panel_prefix=prefix,
+                )
+            except Exception as exc:  # pragma: no cover - defensive
+                self.logger.warning(f"Failed to queue MultiQC ROC/PR curves: {exc}")
 
-        This function is run at the end of each imputation method, but can be run independently to change plot and PCA parameters such as ``n_axes=3`` or ``scale=True``.
+    def plot_history(
+        self,
+        history: Dict[str, List[float] | Dict[str, List[float]] | None] | None,
+    ) -> None:
+        """Plot model history traces. Will be saved to file.
 
-        The imputed and original GenotypeData objects need to be passed to the function as positional arguments.
+        This method plots the deep learning model history traces. The plot is saved to disk as a ``<plot_format>`` file.
 
-        PCA (principal component analysis) scatterplot can have either two or three axes, set with the n_axes parameter.
+        Args:
+            history (Dict[str, List[float]]): Dictionary with lists of history objects. Keys should be "Train" and "Validation".
 
-        The plot is saved as both an interactive HTML file and as a static image. Each population is represented by point shapes. The interactive plot has associated metadata when hovering over the points.
+        Raises:
+            ValueError: nn_method must be either 'ImputeNLPCA', 'ImputeUBP', 'ImputeAutoencoder', 'ImputeVAE'.
+        """
+        if self.model_name not in {
+            "ImputeNLPCA",
+            "ImputeVAE",
+            "ImputeAutoencoder",
+            "ImputeUBP",
+        }:
+            msg = "nn_method must be either 'ImputeNLPCA', 'ImputeVAE', 'ImputeAutoencoder', 'ImputeUBP'."
+            self.logger.error(msg)
+            raise ValueError(msg)
+
+        if self.model_name != "ImputeUBP":
+            fig, ax = plt.subplots(1, 1, figsize=(12, 8))
+            df = pd.DataFrame(history)
+            df = df.iloc[1:]
 
-        Files are saved to a reports directory as <prefix>_output/imputed_pca.<plot_format|html>. Supported image formats include: "pdf", "svg", "png", and "jpeg" (or "jpg").
+            # Plot train accuracy
+            ax.plot(df["Train"], c="blue", lw=3)
 
-        Args:
-            original_genotype_data (GenotypeData): Original GenotypeData object that was input into the imputer.
+            ax.set_title(f"{self.model_name} Loss per Epoch")
+            ax.set_ylabel("Loss")
+            ax.set_xlabel("Epoch")
+            ax.legend(["Train"], loc="best", shadow=True, fancybox=True)
 
-            imputer_object (Any imputer instance): Imputer object created when imputing. Can be any of the imputers, such as: ``ImputePhylo()``, ``ImputeUBP()``, and ``ImputeRandomForest()``.
+        else:
+            fig, ax = plt.subplots(3, 1, figsize=(12, 8))
+
+            # Ensure history is the nested dictionary type for ImputeUBP
+            if not (
+                isinstance(history, dict)
+                and "Train" in history
+                and isinstance(history["Train"], dict)
+            ):
+                msg = "For ImputeUBP, history must be a nested dictionary with phases."
+                self.logger.error(msg)
+                raise TypeError(msg)
+
+            for i, phase in enumerate(range(1, 4)):
+                train = pd.Series(history["Train"][f"Phase {phase}"])
+                train = train.iloc[1:]  # ignore first epoch
+
+                # Plot train accuracy
+                ax[i].plot(train, c="blue", lw=3)
+                ax[i].set_title(f"{self.model_name}: Phase {phase} Loss per Epoch")
+                ax[i].set_ylabel("Loss")
+                ax[i].set_xlabel("Epoch")
+                ax[i].legend([f"Phase {phase}"], loc="best", shadow=True, fancybox=True)
+
+        fn = f"{self.model_name.lower()}_history_plot.{self.plot_format}"
+        fn = self.output_dir / fn
+        fig.savefig(fn)
+
+        if self.show_plots:
+            plt.show()
+        plt.close(fig)
+
+        # ---- MultiQC: training-loss vs epoch linegraphs -------------------
+        if self._multiqc_enabled():
+            try:
+                self._queue_multiqc_history(history=history)
+            except Exception as exc:  # pragma: no cover
+                self.logger.warning(f"Failed to queue MultiQC history plot: {exc}")
 
-            original_012 (pandas.DataFrame, numpy.ndarray, or List[List[int]], optional): Original 012-encoded genotypes (before imputing). Missing values are encoded as -9. This object can be obtained as ``df = GenotypeData.genotypes012_df``.
+    def plot_confusion_matrix(
+        self,
+        y_true_1d: np.ndarray | pd.DataFrame | List[str | int] | torch.Tensor,
+        y_pred_1d: np.ndarray | pd.DataFrame | List[str | int] | torch.Tensor,
+        label_names: Sequence[str] | Dict[str, int] | None = None,
+        prefix: str = "",
+    ) -> None:
+        """Plot a confusion matrix with optional class labels.
 
-            prefix (str, optional): Prefix for report directory. Plots will be save to a directory called <prefix>_output/imputed_pca<html|plot_format>. Report directory will be created if it does not already exist. Defaults to "imputer".
+        This method plots a confusion matrix using true and predicted labels. The plot is saved to disk as a ``<plot_format>`` file.
 
-            n_components (int, optional): Number of principal components to include in the PCA. Defaults to 3.
+        Args:
+            y_true_1d (np.ndarray | pd.DataFrame | list | torch.Tensor): 1D array of true integer labels in [0, n_classes-1].
+            y_pred_1d (np.ndarray | pd.DataFrame | list | torch.Tensor): 1D array of predicted integer labels in [0, n_classes-1].
+            label_names (Sequence[str] | None): Optional sequence of class names (length must equal n_classes). If provided, both the internal label order and displayed tick labels will respect this order (assumed to be 0..n-1).
+            prefix (str): Optional prefix for the output filename.
 
-            center (bool, optional): If True, centers the genotypes to the mean before doing the PCA. If False, no centering is done. Defaults to True.
+        Notes:
+            - If `label_names` is None, the display labels default to the numeric class indices inferred from `y_true_1d ∪ y_pred_1d`.
+        """
+        y_true_1d = misc.validate_input_type(y_true_1d, return_type="array")
+        y_pred_1d = misc.validate_input_type(y_pred_1d, return_type="array")
+
+        if not isinstance(y_true_1d, np.ndarray) or y_true_1d.ndim != 1:
+            msg = "y_true_1d must be a 1D array-like of true labels."
+            self.logger.error(msg)
+            raise TypeError(msg)
+
+        if not isinstance(y_pred_1d, np.ndarray) or y_pred_1d.ndim != 1:
+            msg = "y_pred_1d must be a 1D array-like of predicted labels."
+            self.logger.error(msg)
+            raise TypeError(msg)
+
+        if y_true_1d.ndim > 1:
+            y_true_1d = y_true_1d.flatten()
+
+        if y_pred_1d.ndim > 1:
+            y_pred_1d = y_pred_1d.flatten()
+
+        # Determine class count/order
+        if label_names is not None:
+            n_classes = len(label_names)
+            labels = np.arange(n_classes)  # our y_* are ints 0..n-1
+            display_labels = list(map(str, label_names))
+        else:
+            # Infer labels from data to keep matrix tight
+            labels = np.unique(np.concatenate([y_true_1d, y_pred_1d]))
+            display_labels = labels  # sklearn will convert to strings
 
-            scale (bool, optional): If True, scales the genotypes to unit variance before doing the PCA. If False, no scaling is done. Defaults to False.
+        fig, ax = plt.subplots(1, 1, figsize=(15, 15))
 
-            n_axes (int, optional): Number of principal component axes to plot. Must be set to either 2 or 3. If set to 3, a 3-dimensional plot will be made. Defaults to 2.
+        ConfusionMatrixDisplay.from_predictions(
+            y_true=y_true_1d,
+            y_pred=y_pred_1d,
+            labels=labels,
+            display_labels=display_labels,
+            ax=ax,
+            cmap="viridis",
+            colorbar=True,
+        )
 
-            point_size (int, optional): Point size for scatterplot points. Defaults to 15.
+        # Build a stable panel id before mutating prefix
+        panel_suffix = f"{prefix}_" if prefix else ""
+        panel_id = f"{self.model_name.lower()}_{panel_suffix}confusion_matrix"
 
-            plot_format (str, optional): Plot file format to use. Supported formats include: "pdf", "svg", "png", and "jpeg" (or "jpg"). An interactive HTML file is also created regardless of this setting. Defaults to "pdf".
+        if prefix != "":
+            prefix = f"{prefix}_"
 
-            bottom_margin (int, optional): Adjust bottom margin. If whitespace cuts off some of your plot, lower the corresponding margins. The default corresponds to that of plotly update_layout(). Defaults to 0.
+        out_name = (
+            f"{self.model_name.lower()}_{prefix}confusion_matrix.{self.plot_format}"
+        )
+        fig.savefig(self.output_dir / out_name, bbox_inches="tight")
+        if self.show_plots:
+            plt.show()
+        plt.close(fig)
+
+        # ---- MultiQC: confusion-matrix heatmap ----------------------------
+        if self._multiqc_enabled():
+            try:
+                self._queue_multiqc_confusion(
+                    y_true=y_true_1d,
+                    y_pred=y_pred_1d,
+                    labels=labels,
+                    display_labels=display_labels,
+                    panel_id=panel_id,
+                )
+            except Exception as exc:  # pragma: no cover
+                self.logger.warning(f"Failed to queue MultiQC confusion matrix: {exc}")
 
-            top (int, optional): Adjust top margin. If whitespace cuts off some of your plot, lower the corresponding margins. The default corresponds to that of plotly update_layout(). Defaults to 0.
+    def plot_gt_distribution(
+        self,
+        X: np.ndarray | pd.DataFrame | list | torch.Tensor,
+        is_imputed: bool = False,
+    ) -> None:
+        """Plot genotype distribution (IUPAC or integer-encoded).
 
-            left_margin (int, optional): Adjust left margin. If whitespace cuts off some of your plot, lower the corresponding margins. The default corresponds to that of plotly update_layout(). Defaults to 0.
+        This plots counts for all genotypes present in X. It supports IUPAC single-letter genotypes and integer encodings. Missing markers '-', '.', '?' are normalized to 'N'. Bars are annotated with counts and percentages.
 
-            right_margin (int, optional): Adjust right margin. If whitespace cuts off some of your plot, lower the corresponding margins. The default corresponds to that of plotly update_layout(). Defaults to 0.
+        Args:
+            X (np.ndarray | pd.DataFrame | list | torch.Tensor): Array-like genotype matrix. Rows=loci, cols=samples (any orientation is OK). Elements are IUPAC one-letter genotypes (e.g., 'A','C','G','T','N','R',...) or integers (e.g., 0/1/2[/3]).
+            is_imputed (bool): Whether these genotypes are imputed. Affects the title only. Defaults to False.
+        """
+        # Flatten X to a 1D Series
+        if isinstance(X, pd.DataFrame):
+            arr = X.values
+        elif torch.is_tensor(X):
+            arr = X.detach().cpu().numpy()
+        else:
+            arr = np.asarray(X)
+
+        s = pd.Series(arr.ravel())
+
+        # Detect string vs numeric encodings and normalize
+        if s.dtype.kind in ("O", "U", "S"):  # string-like → IUPAC path
+            s = s.astype(str).str.upper().replace({"-": "N", ".": "N", "?": "N"})
+            x_label = "Genotype (IUPAC)"
+
+            # Define canonical order: N, A/C/T/G, then IUPAC ambiguity codes.
+            canonical = ["A", "C", "T", "G"]
+            iupac_ambiguity = sorted(["M", "R", "W", "S", "Y", "K", "V", "H", "D", "B"])
+            base_order = ["N"] + canonical + iupac_ambiguity
+        else:  # numeric path (e.g., 0/1/2/[3], -1 for missing)
+            # Map common missing sentinels to 'N', keep others as strings for
+            # labeling
+            s = s.astype(float)  # allow NaN comparisons
+            s = s.where(~np.isin(s, [-1, np.nan]), other=np.nan)
+            s = s.fillna("N").astype(int, errors="ignore").astype(str)
+
+            x_label = "Genotype (Integer-encoded)"
+
+            # Support both ternary and quaternary encodings; keep a stable order
+            base_order = ["N", "0", "1", "2", "3"]
+
+        # Include any unexpected symbols at the end (sorted) so nothing is
+        # dropped
+        extras = sorted(set(s.unique()) - set(base_order))
+        full_order = base_order + [e for e in extras if e not in base_order]
+
+        # Count and reindex to show zero-count categories
+        counts = s.value_counts().reindex(full_order, fill_value=0)
+        df = counts.rename_axis("Genotype").reset_index(name="Count")
+        df["Percent"] = df["Count"] / df["Count"].sum() * 100
+
+        title = "Imputed Genotype Counts" if is_imputed else "Genotype Counts"
+
+        # --- Plot ---
+        fig, ax = plt.subplots(figsize=(8, 5))
+        sns.despine(fig=fig)
+
+        ax = sns.barplot(
+            data=df,
+            x="Genotype",
+            y="Percent",
+            hue="Genotype",
+            order=full_order,
+            errorbar=None,
+            ax=ax,
+            palette="Set1",
+            legend=False,
+            fill=True,
+        )
 
-            width (int, optional): Width of plot space. If your plot is cut off at the edges, even after adjusting the margins, increase the width and height. Try to keep the aspect ratio similar. Defaults to 1088.
+        ax.set_xlabel(x_label)
+        ax.set_ylabel("Percent")
+        ax.set_title(title)
+        ax.set_ylim((0.0, 50.0))
 
-            height (int, optional): Height of plot space. If your plot is cut off at the edges, even after adjusting the margins, increase the width and height. Try to keep the aspect ratio similar. Defaults to 700.
+        fig.tight_layout()
 
-        Returns:
-            numpy.ndarray: PCA data as a numpy array with shape (n_samples, n_components).
+        suffix = "imputed" if is_imputed else "original"
+        fn = self.output_dir / f"gt_distributions_{suffix}.{self.plot_format}"
+        fig.savefig(fn, dpi=300)
+
+        if self.show_plots:
+            plt.show()
+        plt.close(fig)
+
+        # ---- MultiQC: genotype-distribution barplot -----------------------
+        if self._multiqc_enabled():
+            try:
+                self._queue_multiqc_gt_distribution(df=df, is_imputed=is_imputed)
+            except Exception as exc:  # pragma: no cover
+                self.logger.warning(
+                    f"Failed to queue MultiQC genotype distribution: {exc}"
+                )
 
-            sklearn.decomposision.PCA: Scikit-learn PCA object from sklearn.decomposision.PCA. Any of the sklearn.decomposition.PCA attributes can be accessed from this object. See sklearn documentation.
+    # --------------------------------------------------------------------- #
+    # MultiQC helper methods                                                #
+    # --------------------------------------------------------------------- #
+    def _multiqc_enabled(self) -> bool:
+        """Return True if MultiQC integration is active."""
+        return bool(self.use_multiqc)
+
+    def _queue_multiqc_tuning(
+        self,
+        *,
+        study: optuna.study.Study,
+        model_name: str,
+        target_name: str,
+    ) -> None:
+        """Queue Optuna tuning results for MultiQC.
 
-        Examples:
-            >>>data = GenotypeData(
-            >>>    filename="snps.str",
-            >>>    filetype="structure2row",
-            >>>    popmapfile="popmap.txt",
-            >>>)
-            >>>
-            >>>ubp = ImputeUBP(genotype_data=data)
-            >>>
-            >>>components, pca = run_and_plot_pca(
-            >>>    data,
-            >>>    ubp,
-            >>>    scale=True,
-            >>>    center=True,
-            >>>    plot_format="png"
-            >>>)
-            >>>
-            >>>explvar = pca.explained_variance_ratio\_
+        Args:
+            study (optuna.study.Study): Optuna study object.
+            model_name (str): Name of the model.
+            target_name (str): Name of the target value.
         """
-        report_path = os.path.join(f"{prefix}_output", "plots")
-        Path(report_path).mkdir(parents=True, exist_ok=True)
+        if not self._multiqc_enabled():
+            return
 
-        if n_axes > 3:
-            raise ValueError(
-                ">3 axes is not supported; n_axes must be either 2 or 3."
-            )
-        if n_axes < 2:
-            raise ValueError(
-                "<2 axes is not supported; n_axes must be either 2 or 3."
+        # trial number vs objective value line graph
+        try:
+            df_trials = study.trials_dataframe(attrs=("number", "value"))
+        except Exception as exc:  # pragma: no cover
+            self.logger.warning(
+                f"Could not extract trials_dataframe for MultiQC: {exc}"
             )
+            return
 
-        imputer = imputer_object.imputed
+        if df_trials.empty or "value" not in df_trials:
+            return
 
-        df = misc.validate_input_type(
-            imputer.genotypes012_df, return_type="df"
-        )
-
-        original_df = misc.validate_input_type(
-            original_genotype_data.genotypes_012(fmt="pandas"),
-            return_type="df",
-        )
-
-        original_df.replace(-9, np.nan, inplace=True)
-
-        if center or scale:
-            # Center data to mean. Scaling to unit variance is off.
-            scaler = StandardScaler(with_mean=center, with_std=scale)
-            pca_df = scaler.fit_transform(df)
-        else:
-            pca_df = df.copy()
-
-        # Run PCA.
-        model = PCA(n_components=n_components)
-        components = model.fit_transform(pca_df)
+        data: Dict[str, Dict[int, int]] = {
+            model_name: {
+                row["number"]: row["value"]
+                for _, row in df_trials.iterrows()
+                if row["value"] is not None
+            }
+        }
 
-        df_pca = pd.DataFrame(
-            components[:, [0, 1, 2]], columns=["Axis1", "Axis2", "Axis3"]
+        if not data[model_name]:
+            return
+
+        SNPioMultiQC.queue_linegraph(
+            data=data,
+            panel_id=f"{self.model_name}_optuna_history",
+            section=self.multiqc_section,
+            title=f"{self.model_name} Optuna Optimization History",
+            index_label="Trial",
+            description=f"Optuna optimization history for {self.model_name} "
+            f"(target={target_name}).",
         )
 
-        df_pca["SampleID"] = original_genotype_data.samples
-        df_pca["Population"] = original_genotype_data.pops
-        df_pca["Size"] = point_size
-
-        _, ind, _, _, _ = imputer.calc_missing(original_df, use_pops=False)
-        df_pca["missPerc"] = ind
-
-        my_scale = [("rgb(19, 43, 67)"), ("rgb(86,177,247)")]  # ggplot default
+        # best-params table
+        try:
+            best_value = study.best_value
+            best_params = study.best_params
+        except Exception:
+            return
+
+        if best_params:
+            series = pd.Series(best_params, name="Best Value")
+            series["objective"] = best_value
+            SNPioMultiQC.queue_table(
+                df=series,
+                panel_id=f"{self.model_name}_optuna_best_params",
+                section=self.multiqc_section,
+                title=f"{self.model_name} Best Optuna Parameters",
+                index_label="Parameter",
+                description="Best Optuna hyperparameters and objective value.",
+            )
 
-        z = "Axis3" if n_axes == 3 else None
-        labs = {
-            "Axis1": f"PC1 ({round(model.explained_variance_ratio_[0] * 100, 2)}%)",
-            "Axis2": f"PC2 ({round(model.explained_variance_ratio_[1] * 100, 2)}%)",
-            "missPerc": "Missing Prop.",
-            "Population": "Population",
-        }
+    def _queue_multiqc_roc_curves(
+        self,
+        *,
+        fpr: dict,
+        tpr: dict,
+        label_names: Sequence[str],
+        panel_prefix: str,
+    ) -> None:
+        """Queue ROC and Precision-Recall curves for MultiQC.
 
-        if z is not None:
-            labs[
-                "Axis3"
-            ] = f"PC3 ({round(model.explained_variance_ratio_[2] * 100, 2)}%)"
-            fig = px.scatter_3d(
-                df_pca,
-                x="Axis1",
-                y="Axis2",
-                z="Axis3",
-                color="missPerc",
-                symbol="Population",
-                color_continuous_scale=my_scale,
-                custom_data=["Axis3", "SampleID", "Population", "missPerc"],
-                size="Size",
-                size_max=point_size,
-                labels=labs,
-            )
-        else:
-            fig = px.scatter(
-                df_pca,
-                x="Axis1",
-                y="Axis2",
-                color="missPerc",
-                symbol="Population",
-                color_continuous_scale=my_scale,
-                custom_data=["Axis3", "SampleID", "Population", "missPerc"],
-                size="Size",
-                size_max=point_size,
-                labels=labs,
-            )
-        fig.update_traces(
-            hovertemplate="<br>".join(
-                [
-                    "Axis 1: %{x}",
-                    "Axis 2: %{y}",
-                    "Axis 3: %{customdata[0]}",
-                    "Sample ID: %{customdata[1]}",
-                    "Population: %{customdata[2]}",
-                    "Missing Prop.: %{customdata[3]}",
-                ]
+        Args:
+            fpr (dict): False positive rates for each class.
+            tpr (dict): True positive rates for each class.
+            label_names (Sequence[str]): Class names.
+            panel_prefix (str): Optional prefix for panel IDs.
+        """
+        if not self._multiqc_enabled():
+            return
+
+        def _curve_to_mapping(
+            x_vals: Sequence[float], y_vals: Sequence[float]
+        ) -> Dict[float, float]:
+            """Return {x: y} mapping expected by MultiQC linegraphs."""
+            return {float(x): float(y) for x, y in zip(x_vals, y_vals)}
+
+        data: Dict[str, Dict[float, float]] = {}
+
+        # Only report the first three classes (MultiQC plot readability) plus micro/macro averages
+        class_keys = sorted(k for k in fpr.keys() if isinstance(k, int))
+        for idx in class_keys[:3]:
+            label = label_names[idx] if idx < len(label_names) else f"Class {idx}"
+            data[label] = _curve_to_mapping(fpr[idx], tpr[idx])
+
+        for agg in ("micro", "macro"):
+            if agg in fpr and agg in tpr:
+                pretty_name = f"{agg.title()} Average"
+                data[pretty_name] = _curve_to_mapping(fpr[agg], tpr[agg])
+
+        if not data:
+            return
+
+        # ROC curves
+        curve_data = cast(Dict[str, Dict[int, int]], data)
+
+        SNPioMultiQC.queue_linegraph(
+            data=curve_data,
+            panel_id=(
+                f"{self.model_name}_{panel_prefix}_roc_curves"
+                if panel_prefix
+                else f"{self.model_name}_roc_curves"
             ),
+            section=self.multiqc_section,
+            title=f"{self.model_name} ROC Curves",
+            index_label="False Positive Rate",
+            description="Multi-class ROC curves for PG-SUI predictions.",
         )
-        fig.update_layout(
-            showlegend=True,
-            margin=dict(
-                b=bottom_margin,
-                t=top_margin,
-                l=left_margin,
-                r=right_margin,
+
+    def _queue_multiqc_pr_curves(
+        self,
+        *,
+        precision: dict,
+        recall: dict,
+        label_names: Sequence[str],
+        panel_prefix: str,
+    ) -> None:
+        """Queue Precision-Recall curves for MultiQC."""
+        if not self._multiqc_enabled():
+            return
+
+        def _curve_to_mapping(
+            x_vals: Sequence[float], y_vals: Sequence[float]
+        ) -> Dict[float, float]:
+            """Return {recall: precision} mapping expected by MultiQC linegraphs."""
+            return {float(x): float(y) for x, y in zip(x_vals, y_vals)}
+
+        data: Dict[str, Dict[float, float]] = {}
+
+        # Only report the first three classes (MultiQC plot readability) plus micro/macro averages
+        class_keys = sorted(k for k in recall.keys() if isinstance(k, int))
+        for idx in class_keys[:3]:
+            if idx not in precision or idx not in recall:
+                continue
+            label = label_names[idx] if idx < len(label_names) else f"Class {idx}"
+            data[label] = _curve_to_mapping(recall[idx], precision[idx])
+
+        for agg in ("micro", "macro"):
+            if agg in precision and agg in recall:
+                pretty_name = f"{agg.title()} Average"
+                data[pretty_name] = _curve_to_mapping(recall[agg], precision[agg])
+
+        if not data:
+            return
+
+        curve_data = cast(Dict[str, Dict[int, int]], data)
+
+        SNPioMultiQC.queue_linegraph(
+            data=curve_data,
+            panel_id=(
+                f"{self.model_name}_{panel_prefix}_pr_curves"
+                if panel_prefix
+                else f"{self.model_name}_pr_curves"
             ),
-            width=width,
-            height=height,
-            legend_orientation="h",
-            legend_title="Population",
-            legend_title_font=dict(size=font_size),
-            legend_title_side="top",
-            font=dict(size=font_size),
-        )
-        fig.write_html(os.path.join(report_path, "imputed_pca.html"))
-        fig.write_image(
-            os.path.join(report_path, f"imputed_pca.{plot_format}"),
+            section=self.multiqc_section,
+            title=f"{self.model_name} Precision-Recall Curves",
+            index_label="Recall",
+            description="Multi-class Precision-Recall curves for PG-SUI predictions.",
         )
 
-        return components, model
-
-    @staticmethod
-    def plot_history(lod, nn_method, prefix="imputer"):
-        """Plot model history traces. Will be saved to file.
+    def _queue_multiqc_metrics(
+        self,
+        *,
+        metrics: Dict[str, float],
+        roc_auc: Dict[object, float],
+        average_precision: Dict[object, float],
+        label_names: Sequence[str],
+        panel_prefix: str,
+    ) -> None:
+        """Queue summary metrics and per-class AUC/AP for MultiQC.
 
         Args:
-            lod (List[tf.keras.callbacks.History]): List of history objects.
-            nn_method (str): Neural network method to plot. Possible options include: 'NLPCA', 'UBP', or 'VAE'. NLPCA and VAE get plotted the same, but UBP does it differently due to its three phases.
-            prefix (str, optional): Prefix to use for output directory. Defaults to 'imputer'.
-
-        Raises:
-            ValueError: nn_method must be either 'NLPCA', 'UBP', or 'VAE'.
+            metrics (Dict[str, float]): Summary metrics (accuracy, F1, etc.).
+            roc_auc (Dict[object, float]): Per-class and aggregate ROC-AUC values.
+            average_precision (Dict[object, float]): Per-class and aggregate average precision values.
+            label_names (Sequence[str]): Class names.
+            panel_prefix (str): Optional prefix for panel IDs.
         """
-        if nn_method == "NLPCA" or nn_method == "VAE" or nn_method == "SAE":
-            title = nn_method
-            fn = os.path.join(
-                f"{prefix}_output",
-                "plots",
-                "Unsupervised",
-                nn_method,
-                "histplot.pdf",
+        if not self._multiqc_enabled():
+            return
+
+        # Summary metrics table (accuracy, F1, etc.)
+        if metrics:
+            series = pd.Series(metrics, name="Value")
+            SNPioMultiQC.queue_table(
+                df=series,
+                panel_id=f"{self.model_name}_summary_metrics",
+                section=self.multiqc_section,
+                title=f"{self.model_name} Summary Metrics",
+                index_label="Metric",
+                description="Global evaluation metrics produced by PG-SUI.",
             )
 
-            if nn_method == "VAE":
-                fig, axes = plt.subplots(2, 2)
-                ax1 = axes[0, 0]
-                ax2 = axes[0, 1]
-                # ax3 = axes[1, 0]
-                # ax4 = axes[1, 1]
-            else:
-                fig, (ax1, ax2) = plt.subplots(1, 2)
-            fig.suptitle(title)
-            fig.tight_layout(h_pad=3.0, w_pad=3.0)
-            history = lod[0]
-
-            acctrain = (
-                "categorical_accuracy" if nn_method == "NLPCA" else "accuracy"
+        # Per-class ROC-AUC and AP heatmap
+        rows: List[Dict[str, float | str]] = []
+
+        # integer keys are classes; others are 'micro', 'macro'
+        class_keys = [k for k in roc_auc.keys() if isinstance(k, int)]
+        class_keys_sorted = sorted(class_keys)
+
+        for i in class_keys_sorted:
+            class_name = label_names[i] if i < len(label_names) else f"Class {i}"
+            rows.append(
+                {
+                    "Class": str(class_name),
+                    "ROC_AUC": float(roc_auc.get(i, np.nan)),
+                    "AveragePrecision": float(average_precision.get(i, np.nan)),
+                }
             )
 
-            # if nn_method == "VAE":
-            #     accval = "val_accuracy"
-            #     # recon_loss = "reconstruction_loss"
-            #     # kl_loss = "kl_loss"
-            #     # val_recon_loss = "val_reconstruction_loss"
-            #     # val_kl_loss = "val_kl_loss"
-            #     lossval = "val_loss"
+        for agg in ("micro", "macro"):
+            if agg in roc_auc:
+                rows.append(
+                    {
+                        "Class": agg,
+                        "ROC_AUC": float(roc_auc.get(agg, np.nan)),
+                        "AveragePrecision": float(average_precision.get(agg, np.nan)),
+                    }
+                )
 
-            if nn_method == "SAE":
-                accval = "val_accuracy"
-                lossval = "val_loss"
+        if not rows:
+            return
+
+        df = pd.DataFrame(rows).set_index("Class")
+        suffix = f"{panel_prefix}_" if panel_prefix else ""
+        panel_id = f"{self.model_name}_{suffix}roc_pr_summary"
+
+        SNPioMultiQC.queue_heatmap(
+            df=df,
+            panel_id=panel_id,
+            section=self.multiqc_section,
+            title=f"{self.model_name} ROC-AUC and Average Precision",
+            index_label="Class",
+            description=(
+                "Per-class ROC-AUC and average precision for PG-SUI predictions (including micro/macro averages where available)."
+            ),
+        )
 
-            # Plot train accuracy
-            ax1.plot(history[acctrain])
-            ax1.set_title("Model Accuracy")
-            ax1.set_ylabel("Accuracy")
-            ax1.set_xlabel("Epoch")
-            ax1.set_ylim(bottom=0.0, top=1.0)
-            ax1.set_yticks([0.0, 0.25, 0.5, 0.75, 1.0])
-
-            labels = ["Train"]
-            if nn_method == "SAE":
-                # Plot validation accuracy
-                ax1.plot(history[accval])
-                labels.append("Validation")
-
-            ax1.legend(labels, loc="best")
-
-            # Plot model loss
-            # if nn_method == "VAE":
-            #     # Reconstruction loss only.
-            #     ax2.plot(history["loss"])
-            # ax2.plot(history[val_recon_loss])
-
-            # # KL Loss
-            # ax3.plot(history[kl_loss])
-            # ax3.plot(history[val_kl_loss])
-            # ax3.set_title("KL Divergence Loss")
-            # ax3.set_ylabel("Loss")
-            # ax3.set_xlabel("Epoch")
-            # ax3.legend(labels, loc="best")
-
-            # Total Loss (Reconstruction Loss + KL Loss)
-            # ax4.plot(history["loss"])
-            # ax4.plot(history[lossval])
-            # ax4.set_title("Total Loss (Recon. + KL)")
-            # ax4.set_ylabel("Loss")
-            # ax4.set_xlabel("Epoch")
-            # ax4.legend(labels, loc="best")
-
-            # else:
-            ax2.plot(history["loss"])
-
-            if nn_method == "SAE":
-                ax2.plot(history[lossval])
-
-            ax2.set_title("Total Loss")
-            ax2.set_ylabel("Loss")
-            ax2.set_xlabel("Epoch")
-            ax2.legend(labels, loc="best")
-
-            fig.savefig(fn, bbox_inches="tight", facecolor="white")
+    def _queue_multiqc_history(
+        self,
+        *,
+        history: Dict[str, List[float] | Dict[str, List[float]] | None] | None,
+    ) -> None:
+        """Queue training history (loss vs epoch) for MultiQC.
 
-            plt.close()
-            plt.clf()
-
-        elif nn_method == "UBP":
-            fig = plt.figure(figsize=(12, 16))
-            fig.suptitle(nn_method)
-            fig.tight_layout(h_pad=2.0, w_pad=2.0)
-            fn = os.path.join(
-                f"{prefix}_output",
-                "plots",
-                "Unsupervised",
-                nn_method,
-                "histplot.pdf",
-            )
+        Args:
+            history (Dict[str, List[float]] | None): Dictionary with lists of history objects. Keys should be "Train" and "Validation".
+        """
+        if not self._multiqc_enabled() or history is None:
+            return
 
-            idx = 1
-            for i, history in enumerate(lod, start=1):
-                plt.subplot(3, 2, idx)
-                title = f"Phase {i}"
-
-                # Plot model accuracy
-                ax = plt.gca()
-                ax.plot(history["categorical_accuracy"])
-                ax.set_title(f"{title} Accuracy")
-                ax.set_ylabel("Accuracy")
-                ax.set_xlabel("Epoch")
-                ax.set_yticks([0.0, 0.25, 0.5, 0.75, 1.0])
-                ax.legend(["Training"], loc="best")
-
-                # Plot model loss
-                plt.subplot(3, 2, idx + 1)
-                ax = plt.gca()
-                ax.plot(history["loss"])
-                ax.set_title(f"{title} Loss")
-                ax.set_ylabel("Loss (MSE)")
-                ax.set_xlabel("Epoch")
-                ax.legend(["Train"], loc="best")
-
-                idx += 2
-
-            plt.savefig(fn, bbox_inches="tight", facecolor="white")
+        data: Dict[str, Dict[int, int]] = {}
 
-            plt.close()
-            plt.clf()
+        if self.model_name != "ImputeUBP":
+            if not isinstance(history, dict) or "Train" not in history:
+                return
 
-        else:
-            raise ValueError(
-                f"nn_method must be either 'NLPCA', 'UBP', or 'VAE', but got {nn_method}"
-            )
+            train_vals = pd.Series(history["Train"]).iloc[1:]
 
-    @staticmethod
-    def plot_certainty_heatmap(
-        y_certainty, sample_ids=None, nn_method="VAE", prefix="imputer"
-    ):
-        fig = plt.figure()
-        hm = sns.heatmap(
-            data=y_certainty,
-            cmap="viridis",
-            vmin=0.0,
-            vmax=1.0,
-            cbar_kws={"label": "Prob."},
-        )
-        hm.set_xlabel("Site")
-        hm.set_ylabel("Sample")
-        hm.set_title("Probabilities of Uncertain Sites")
-        fig.tight_layout()
-        fig.savefig(
-            os.path.join(
-                f"{prefix}_output",
-                "plots",
-                "Unsupervised",
-                nn_method,
-                "uncertainty_plot.png",
-            ),
-            bbox_inches="tight",
-            facecolor="white",
+            data["Train"] = {
+                epoch: val for epoch, val in enumerate(train_vals.values, start=1)
+            }
+        else:
+            if not (
+                isinstance(history, dict)
+                and "Train" in history
+                and isinstance(history["Train"], dict)
+            ):
+                return
+            for phase in range(1, 4):
+                key = f"Phase {phase}"
+                if key not in history["Train"]:
+                    continue
+                series = pd.Series(history["Train"][key]).iloc[1:]
+                data[key] = {
+                    epoch: val for epoch, val in enumerate(series.values, start=1)
+                }
+
+        if not data:
+            return
+
+        SNPioMultiQC.queue_linegraph(
+            data=data,
+            panel_id=f"{self.model_name}_training_history",
+            section=self.multiqc_section,
+            title=f"{self.model_name} Training Loss per Epoch",
+            index_label="Epoch",
+            description="Training loss trajectory by epoch as recorded by PG-SUI.",
         )
 
-    @staticmethod
-    def plot_confusion_matrix(
-        y_true_1d, y_pred_1d, nn_method, prefix="imputer"
-    ):
-        fig, ax = plt.subplots(1, 1, figsize=(15, 15))
-        ConfusionMatrixDisplay.from_predictions(
-            y_true=y_true_1d, y_pred=y_pred_1d, ax=ax
-        )
+    def _queue_multiqc_confusion(
+        self,
+        *,
+        y_true: np.ndarray,
+        y_pred: np.ndarray,
+        labels: np.ndarray,
+        display_labels: List[str] | np.ndarray,
+        panel_id: str,
+    ) -> None:
+        """Queue confusion-matrix heatmap for MultiQC.
 
-        outfile = os.path.join(
-            f"{prefix}_output",
-            "plots",
-            "Unsupervised",
-            nn_method,
-            f"confusion_matrix_{nn_method}.png",
+        Args:
+            y_true (np.ndarray): 1D array of true integer labels.
+            y_pred (np.ndarray): 1D array of predicted integer labels.
+            labels (np.ndarray): Array of label indices to index the confusion matrix.
+            display_labels (List[str] | np.ndarray): Labels to display on axes.
+            panel_id (str): Panel ID for MultiQC.
+        """
+        if not self._multiqc_enabled():
+            return
+
+        cm = confusion_matrix(y_true, y_pred, labels=labels)
+        df_cm = pd.DataFrame(cm, index=display_labels, columns=display_labels)
+
+        SNPioMultiQC.queue_heatmap(
+            df=df_cm,
+            panel_id=panel_id,
+            section=self.multiqc_section,
+            title=f"{self.model_name} Confusion Matrix",
+            index_label="True Label",
+            description=(
+                "Confusion matrix for PG-SUI predictions. Rows correspond to true "
+                "labels; columns correspond to predicted labels."
+            ),
         )
 
-        if os.path.isfile(outfile):
-            os.remove(outfile)
-
-        fig.savefig(outfile, facecolor="white")
-
-    @staticmethod
-    def plot_gt_distribution(df, plot_path):
-        df = misc.validate_input_type(df, return_type="df")
-        df_melt = pd.melt(df, value_name="Count")
-        cnts = df_melt["Count"].value_counts()
-        cnts.index.names = ["Genotype"]
-        cnts = pd.DataFrame(cnts).reset_index()
-        cnts.sort_values(by="Genotype", inplace=True)
-        cnts["Genotype"] = cnts["Genotype"].astype(str)
-
-        fig, ax = plt.subplots(1, 1, figsize=(15, 15))
-        g = sns.barplot(x="Genotype", y="Count", data=cnts, ax=ax)
-        g.set_xlabel("Integer-encoded Genotype")
-        g.set_ylabel("Count")
-        g.set_title("Genotype Counts")
-        for p in g.patches:
-            g.annotate(
-                f"{p.get_height():.1f}",
-                (p.get_x() + 0.25, p.get_height() + 0.01),
-                xytext=(0, 1),
-                textcoords="offset points",
-                va="bottom",
-            )
+    def _queue_multiqc_gt_distribution(
+        self,
+        *,
+        df: pd.DataFrame,
+        is_imputed: bool,
+    ) -> None:
+        """Queue genotype-distribution barplot for MultiQC.
 
-        fig.savefig(
-            os.path.join(plot_path, "genotype_distributions.png"),
-            bbox_inches="tight",
-            facecolor="white",
+        Args:
+            df (pd.DataFrame): DataFrame with 'Genotype' and 'Percent' columns
+            is_imputed (bool): Whether these genotypes are imputed.
+        """
+        if not self._multiqc_enabled():
+            return
+
+        if "Genotype" not in df.columns or "Percent" not in df.columns:
+            return
+
+        series = df.set_index("Genotype")["Percent"]
+        suffix = "imputed" if is_imputed else "original"
+        title = (
+            f"{self.model_name} Imputed Genotype Distribution"
+            if is_imputed
+            else f"{self.model_name} Genotype Distribution"
         )
-        plt.close()
-
-    @staticmethod
-    def plot_label_clusters(z_mean, labels, prefix="imputer"):
-        """Display a 2D plot of the classes in the latent space."""
-        fig, ax = plt.subplots(1, 1, figsize=(15, 15))
-
-        sns.scatterplot(x=z_mean[:, 0], y=z_mean[:, 1], ax=ax)
-        ax.set_xlabel("Latent Dimension 1")
-        ax.set_ylabel("Latent Dimension 2")
 
-        outfile = os.path.join(
-            f"{prefix}_output",
-            "plots",
-            "Unsupervised",
-            "VAE",
-            "label_clusters.png",
+        SNPioMultiQC.queue_barplot(
+            df=series,
+            panel_id=f"{self.model_name}_gt_distribution_{suffix}",
+            section=self.multiqc_section,
+            title=title,
+            index_label="Genotype",
+            value_label="Percent",
+            description=(
+                "Genotype frequency distribution (percent of calls per genotype) computed by PG-SUI."
+            ),
         )
-
-        if os.path.isfile(outfile):
-            os.remove(outfile)
-
-        fig.savefig(outfile, facecolor="white", bbox_inches="tight")