data_visualization.py

import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.io as pio
from matplotlib import pyplot as plt
from plotly.subplots import make_subplots


def get_tab10_color(i: int) -> tuple:
    """
    returns colors based on tableau color scheme
    :param i: index of tableau color scheme
    :return: tuple of one of the 10 colors based on the tableau color scheme
    """
    cmap = plt.get_cmap('tab10')
    return cmap(i % 10)


def get_colorblindfriendly_color(i: int) -> tuple:
    """
    returns colorblindfriendly colors
    """
    list = ['rgb(253,174,97)', 'rgb(116,173,209)', 'rgb(215,48,39)', 'rgb(244,109,67)', 'rgb(254,224,144)',
            'rgb(255,255,191)',
            'rgb(224,243,248)', 'rgb(171,217,233)', 'rgb(69,117,180)']
    return list[i]


def add_p_value_annotation(fig, array_columns, annotations=None, subplot=None,
                           _format=dict(interline=0.07, text_height=1.07, color='black')):
    ''' Adds notations to boxplots
    Parameters:
    ----------
    fig: figure
        plotly boxplot figure
    array_columns: np.array
        array of which columns to compare
        e.g.: [[0,1], [1,2]] compares column 0 with 1 and 1 with 2
    subplot: None or int
        specifies if the figures has subplots and what subplot to add the notation to
    _format: dict
        format characteristics for the lines

    Returns:
    -------
    fig: figure
        figure with the added notation
    '''
    # Specify in what y_range to plot for each pair of columns
    y_range = np.zeros([len(array_columns), 2])
    for i in range(len(array_columns)):
        y_range[i] = [1 + 0.4 * i * _format['interline'], 1.01 + 0.4 * i * _format['interline']]

    # Get values from figure
    fig_dict = fig.to_dict()

    # Get indices if working with subplots
    if subplot:
        if subplot == 1:
            subplot_str = ''
        else:
            subplot_str = str(subplot)
        indices = []  # Change the box index to the indices of the data for that subplot
        for index, data in enumerate(fig_dict['data']):
            # print(index, data['xaxis'], 'x' + subplot_str)
            if data['xaxis'] == 'x' + subplot_str:
                indices = np.append(indices, index)
        indices = [int(i) for i in indices]
        print((indices))
    else:
        subplot_str = ""

    # Print the p-values
    for index, column_pair in enumerate(array_columns):
        # Mare sure it is selecting the data and subplot you want
        # print('0:', fig_dict['data'][data_pair[0]]['name'], fig_dict['data'][data_pair[0]]['xaxis'])
        # print('1:', fig_dict['data'][data_pair[1]]['name'], fig_dict['data'][data_pair[1]]['xaxis'])

        # Get the p-value
        symbol = annotations[index]
        # Vertical line
        fig.add_shape(type="line",
                      xref="x" + subplot_str, yref="y" + subplot_str + " domain",
                      x0=column_pair[0], y0=y_range[index][0],
                      x1=column_pair[0], y1=y_range[index][1],
                      line=dict(color=_format['color'], width=2, )
                      )
        # Horizontal line
        fig.add_shape(type="line",
                      xref="x" + subplot_str, yref="y" + subplot_str + " domain",
                      x0=column_pair[0], y0=y_range[index][1],
                      x1=column_pair[1], y1=y_range[index][1],
                      line=dict(color=_format['color'], width=2, )
                      )
        # Vertical line
        fig.add_shape(type="line",
                      xref="x" + subplot_str, yref="y" + subplot_str + " domain",
                      x0=column_pair[1], y0=y_range[index][0],
                      x1=column_pair[1], y1=y_range[index][1],
                      line=dict(color=_format['color'], width=2, )
                      )
        ## add text at the correct x, y coordinates || * _format['text_height']
        ## for bars, there is a direct mapping from the bar number to 0, 1, 2...
        fig.add_annotation(dict(font=dict(color=_format['color'], size=14),
                                x=(column_pair[0] + column_pair[1]) / 2,
                                y=y_range[index][1] + 0.0225 * y_range[index][1],
                                showarrow=False,
                                text=symbol,
                                textangle=0,
                                xref="x" + subplot_str,
                                yref="y" + subplot_str + " domain"
                                ))
    return fig


