plots.py

"""
Contains the functions that plot various results using matplotlib and seaborn.
"""

from analysis import Analysis

import matplotlib
import seaborn as sns
import numpy as np
import pandas as pd
import string
import plotly.graph_objs as go

matplotlib.use('TkAgg')
sns.set_theme()
sns.set_context("paper")
sns.set(font_scale=1.2)


def plot_all_scenarios_all_analyses(analyses: [Analysis], color_cycle: int = None):
    """
    Plots the results of many model outputs in one figure. In one figure, creates four columns, one for each data type.
    Each row is an analysis. Returns a figure.
    """

    # Get the right color and dashes for plot decomposition
    cm = sns.color_palette("Paired")
    cp = ["#70381D", cm[7], "#74B6C2"]
    cp = cp if color_cycle is None else [cp[color_cycle]]

    # Build the dataframes for each category then melt them
    all_data = pd.DataFrame()
    for analysis in analyses:
        data = analysis.model_result.get_dataframe(analysis.scenario.title, analysis._variable_title)
        all_data = pd.concat([all_data, data]).reset_index(drop=True)

    melted_data = all_data.melt(id_vars=["Years from start", "Scenario", "Analysis type"],
                                value_vars=["POC", "DOC", "DIC", "Cells"], var_name="Data type")

    # Plot
    grid = sns.relplot(data=melted_data, x="Years from start", y="value", palette=cp, height=10, aspect=1.5,
                       col="Data type", row="Analysis type", hue="Scenario", kind="line", #markers=True, style="Scenario",
                       facet_kws={'sharey': False, 'sharex': False, "margin_titles": True, "legend_out": False})

    grid.set(xscale="log")
    grid.set_titles(template="{col_name} over time", row_template="{row_name}", col_template="{col_name} over time")
    grid._legend._loc = 4


    # Tune relplot labels
    y_labels = ["Femtograms C/mL"] * 3 + ["Cells/mL"]
    y_lim = [[1, 10 ** 14]] * 3 + [[1, 10 ** 9]]
    for i, ax in enumerate(grid.axes.ravel()):
        if i < 4:
            label = y_labels[i % len(y_labels)]
            ax.set_ylabel(label)
            ax.set_xlabel("Years from start")
        else:
            ax.set_ylabel("")
            ax.set_xlabel("")

        ax.set_yscale("log")
        ax.set_xscale("log")

        ylim = y_lim[i % len(y_lim)]
        xlim = [1, 4 * 10 ** 4]
        ax.set_ylim(ylim)
        ax.set_xlim(xlim)

        # Add panel title to row
        if i % 4 == 0:
            ax.text(-0.1, 1.1, string.ascii_uppercase[i//4], transform=ax.transAxes, size=20, weight='bold')

    grid.tight_layout()
    grid.fig.subplots_adjust(hspace=0.3, wspace=0.4)

    return grid


def plot_one_result_type_all_analyses(analyses: [Analysis], data_type: str, main_index: int = 0, color_cycle: int = None):
    """
    Plots one data type result from each analysis.  Data_type string sets which result type to plot.
    If main_index is set to an index in Analyses, that result is plotted across multiple rows in its own column.
    Color_cycle shifts the color map to the selected index, used to plot one type of scenario.
    Returns a figure.
    """
    # Get the right color and dashes for plot decomposition
    cp = ["#70381D", "#FF7F00", "#74B6C2"]
    cp = cp if color_cycle is None else [cp[color_cycle]]

    # Move the main_analysis to front
    analyses.insert(0, analyses.pop(main_index))

    # Build the dataframes for each category then melt them
    all_data = pd.DataFrame()
    for analysis in analyses:
        data = analysis.model_result.get_dataframe(analysis.scenario.title, analysis._variable_title)
        all_data = pd.concat([all_data, data]).reset_index(drop=True)

    melted_data = all_data.melt(id_vars=["Years from start", "Scenario", "Analysis type"],
                                value_vars=[data_type], var_name="Data type")

    # Plot
    grid = sns.relplot(data=melted_data, x="Years from start", y="value", palette=cp, height=10, aspect=1.5,
                       col="Analysis type", hue="Scenario", kind="line", col_wrap=4, #markers=True, style="Scenario",
                       facet_kws={'sharey': True, 'sharex': True, "legend_out": False})

    grid._legend._loc = 4
    grid.set_titles(template="{col_name}")
    grid.set(xscale="log")

    # Tune relplot labels
    y_label = "Cells/mL" if data_type == "Cells" else "Femtograms C/mL"
    y_lim = [1, 10 ** 9] if data_type == "Cells" else [1, 10 ** 14]
    for i, ax in enumerate(grid.axes.ravel()):
        ax.set_ylabel(y_label)

        ax.set_yscale("log")
        ax.set_xscale("log")
        ax.set_ylim(y_lim)
        ax.set_xlim([1, 40000])

        # Add panel title
        ax.text(-0.1, 1.1, string.ascii_uppercase[i], transform=ax.transAxes, size=20, weight='bold')

    grid.tight_layout()
    grid.fig.subplots_adjust(wspace=0.2)

    return grid


def plot_sensitivity(analysis: Analysis):
    """
    Plots the results of a sensitivity analysis as a bar growth per variable. Includes both total and first order
    Sobol indices. Returns a figure.
    """

    # Check there are results
    if analysis.sensitivity_analysis_result is None:
        return None

    # Get dataframes
    df = analysis.sensitivity_analysis_result.get_dataframe(analysis.scenario).set_index('Parameter')

    # We're only plotting mean
    df_mean = df[df.Output.eq('Mean')]

    vals = df_mean[['Total-effect', 'First-order']]
    yerr = df_mean[['Total Error', 'First Error']]
    yerr.columns = vals.columns

    vals = vals.apply(lambda x: np.where(x < 0.01, 0, x))

    # Use pandas plot function
    ax = vals.plot(kind='bar', yerr=yerr, figsize=(10, 6), ylabel="Value", width=0.8, rot=0)

    # Add bar label
    label_colors = ["white", "white", "white", "white", "white", "white", "black",
                    "white", "white", "black", "black", "black", "white", "black"]
    for c in ax.containers:
        if type(c) == matplotlib.container.BarContainer:
            labels = [f'{h:.2f}' if (h := v.get_height()) >= 0.01 else "<0.01" for v in c]
            ax.bar_label(c, labels=labels, label_type='center', padding=5, fontsize=12)

    for i, text in enumerate(ax.texts):
        text.set_color(label_colors[i])

    return ax.get_figure()


def plot_one_result_type_all_analyses_interactive_plotly(analyses, data_type, main_index=0, color_cycle=None):
    # Your color palette
    cp = ["#70381D", "#FF7F00", "#74B6C2", "#9E8576", "#FF9999", "#A3CCE9"]
    cp = cp if color_cycle is None else [cp[color_cycle]]

    # Move the main_analysis to front
    analyses.insert(0, analyses.pop(main_index))

    # Initialize a Plotly figure
    fig = go.Figure()

    # Iterate over analyses and scenarios to add traces to the figure
    for idx, analysis in enumerate(analyses):
        data = analysis.model_result.get_dataframe(analysis.scenario.title, analysis._variable_title)
        for scenario in data['Scenario'].unique():
            scenario_data = data[data['Scenario'] == scenario]

            # Use colors from your palette
            color = cp[idx % len(cp)]

            fig.add_trace(go.Scatter(
                x=scenario_data['Years from start'],
                y=scenario_data[data_type],
                mode='lines+markers',
                name=analysis.scenario.title,
                line=dict(color=color)
            ))

    # Set log scale for axes
    fig.update_layout(yaxis_type="log")  # xaxis_type="log"

    # Set axes to use scientific notation
    fig.update_xaxes(tickformat=".5e")
    fig.update_yaxes(tickformat=".5e")

    # Customize layout, titles, etc.
    fig.update_layout(title="", xaxis_title="Years from start", yaxis_title=data_type)

    return fig