plotting.py

"""
Author: Benjamin Garrard

This script houses code to make it easier to work with the TSNE matrix
that is generated by scikit-learn.  However this script is solely for
handling the plotting and prep for plotting.  It does not do the computation
of TSNE.  That needs to be used in another script.
"""
import pandas as pd
import plotly.offline as pyo
import plotly.graph_objs as go

from cvjson.cvj import CVJ

# These are colors to attempt to stay uniform through out plotting.
colors = {0 : "#d64933",
          1 : "#7c7c7c",
          2 : "#eee5e9",
          3 : "#92dce5",
          4 : "#000000",
          5 : "#f56960",
          6 : "#c46d5e",
          7 : "#daa588",
          8 : "#92dce5",
          9 : "#f3c98b",
          10 : "#9b5094",
          11 : "#9a94bc",
          12 : "#daa588",
          13 : "#92dce5",
          14 : "#cdf7f6",
          15 : "#4c5b61",
          16 : "#c5c5c5",
          17 : "#009fb7",
          18 : "#272727",
          19 : "#fed766",
          20 : "#610345",
          21 : "#107e7d",
          22 : "#044b7f",
          23 : "#e3b505",
          24 : "#dbd56e",
          25 : "#fc7753",
          26 : "#f2efea",
          27 : "#66d7d1",
          28 : "#080357",
          29 : "#d6ffb7",
          30 : "#f5ff90",
          31 : "#ff9f1c"}

marker_size = 2

def threeD_data(data, labels):
    """
    This method needs to receive the data from the TSNE matrix when the components are
    reduced to three.  This method just helps organize the three dimensional data.

    Parameters
    ----------
    data: numpy array
            This is the matrix from TSNE using scikit-learn

    labels: numpy array
            This is the label data from the y_train or y_test set.

    Returns
    -------
    df: DataFrame object
        This dataframe object comes from pandas and houses the x data given
        from TSNe and the y data given from the labels.
    """
    df = pd.DataFrame(data)
    df.columns = ["x", "y", "z"]

    df['labels'] = labels

    return df

def twoD_data(data, labels):
    """
    This method needs to receive the data from the TSNE matrix when the components are
    reduced to two.  This method just helps organize the two dimensional data.

    Parameters
    ----------
    data: numpy array
            This is the matrix from TSNE using scikit-learn

    labels: numpy array
            This is the label data from the y_train or y_test set.

    Returns
    -------
    df: DataFrame object
        This dataframe object comes from pandas and houses the x data given
        from TSNe and the y data given from the labels.
    """
    df = pd.DataFrame(data)
    df.columns = ["x", "y"]

    df['labels'] = labels

    return df

def plot_3d(df, layout=None, file_name=None, jupyter=True, auto=False):
    """
    This method needs to receives the dataframe from the ThreeD_data() method.  
    It then creates a 3D plot using the plotly library.

    Parameters
    ----------
    df: DataFrame object
            This dataframe object comes from ThreeD_data()

    layout: layout object
            This a layout object from plotly.  This is used if someone wants to manipulate the
            layout that is generated using plotly.

    file_name: str
            This the file name/file path where the file should be saved.

    jupyter: bool
            If this is true then this method will try to plot inside of a jupyter notebook.
    
    auto: bool
            If this is true then the plot from plotly will automatically open in the browser.

    """
    labels = df["labels"].unique()
    traces = []
    for label in labels:
        traces.append(go.Scatter3d(x=df[df['labels'] == label]['x'], y=df[df['labels'] == label]['y'], z=df[df['labels'] == label]['z'],
                        mode='markers', name=str(label), marker={"color": colors[label],
                                                                 "size" : marker_size}
                                  ))

    data = go.Data(traces)

    if layout != None:
        data = go.Figure(data=data, layout=layout)

    if file_name == None:
        if layout != None:
            file_name = layout.title
        else:
            file_name = "3D_plot"

    if jupyter:
        pyo.iplot(data, filename=file_name +".html", auto_open=auto)
    else:
        pyo.plot(data, filename=file_name +".html", auto_open=auto)
    
def plot_2d(df, layout=None, file_name=None, jupyter=True, auto=False):
    """
    This method needs to receives the dataframe from the ThreeD_data() method.  
    It then creates a 2D plot using the plotly library.

    Parameters
    ----------
    df: DataFrame object
            This dataframe object comes from TwoD_data()

    layout: layout object
            This a layout object from plotly.  This is used if someone wants to manipulate the
            layout that is generated using plotly.

    file_name: str
            This the file name/file path where the file should be saved.

    jupyter: bool
            If this is true then this method will try to plot inside of a jupyter notebook.
    
    auto: bool
            If this is true then the plot from plotly will automatically open in the browser.

    """
    labels = df["labels"].unique()
    traces = []
    for label in labels:
        traces.append(go.Scatter(x=df[df['labels'] == label]['x'], y=df[df['labels'] == label]['y'],
                        mode='markers', name=str(label), marker={"color": colors[label],
                                                                 "size" : marker_size}
                                  ))

    data = go.Data(traces)

    if layout != None:
        data = go.Figure(data=data, layout=layout)

    if file_name == None:
        if layout != None:
            file_name = layout.title
        else:
            file_name = "2D_plot"

    if jupyter:
        pyo.iplot(data, filename=file_name +".html", auto_open=auto)
    else:
        pyo.plot(data, filename=file_name +".html", auto_open=auto)