def plot_multiple_box_plotly(data: list, rows: int, cols: int, x_labels: list, yaxis_title: str,
                             xaxis_title: str = None, theme: str = 'plotly_white',
                             titles: list = None, outliers='outliers', ymin=None, ymax=None, outlier_opacity=1,
                             annotations_pos=None, annotations_text=None, x_label_rotation=0, legend=False,
                             savefig=False, path=None, width=550, height=400, fontsize=9):
    fig = make_subplots(rows=rows, cols=cols, shared_yaxes=True, subplot_titles=titles, horizontal_spacing=0.01)

    # If p-value annotations are given change position of subplot titles
    if annotations_pos:
        ann_list = [x for x in annotations_pos if x is not None]
        list_len = [len(i) for i in ann_list]
        y_offset = max(list_len) * 0.03
        for i in fig['layout']['annotations']:
            i['y'] = 1 + y_offset
            i['font'] = dict(family="CMU Serif", size=fontsize * 0.8, color='black')

    for j in range(len(data)):
        for i in range(len(data[j])):
            # print(len(data[j][i]))
            showlegend = False
            if legend:
                showlegend = True if j == 0 else False
            if len(data[j][i]) == 0:
                print('empty data in [{j}, {i}]'.format(j=j, i=i))
                fig.add_trace(
                    go.Box(y=[0], name=x_labels[i], boxpoints=outliers,
                           line=dict(color=get_colorblindfriendly_color(i)),
                           marker=dict(opacity=outlier_opacity), boxmean=True, showlegend=showlegend),
                    row=1, col=j + 1)

            else:
                fig.add_trace(
                    go.Box(y=data[j][i], name=x_labels[i], boxpoints=outliers,
                           line=dict(color=get_colorblindfriendly_color(i)),
                           marker=dict(opacity=outlier_opacity), boxmean=True, showlegend=showlegend),
                    row=1, col=j + 1)
        if annotations_pos:
            if annotations_pos[j]:
                fig = add_p_value_annotation(fig, annotations_pos[j], annotations_text[j], subplot=j + 1)

    fig.update_xaxes(tickangle=x_label_rotation)
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')

    # Set font and color for annotations
    for i in fig['layout']['annotations']:
        i['font'] = dict(family="CMU Serif", size=20, color='black')

    fig.update_layout(font=dict(family="CMU Serif", size=fontsize, color='black'),
                      yaxis_title=yaxis_title, xaxis_title=xaxis_title, yaxis_range=[ymin, ymax],
                      template=theme)  # , margin=dict(l=0, r=0, t=20, b=0)
    if savefig:
        pio.write_image(fig=fig, file=path, format='pdf', width=width, height=height)
    fig.show()


def plot_box_plotly(data: list, labels: list, yaxis_title: str = None, xaxis_title: str = None,
                    theme: str = 'plotly_white', title: str = None, outliers='outliers', showlegend: bool = False,
                    horizontal: bool = False, colors=None, ymin=None, ymax=None, annotations_pos=None,
                    outlier_opacity=1,
                    annotations_text=None, plot_stats=False, savefig=False, path=None, width=1000, height=800):
    fig = go.Figure()

    if horizontal:
        for i in range(len(data)):
            fig.add_trace(
                go.Box(x=data[i], name=labels[i], boxpoints=outliers, boxmean=True,
                       marker=dict(opacity=outlier_opacity),
                       line=dict(color=colors[i])))
    else:
        for i in range(len(data)):
            print(colors[i])
            fig.add_trace(
                go.Box(y=data[i], name=labels[i], boxpoints=outliers, boxmean=True,
                       marker=dict(opacity=outlier_opacity),
                       line=dict(color=colors[i])))
            print(data[i].quantile([0.25, 0.5, 0.75]).values)

            if plot_stats:
                ## loop through the values you want to label and add them as annotations
                for x in zip(["q1", "med", "q3"], data[i].quantile([0.25, 0.5, 0.75]).values.round(2)):
                    fig.add_annotation(

                        font=dict(size=18),
                        x=i + 0.25,
                        y=x[1],
                        text=x[0] + ":" + str(x[1]),
                        showarrow=False,
                        xanchor="left",
                    )

    if annotations_pos:
        fig = add_p_value_annotation(fig, annotations_pos, annotations_text)

    fig.update_layout(font=dict(family="CMU Serif", size=28, color='black'),
                      title=dict(text=title, xanchor='center', x=0.5),
                      yaxis=dict(title=yaxis_title, showgrid=True, mirror=True, showline=True, ticks='outside'),
                      xaxis=dict(title=xaxis_title, showgrid=True, mirror=True, showline=True, ticks='outside'),
                      template=theme, bargap=0, yaxis_range=[ymin, ymax], showlegend=showlegend,
                      margin=go.layout.Margin(l=0, r=0, b=50, t=80, ))

    if savefig:
        pio.write_image(fig=fig, file=path, format='pdf', width=width, height=height)

    fig.show()


