polynet 1.0.1__py3-none-any.whl

polynet/app/components/forms/analyse_results.py

@@ -0,0 +1,825 @@
+import re
+
+import pandas as pd
+import streamlit as st
+
+from polynet.app.options.state_keys import AnalyseResultsStateKeys, PlotCustomiserStateKeys
+from polynet.config.column_names import (
+    get_iterator_name,
+    get_predicted_label_column_name,
+    get_true_label_column_name,
+)
+from polynet.config.constants import DataSet, ResultColumn
+from polynet.config.display_names import prettify_label, prettify_metric
+from polynet.config.enums import Plot, ProblemType, SplitType
+from polynet.config.schemas import DataConfig
+from polynet.utils.statistical_analysis import (
+    METRIC_COMPARISON_TESTS,
+    MIN_SPLITS_FOR_WILCOXON,
+    MULTIPLE_COMPARISON_METHODS,
+    correct_pvalue_matrix,
+    mcnemar_pvalue_matrix,
+    metrics_pvalue_matrix,
+    min_wilcoxon_pvalue,
+    regression_pvalue_matrix,
+)
+from polynet.visualization import (
+    plot_bootstrap_boxplots,
+    plot_confusion_matrix,
+    plot_parity,
+    plot_pvalue_matrix,
+)
+
+
+def _prettify_series(series: pd.Series, column_name: str, abbreviate: bool) -> pd.Series:
+    """Map model identifiers in a hue/style series to display labels.
+
+    Only the model column is translated (its values are raw model ids like
+    ``"random_forest-polybert"``); other grouping columns (set, iteration) are
+    returned unchanged so their legends keep their original values.
+    """
+    if column_name == ResultColumn.MODEL:
+        return series.map(lambda v: prettify_label(v, abbreviate=abbreviate))
+    return series
+
+
+def _multiple_comparison_selector(key: str) -> str | None:
+    """Render a correction-method selectbox and return the ``multipletests`` name.
+
+    Pairwise comparisons across many models inflate the family-wise error rate,
+    so a correction is applied by default (Holm-Bonferroni). The user can switch
+    to Bonferroni, Benjamini-Hochberg (FDR), or disable correction.
+    """
+    labels = list(MULTIPLE_COMPARISON_METHODS.keys())
+    choice = st.selectbox(
+        "Multiple-comparison correction",
+        options=labels,
+        index=0,  # Holm-Bonferroni
+        key=key,
+        help=(
+            "Adjusts the pairwise p-values for the number of simultaneous model "
+            "comparisons. Holm-Bonferroni and Bonferroni control the family-wise "
+            "error rate; Benjamini-Hochberg controls the false discovery rate."
+        ),
+    )
+    return MULTIPLE_COMPARISON_METHODS[choice]
+
+
+def compare_predictions_form(
+    predictions_df: pd.DataFrame, target_variable_name: str, data_options: DataConfig
+):
+
+    st.write("### Model predictions pairwise comparison")
+
+    analyse_set = st.selectbox(
+        "Select a set to analyse",
+        options=[DataSet.Training, DataSet.Validation, DataSet.Test, "All sets"],
+        index=2,
+        key="analyse_set",
+    )
+
+    if analyse_set != "All sets":
+        predictions_df = predictions_df.loc[predictions_df[ResultColumn.SET] == analyse_set]
+
+    label_col_name = get_true_label_column_name(target_variable_name=target_variable_name)
+    true_vals = predictions_df[label_col_name].to_numpy()
+
+    trained_models = [
+        col.split(" ")[0] for col in predictions_df.columns if ResultColumn.PREDICTED in col
+    ]
+    compare_models = st.multiselect(
+        "Select models to test statistically difference in predictions",
+        options=trained_models,
+        format_func=prettify_label,
+    )
+
+    if len(compare_models) > 1:
+
+        predicted_cols = [
+            get_predicted_label_column_name(
+                target_variable_name=target_variable_name, model_name=model
+            )
+            for model in compare_models
+        ]
+        predictions_array = predictions_df[predicted_cols].to_numpy().T
+
+        if data_options.problem_type == ProblemType.Classification:
+            p_matrix = mcnemar_pvalue_matrix(y_true=true_vals, predictions=predictions_array)
+        elif data_options.problem_type == ProblemType.Regression:
+            p_matrix = regression_pvalue_matrix(y_true=true_vals, predictions=predictions_array)
+
+        correction = _multiple_comparison_selector(key="predictions_pvalue_correction")
+        p_matrix = correct_pvalue_matrix(p_matrix, method=correction)
+
+        config = get_plot_customisation_form(
+            plot_type=Plot.MatrixPlot,
+            data=None,
+            data_options=data_options,
+            color_by_opts=None,
+            model_pred_col=None,
+            model_true_cols=None,
+            key_prefix="predictions_",
+        )
+        abbreviate = config.pop("abbreviate_labels", False)
+        model_labels = [prettify_label(m, abbreviate=abbreviate) for m in compare_models]
+        plot = plot_pvalue_matrix(p_matrix=p_matrix, model_names=model_labels, **config)
+
+        return plot
+
+
+def _metric_test_selector(
+    n_splits: int, analyse_set: str, split_sizes: dict[str, float] | None
+) -> tuple[str, float | None] | None:
+    """
+    Render the statistical-test choice for the metric-level comparison.
+
+    Warns when Wilcoxon cannot reach significance with the available number of
+    splits, and computes ``n_test / n_train`` for the Nadeau–Bengio corrected
+    t-test from the saved split sizes.
+
+    Returns
+    -------
+    tuple[str, float | None] | None
+        ``(test, test_train_ratio)`` for ``metrics_pvalue_matrix``, or ``None``
+        when the chosen test cannot be run.
+    """
+    label = st.selectbox(
+        "Statistical test",
+        options=list(METRIC_COMPARISON_TESTS),
+        index=0,
+        key="metrics_statistical_test",
+        help=(
+            "Metrics come from repeated random splits of the same dataset, whose training "
+            "sets overlap, so the per-split values are not independent. The Nadeau–Bengio "
+            "corrected resampled t-test accounts for this overlap by inflating the variance "
+            "with n_test / n_train; the Wilcoxon signed-rank test does not."
+        ),
+    )
+    test = METRIC_COMPARISON_TESTS[label]
+
+    if test == "wilcoxon":
+        if n_splits < MIN_SPLITS_FOR_WILCOXON:
+            st.warning(
+                f"Only {n_splits} split(s) are available. With {n_splits} paired values the "
+                f"smallest two-sided p-value the Wilcoxon test can return is "
+                f"{min_wilcoxon_pvalue(n_splits):.4f}, so no difference can be significant at "
+                f"0.05. Train with at least {MIN_SPLITS_FOR_WILCOXON} repeated splits, or use "
+                "the Nadeau–Bengio corrected t-test."
+            )
+        return test, None
+
+    if n_splits < 2:
+        st.warning("The corrected t-test needs at least 2 splits.")
+        return None
+    if not split_sizes or DataSet.Training not in split_sizes or analyse_set not in split_sizes:
+        st.error(
+            "Split sizes are unavailable (split_indices.json not found), so the corrected "
+            "t-test cannot be computed. Re-run training to save the split indices."
+        )
+        return None
+    ratio = split_sizes[analyse_set] / split_sizes[DataSet.Training]
+    st.caption(
+        f"Corrected t-test with n_test / n_train = {split_sizes[analyse_set]:.0f} / "
+        f"{split_sizes[DataSet.Training]:.0f} = {ratio:.3f} (mean set sizes across splits; "
+        "n_train is the training set without validation)."
+    )
+    if analyse_set == DataSet.Training:
+        st.info(
+            "The corrected t-test is meant for held-out data. On training-set metrics "
+            "n_test / n_train = 1, which makes it very conservative."
+        )
+    return test, ratio
+
+
+def compare_metrics_form(
+    metrics: dict,
+    data_options: DataConfig | None = None,
+    split_sizes: dict[str, float] | None = None,
+):
+
+    st.write("### Model metrics comparison")
+
+    analyse_set = st.selectbox(
+        "Select a set to analyse",
+        options=[DataSet.Training, DataSet.Validation, DataSet.Test],
+        index=2,
+        key="analyse_set_metrics",
+    )
+
+    # Build: data_dict[model][metric_name] -> list of bootstrap values
+    data_dict = {}
+    metrics_name = set()
+    for _, dictio in metrics.items():
+        for model, per_model_dict in dictio.items():
+            data_dict.setdefault(model, {})
+            for metric_name, val in per_model_dict[analyse_set].items():
+                data_dict[model].setdefault(metric_name, [])
+                data_dict[model][metric_name].append(val)
+                metrics_name.add(metric_name)
+
+    models = st.multiselect(
+        "Select models to compare",
+        options=list(data_dict.keys()),
+        key="comparemodelmetrics",
+        format_func=prettify_label,
+    )
+
+    metric_choice = st.selectbox(
+        "Select a metric to compare",
+        options=sorted(metrics_name),
+        key="metriccomparemetric",
+        format_func=prettify_metric,
+    )
+
+    selected = {m: data_dict[m][metric_choice] for m in models if m in data_dict}
+
+    plot_type = st.selectbox("Select a type of plot", options=["P-value Matrix", "Box Plot"])
+
+    if len(selected) > 1:
+
+        if plot_type == "P-value Matrix":
+            config = get_plot_customisation_form(
+                plot_type=Plot.MatrixPlot,
+                data=None,
+                data_options=data_options,
+                color_by_opts=None,
+                model_pred_col=None,
+                model_true_cols=None,
+                key_prefix="metrics_",
+            )
+            n_splits = min(len(v) for v in selected.values())
+            test_choice = _metric_test_selector(
+                n_splits=n_splits, analyse_set=analyse_set, split_sizes=split_sizes
+            )
+            if test_choice is None:
+                return None
+            test, test_train_ratio = test_choice
+            p_matrix, order = metrics_pvalue_matrix(
+                selected, test=test, test_train_ratio=test_train_ratio
+            )
+            correction = _multiple_comparison_selector(key="metrics_pvalue_correction")
+            p_matrix = correct_pvalue_matrix(p_matrix, method=correction)
+            abbreviate = config.pop("abbreviate_labels", False)
+            order_labels = [prettify_label(m, abbreviate=abbreviate) for m in order]
+            plot = plot_pvalue_matrix(p_matrix=p_matrix, model_names=order_labels, **config)
+        elif plot_type == "Box Plot":
+            config = get_plot_customisation_form(
+                plot_type=Plot.MetricsBoxPlot,
+                data=None,
+                data_options=data_options,
+                color_by_opts=None,
+                model_pred_col=None,
+                model_true_cols=None,
+                key_prefix="metrics_",
+            )
+            abbreviate = config.pop("abbreviate_labels", False)
+            selected = {
+                prettify_label(m, abbreviate=abbreviate): vals for m, vals in selected.items()
+            }
+            plot = plot_bootstrap_boxplots(
+                metrics_dict=selected, metric_name=prettify_metric(metric_choice), **config
+            )
+        return plot
+
+
+def get_plot_customisation_form(
+    plot_type: str,
+    data: pd.DataFrame,
+    data_options: DataConfig,
+    color_by_opts=None,
+    model_pred_col: str = None,
+    model_true_cols: str = None,
+    key_prefix: str = "",
+):
+    """
+    Renders a dynamic Streamlit form for customizing a plot, based on plot_type.
+    Returns a dictionary of kwargs for the plotting function.
+    """
+    config = {}
+    widget_suffix = key_prefix + plot_type
+
+    with st.expander("Customise Plot", expanded=False):
+
+        # Display-only label style. Affects model/descriptor labels shown on the
+        # plot (legend, axis ticks); the underlying data keeps its raw names.
+        abbreviate_labels = st.checkbox(
+            "Use abbreviated labels",
+            value=False,
+            key=PlotCustomiserStateKeys.AbbreviateLabels + widget_suffix,
+            help="Show short model/descriptor names (e.g. 'RF polyBERT') instead of "
+            "full names (e.g. 'Random Forest polyBERT').",
+        )
+        config["abbreviate_labels"] = abbreviate_labels
+
+        config["title"] = st.text_input(
+            "Plot Title",
+            value=f"{data_options.target_variable_name} {plot_type.title()} Plot",
+            key=PlotCustomiserStateKeys.PlotTitle + widget_suffix,
+            help="Enter a title for the plot.",
+        )
+
+        config["title_fontsize"] = st.slider(
+            "Title Font Size",
+            min_value=8,
+            max_value=30,
+            value=16,
+            step=1,
+            key=PlotCustomiserStateKeys.PlotFontSize + widget_suffix,
+            help="Select the font size for the plot title.",
+        )
+
+        columns = st.columns(2)
+        with columns[0]:
+            config["height"] = st.slider(
+                "Plot height",
+                min_value=4,
+                max_value=14,
+                value=7,
+                step=1,
+                key="plot_height" + widget_suffix,
+            )
+            if plot_type not in [Plot.MatrixPlot, Plot.MetricsBoxPlot]:
+                config["x_label"] = st.text_input(
+                    "X-axis Label",
+                    value=model_true_cols,
+                    key=PlotCustomiserStateKeys.PlotXLabel + widget_suffix,
+                    help="Enter a label for the X-axis.",
+                )
+            config["x_label_fontsize"] = st.slider(
+                "X-axis Label Font Size",
+                min_value=8,
+                max_value=30,
+                value=14,
+                step=1,
+                key=PlotCustomiserStateKeys.PlotXLabelFontSize + widget_suffix,
+                help="Select the font size for the X-axis label.",
+            )
+        with columns[1]:
+
+            config["width"] = st.slider(
+                "Plot width",
+                min_value=4,
+                max_value=14,
+                value=7,
+                step=1,
+                key="plot_width" + widget_suffix,
+            )
+
+            if plot_type not in [Plot.MatrixPlot, Plot.MetricsBoxPlot]:
+                config["y_label"] = st.text_input(
+                    "Y-axis Label",
+                    value=model_pred_col,
+                    key=PlotCustomiserStateKeys.PlotYLabel + widget_suffix,
+                    help="Enter a label for the Y-axis.",
+                )
+            config["y_label_fontsize"] = st.slider(
+                "Y-axis Label Font Size",
+                min_value=8,
+                max_value=30,
+                value=14,
+                step=1,
+                key=PlotCustomiserStateKeys.PlotYLabelFontSize + widget_suffix,
+                help="Select the font size for the Y-axis label.",
+            )
+
+        config["tick_size"] = st.select_slider(
+            "Tick Size",
+            options=list(range(8, 21)),
+            value=12,
+            key=PlotCustomiserStateKeys.TickSize + widget_suffix,
+            help="Select the font size for the tick labels.",
+        )
+
+        if plot_type == Plot.Parity:
+
+            config["grid"] = st.checkbox(
+                "Show Grid",
+                value=True,
+                key=PlotCustomiserStateKeys.PlotGrid + widget_suffix,
+                help="Check to display grid lines on the plot.",
+            )
+
+            config["point_size"] = st.slider(
+                "Point Size",
+                min_value=1,
+                max_value=100,
+                value=50,
+                step=1,
+                key=PlotCustomiserStateKeys.PlotPointSize + widget_suffix,
+                help="Select the size of the points in the plot.",
+            )
+
+            config["border_color"] = st.color_picker(
+                "Point Border Color",
+                value="#000000",
+                key=PlotCustomiserStateKeys.PlotPointBorderColour + widget_suffix,
+                help="Select a color for the border of the plot.",
+            )
+
+            config["legend"] = st.checkbox(
+                "Show Legend",
+                value=True,
+                key=PlotCustomiserStateKeys.ShowLegend + widget_suffix,
+                help="Check to display the legend on the plot.",
+            )
+
+            if config["legend"]:
+                config["legend_fontsize"] = st.slider(
+                    "Legend Font Size",
+                    min_value=8,
+                    max_value=20,
+                    value=12,
+                    step=1,
+                    key=PlotCustomiserStateKeys.LegendFontSize + widget_suffix,
+                    help="Select the font size for the legend.",
+                )
+            else:
+                config["legend_fontsize"] = None
+
+            if len(color_by_opts) > 1:
+
+                hue = st.selectbox(
+                    "Color by",
+                    options=color_by_opts,
+                    key=PlotCustomiserStateKeys.ColourBy + widget_suffix,
+                    index=1,
+                    help="Select a column to color the points in the plot.",
+                )
+                config["hue"] = _prettify_series(data[hue], hue, abbreviate_labels) if hue else None
+
+                style_by = st.selectbox(
+                    "Style by",
+                    options=color_by_opts,
+                    key=PlotCustomiserStateKeys.StyleBy + widget_suffix,
+                    index=0,
+                    help="Select a column to style the points in the plot.",
+                )
+                config["style_by"] = (
+                    _prettify_series(data[style_by], style_by, abbreviate_labels)
+                    if style_by
+                    else None
+                )
+            else:
+                config["hue"] = None
+                config["style_by"] = None
+
+            if config["hue"] is not None:
+                cmap = st.selectbox(
+                    "Color Map",
+                    options=["tab10", "pastel", "colorblind", "Spectral", "Custom"],
+                    key=PlotCustomiserStateKeys.CMap + widget_suffix,
+                    help="Select a color map for the points in the plot.",
+                )
+                if cmap == "Custom":
+                    cmap = []
+                    for cat in config["hue"].unique():
+                        color = st.color_picker(
+                            f"Color for {cat}",
+                            value="#1f77b4",
+                            key=f"color_{cat}" + widget_suffix,
+                            help=f"Select a color for the category '{cat}' in the plot.",
+                        )
+                        cmap.append(color)
+                config["palette"] = cmap
+                config["point_color"] = None
+            else:
+                config["palette"] = None
+                config["point_color"] = st.color_picker(
+                    "Point Fill Color",
+                    value="#1f77b4",
+                    key="point_color" + widget_suffix,
+                    help="Select a color for the points in the plot.",
+                )
+
+        elif plot_type == Plot.ConfusionMatrix:
+            config["cmap"] = st.selectbox(
+                "Color Map",
+                options=["Blues", "Greens", "Oranges"],
+                key=PlotCustomiserStateKeys.CMap + widget_suffix,
+            )
+
+            config["label_fontsize"] = st.slider(
+                "Label Font Size",
+                min_value=8,
+                max_value=30,
+                value=16,
+                step=1,
+                key=PlotCustomiserStateKeys.PlotLabelFontSize + widget_suffix,
+                help="Select the font size for the labels in the confusion matrix.",
+            )
+
+            labels = data_options.class_names.keys()
+            config["display_labels"] = []
+            for label in labels:
+                name = st.text_input(
+                    f"Name for label for {label}",
+                    value=data_options.class_names[label],
+                    key=f"{PlotCustomiserStateKeys.LabelNames}_{label}" + widget_suffix,
+                    help=f"Enter the label for the class '{label}'.",
+                )
+                config["display_labels"].append(name)
+
+        elif plot_type == Plot.MatrixPlot:
+
+            config["mask_upper_triangle"] = st.checkbox(
+                "Mask upper triangle", key="mask_upper_triangle" + widget_suffix
+            )
+
+            col1, col2 = st.columns(2)
+
+            with col1:
+                config["non_signifficant_colour"] = st.color_picker(
+                    "Non-signifficant colour",
+                    value="#4575B4",
+                    key="non_signifficant" + widget_suffix,
+                )
+            with col2:
+                config["signifficant_colour"] = st.color_picker(
+                    "Non-signifficant colour", value="#D73027", key="signifficant" + widget_suffix
+                )
+
+        elif plot_type == Plot.MetricsBoxPlot:
+
+            cols = st.columns(3)
+
+            with cols[0]:
+                fill_colour = st.color_picker(
+                    "Box Fill Colour",
+                    value="#9aa0a5",
+                    key="box_fill" + widget_suffix,
+                    help="Select a fill color for the box in the box plot.",
+                )
+            with cols[1]:
+                border_colour = st.color_picker(
+                    "Box Border Colour",
+                    value="#000000",
+                    key="box_border" + widget_suffix,
+                    help="Select a border color for the box in the box plot.",
+                )
+            with cols[2]:
+                median_colour = st.color_picker(
+                    "Median Line Colour",
+                    value="#FF0000",
+                    key="median_color" + widget_suffix,
+                    help="Select a color for the median line in the box plot.",
+                )
+
+            config["fill_colour"] = fill_colour
+            config["border_colour"] = border_colour
+            config["median_colour"] = median_colour
+
+        config["dpi"] = st.slider(
+            "DPI",
+            min_value=100,
+            max_value=600,
+            value=300,
+            step=10,
+            key="dpi" + widget_suffix,
+            help="Select the DPI for the plot.",
+        )
+
+    return config
+
+
+def confusion_matrix_plot_form(
+    predictions_df: pd.DataFrame, split_type: SplitType | str, data_options: DataConfig
+):
+    """
+    Renders a Streamlit form for customizing and displaying a parity plot based on model predictions.
+    Returns a matplotlib figure object.
+    """
+
+    iterator = get_iterator_name(split_type)
+
+    iteration = st.multiselect(
+        f"Select the {iterator} to display parity plot",
+        options=predictions_df[iterator].unique(),
+        key=AnalyseResultsStateKeys.PlotIterations,
+        default=predictions_df[iterator].unique().tolist()[0],
+        help="Select the iteration or fold for which you want to display the parity plot.",
+    )
+
+    # Derive the available models from the prediction columns so both GNN and
+    # TML models are listed (predictions.csv merges both frameworks).
+    models = [
+        col_name.split(" ")[0]
+        for col_name in predictions_df.columns
+        if ResultColumn.PREDICTED in col_name
+    ]
+    model_name = st.multiselect(
+        "Select the model to display the confusion matrix",
+        options=models,
+        default=models[0],
+        key=AnalyseResultsStateKeys.PlotModels,
+        format_func=prettify_label,
+        help="Select the model for which you want to display the confusion matrix.",
+    )
+
+    model_pred_cols = [
+        get_predicted_label_column_name(
+            target_variable_name=data_options.target_variable_name, model_name=name
+        )
+        for name in model_name
+    ]
+
+    model_true_cols = get_true_label_column_name(
+        target_variable_name=data_options.target_variable_name
+    )
+
+    set_name = st.multiselect(
+        "Select the set to display parity plot",
+        options=predictions_df[ResultColumn.SET].unique(),
+        key=AnalyseResultsStateKeys.PlotSet,
+        default=[DataSet.Test],
+        help="Select the set for which you want to display the parity plot.",
+    )
+
+    plot_data = reshape_predictions(
+        df=predictions_df,
+        pred_cols=model_pred_cols,
+        target_variable_name=data_options.target_variable_name,
+    )
+
+    model_pred_col = get_predicted_label_column_name(
+        target_variable_name=data_options.target_variable_name, model_name=None
+    )
+
+    plot_data = plot_data[
+        (plot_data[iterator].isin(iteration))
+        & (plot_data[ResultColumn.SET].isin(set_name))
+        & (plot_data[ResultColumn.MODEL].isin(model_name))
+    ]
+
+    with st.expander("Show Data", expanded=False):
+        st.markdown(
+            f"**Showing data for model: {model_name}, set: {set_name}, iterations: {iteration}**"
+        )
+        st.markdown(
+            f"**True label column:** `{model_true_cols}` | **Predicted label column:** `{model_pred_col}`**"
+        )
+        st.dataframe(plot_data)
+
+    plot_config = get_plot_customisation_form(
+        plot_type=Plot.ConfusionMatrix,
+        data=plot_data,
+        data_options=data_options,
+        color_by_opts=[None],
+        model_pred_col=model_pred_col,
+        model_true_cols=model_true_cols,
+    )
+    # Not a label surface for this single-model plot; drop so it isn't splatted.
+    plot_config.pop("abbreviate_labels", None)
+
+    if not plot_data.empty:
+
+        # invert preds and true to match x and y axes title
+        fig = plot_confusion_matrix(
+            y_true=plot_data[model_pred_col], y_pred=plot_data[model_true_cols], **plot_config
+        )
+
+        return fig
+
+
+def parity_plot_form(
+    predictions_df: pd.DataFrame, split_type: SplitType | str, data_options: DataConfig
+):
+
+    color_by_opts = [None]
+
+    iterator = get_iterator_name(split_type)
+
+    iteration = st.multiselect(
+        f"Select the {iterator} to display parity plot",
+        options=predictions_df[iterator].unique(),
+        key=AnalyseResultsStateKeys.PlotIterations,
+        default=predictions_df[iterator].unique().tolist()[0],
+        help="Select the iteration or fold for which you want to display the parity plot.",
+    )
+
+    if len(iteration) > 1:
+        color_by_opts.append(iterator)
+
+    # Derive the available models from the prediction columns so both GNN and
+    # TML models are listed (predictions.csv merges both frameworks).
+    models = [
+        col_name.split(" ")[0]
+        for col_name in predictions_df.columns
+        if ResultColumn.PREDICTED in col_name
+    ]
+
+    model_name = st.multiselect(
+        "Select the model to display parity plot",
+        options=models,
+        default=models[0],
+        key=AnalyseResultsStateKeys.PlotModels,
+        format_func=prettify_label,
+        help="Select the model for which you want to display the parity plot.",
+    )
+
+    if len(model_name) > 1:
+        color_by_opts.append(ResultColumn.MODEL)
+
+    model_pred_cols = [
+        get_predicted_label_column_name(
+            target_variable_name=data_options.target_variable_name, model_name=name
+        )
+        for name in model_name
+    ]
+
+    model_true_cols = get_true_label_column_name(
+        target_variable_name=data_options.target_variable_name
+    )
+
+    set_name = st.multiselect(
+        "Select the set to display parity plot",
+        options=predictions_df[ResultColumn.SET].unique(),
+        key=AnalyseResultsStateKeys.PlotSet,
+        default=[DataSet.Test],
+        help="Select the set for which you want to display the parity plot.",
+    )
+
+    if len(set_name) > 1:
+        color_by_opts.append(ResultColumn.SET)
+
+    plot_data = reshape_predictions(
+        df=predictions_df,
+        pred_cols=model_pred_cols,
+        target_variable_name=data_options.target_variable_name,
+    )
+
+    model_pred_col = get_predicted_label_column_name(
+        target_variable_name=data_options.target_variable_name, model_name=None
+    )
+
+    plot_data = plot_data[
+        (plot_data[iterator].isin(iteration))
+        & (plot_data[ResultColumn.SET].isin(set_name))
+        & (plot_data[ResultColumn.MODEL].isin(model_name))
+    ]
+
+    with st.expander("Show Data", expanded=False):
+        st.markdown(
+            f"**Showing data for model: {model_name}, set: {set_name}, iterations: {iteration}**"
+        )
+        st.markdown(
+            f"**True label column:** `{model_true_cols}` | **Predicted label column:** `{model_pred_col}`**"
+        )
+        st.dataframe(plot_data)
+
+    plot_config = get_plot_customisation_form(
+        plot_type=Plot.Parity,
+        data=plot_data,
+        data_options=data_options,
+        color_by_opts=color_by_opts,
+        model_pred_col=model_pred_col,
+        model_true_cols=model_true_cols,
+    )
+    # Already applied to the hue/style legend inside the customiser; drop the
+    # flag so it isn't splatted into plot_parity.
+    plot_config.pop("abbreviate_labels", None)
+
+    if not plot_data.empty:
+
+        fig = plot_parity(
+            y_true=plot_data[model_true_cols], y_pred=plot_data[model_pred_col], **plot_config
+        )
+
+        return fig
+
+
+def extract_model_name_and_target(col_name, target_variable_name):
+    """Extract model name from column like 'Model Predicted target'."""
+    pattern = rf"^(.*?)\s+Predicted\s+{re.escape(target_variable_name)}$"
+    match = re.match(pattern, col_name)
+    if match:
+        return match.group(1)
+    return None
+
+
+def reshape_predictions(
+    df: pd.DataFrame, pred_cols: list, target_variable_name: str
+) -> pd.DataFrame:
+    # Columns to keep (everything that’s not a prediction column)
+    id_vars = [col for col in df.columns if col not in pred_cols]
+
+    # Melt the DataFrame
+    melted_df = df.melt(
+        id_vars=id_vars, value_vars=pred_cols, var_name="Prediction_Column", value_name="Prediction"
+    )
+
+    # Extract model name from the column name
+    melted_df[ResultColumn.MODEL] = melted_df["Prediction_Column"].apply(
+        lambda col: extract_model_name_and_target(col, target_variable_name)
+    )
+
+    new_col_name = get_predicted_label_column_name(
+        target_variable_name=target_variable_name, model_name=None
+    )
+
+    # Create new column name (without model name)
+    melted_df[new_col_name] = melted_df["Prediction"]
+
+    # Drop the old column
+    melted_df = melted_df.drop(columns=["Prediction_Column", "Prediction"])
+
+    return melted_df