if __name__ == "__main__":
    # Set plot colors
    plot_colors = ['#b8b0ac', '#6a9f58', '#e49444']

    # Generate plot for scores
    scores = pd.read_csv("data/scores_processed.csv")
    metric = 'score'
    plot_box_plotly([scores[scores['ams'] == 'no_ams'][metric],
                     scores[scores['ams'] == 'gaze_bubble'][metric],
                     scores[scores['ams'] == 'ambient'][metric]],
                    labels=['No AMS', 'Gaze Bubble', 'Ambient'], yaxis_title='Score',
                    outliers='suspectedoutliers', ymin=0, ymax=90,
                    annotations_pos=[[0, 1], [1, 2]],
                    annotations_text=['***', '***'], outlier_opacity=1,
                    colors=plot_colors,
                    plot_stats=False,
                    savefig=False,
                    path='gfx/score_comparison.pdf',
                    width=600, height=800)

    # Generate plot for pupil size
    pupils = pd.read_csv("data/pupils_processed.csv")
    metric = 'pupil_size'
    plot_box_plotly([pupils[pupils['ams'] == 'no_ams'][metric],
                     pupils[pupils['ams'] == 'gaze_bubble'][metric],
                     pupils[pupils['ams'] == 'ambient'][metric]],
                    labels=['No AMS', 'Gaze Bubble', 'Ambient'], yaxis_title='Pupil Size in mm',
                    outliers='suspectedoutliers', ymin=2, ymax=4,
                    annotations_pos=[],
                    annotations_text=[], outlier_opacity=1,
                    colors=plot_colors,
                    plot_stats=False,
                    savefig=False,
                    path='gfx/pupil_size_comparison.pdf',
                    width=600, height=800)

    # Generate plot for AMS adherence
    plot_colors_two = ['#6a9f58', '#e49444']
    follows = pd.read_csv("data/follows_processed.csv")
    metric = 'ams_follow'
    plot_box_plotly([follows[follows['ams'] == 'gaze_bubble'][metric],
                     follows[follows['ams'] == 'ambient'][metric]],
                    labels=['Gaze Bubble', 'Ambient'], yaxis_title='Following Rate in %',
                    outliers='suspectedoutliers', ymin=20, ymax=90,
                    annotations_pos=[[0, 1]],
                    annotations_text=['***'], outlier_opacity=1,
                    colors=plot_colors_two,
                    plot_stats=False,
                    savefig=False,
                    path='gfx/follows_comparison.pdf',
                    width=600, height=800)


    # Score Plot Interactions
    metric = 'score'
    fig1 = [scores[scores['condition'] == 'no_ams_e'][metric], scores[scores['condition'] == 'no_ams_h'][metric]]
    fig2 = [scores[scores['condition'] == 'gaze_bubble_e'][metric], scores[scores['condition'] == 'gaze_bubble_h'][metric]]
    fig3 = [scores[scores['condition'] == 'ambient_e'][metric], scores[scores['condition'] == 'ambient_h'][metric]]


    plot_multiple_box_plotly([fig1, fig2, fig3], rows=1, cols=3,
                             x_labels=['Easy', 'Hard'], yaxis_title='Score',
                             titles=['No AMS', 'Gaze Bubble', 'Ambient'],
                             outliers='suspectedoutliers', fontsize=20, ymin=0, ymax=90,
                             outlier_opacity=1,
                             savefig=False,
                             path='gfx/test.pdf',
                             width=600, height=800)

    # Pupil Size  Plot Interactions
    metric = 'pupil_size'
    fig1 = [pupils[pupils['condition'] == 'no_ams_e'][metric], pupils[pupils['condition'] == 'no_ams_h'][metric]]
    fig2 = [pupils[pupils['condition'] == 'gaze_bubble_e'][metric], pupils[pupils['condition'] == 'gaze_bubble_h'][metric]]
    fig3 = [pupils[pupils['condition'] == 'ambient_e'][metric], pupils[pupils['condition'] == 'ambient_h'][metric]]


    plot_multiple_box_plotly([fig1, fig2, fig3], rows=1, cols=3,
                             x_labels=['Easy', 'Hard'], yaxis_title='Pupil Size in mm',
                             titles=['No AMS', 'Gaze Bubble', 'Ambient'],
                             outliers='suspectedoutliers', fontsize=20, ymin=1.5, ymax=4,
                             savefig=False,
                             path='gfx/test.pdf',
                             width=600, height=800)


    # Following Rate Interactions
    metric = 'ams_follow'
    fig1 = [follows[follows['condition'] == 'gaze_bubble_e'][metric], follows[follows['condition'] == 'gaze_bubble_h'][metric]]
    fig2 = [follows[follows['condition'] == 'ambient_e'][metric], follows[follows['condition'] == 'ambient_h'][metric]]


    plot_multiple_box_plotly([fig1, fig2], rows=1, cols=2,
                             x_labels=['Easy', 'Hard'], yaxis_title='Follwoing Rate in %',
                             titles=['Gaze Bubble', 'Ambient'],
                             outliers='suspectedoutliers', fontsize=20, ymin=20, ymax=90,
                             savefig=False,
                             path='gfx/test.pdf',
                             width=600, height=800)