utils/graphs.py

import altair as altair
import pandas as pd

class Altair():
    def plot_single_bar(self, df, x, y, x_legend, y_legend, width, color, title):
        if color is None:
            color = '#d62728'
        graph1 = altair.Chart(df).mark_bar(opacity=0.7, cornerRadius=7).encode(
            altair.X('{0}:N'.format(x), title=x_legend),
            altair.Y('{0}:Q'.format(y), title=y_legend),
            altair.OpacityValue(0.77),
            color=altair.value(color),
            tooltip=[altair.Tooltip('{0}:N'.format(x)),
                     altair.Tooltip('{0}:Q'.format(y))]).properties(width=width, height=185)
        line = altair.Chart(df).mark_line(color='red').transform_window(
            rolling_mean='mean({0})'.format(y), frame=[-9, 0]).encode(
            x='{0}:O'.format(x),
            y='rolling_mean:Q')

        return altair.layer(graph1 + line)

    def plot_neg_pos_bars(self, df, x, y, x_legend, y_legend, width, title):
        graph = altair.Chart(df).mark_bar(opacity=0.78, cornerRadius=7).encode(
            altair.X('{0}:N'.format(x), title=x_legend),
            altair.Y('{0}:Q'.format(y), title=y_legend),
            color=altair.condition(
                altair.datum.y > 0,
                altair.value("steelblue"),  # The positive value color
                altair.value("pink")      # The negative value color
            ),
            tooltip = [altair.Tooltip('{0}:N'.format(x)),
                   altair.Tooltip('{0}:Q'.format(y))]).properties(width=width, height=185)
        line = altair.Chart(df).mark_line(color='red').transform_window(
            rolling_mean='mean({0})'.format(y), frame=[-9, 0]).encode(
            x='{0}:O'.format(x),
            y='rolling_mean:Q')
        return altair.layer(graph + line)


    def plot_single_line(self, df, x, y, x_legend, y_legend, width, color, title):
        if color is None:
            color = '#d62728'
        graph = altair.Chart(df).mark_line(opacity=0.7, cornerRadius=7).encode(
            altair.X('{0}:N'.format(x), title=x_legend),
            altair.Y('{0}:Q'.format(y), title=y_legend),
            altair.OpacityValue(0.77),
            color=altair.value(color),
            strokeWidth=altair.value(4),
            tooltip=[altair.Tooltip('{0}:N'.format(x)),
                     altair.Tooltip('{0}:Q'.format(y))]).properties(width=width)
        return altair.layer(graph)

    def plot_single_area(self, df, x, y, x_legend, y_legend, width, color, title):
        if color is None:
            color = '#d62728'
        graph = altair.Chart(df).mark_area(opacity=0.7, cornerRadius=7).encode(
            altair.X('{0}:N'.format(x), title=x_legend),
            altair.Y('{0}:Q'.format(y), title=y_legend),
            altair.OpacityValue(0.77),
            color=altair.value(color),
            strokeWidth=altair.value(4),
            tooltip=[altair.Tooltip('{0}:N'.format(x)),
                     altair.Tooltip('{0}:Q'.format(y))]).properties(width=width, height=200)

        line = altair.Chart(df).mark_line(color='red').transform_window(
            rolling_mean='mean({0})'.format(y), frame=[-9, 0]).encode(
            x='{0}:O'.format(x),
            y='rolling_mean:Q')
        return altair.layer(graph + line)

    def plot_donut(self, df, category, value, inner_radius, outer_radius):
        graph = altair.Chart(df).mark_arc(innerRadius=inner_radius, outerRadius=outer_radius).encode(
            theta=value,
            color="{0}:N".format(category),
            tooltip=[altair.Tooltip('{0}:N'.format(category)),
                     altair.Tooltip('{0}:Q'.format(value))])
        return altair.layer(graph)

    def generate_date_cols(self, df, dt_col_nm):
        dt_col_nm_frm = '{0}_frm'.format(dt_col_nm)
        df[dt_col_nm_frm] = pd.to_datetime(df[dt_col_nm], format='%Y-%m-%d')
        df['year'] = df[dt_col_nm_frm].dt.year.astype(str)
        df['month'] = df[dt_col_nm_frm].dt.month.astype(str)
        df['month'] = df['month'].str.zfill(2)
        df['year_month'] = df['year'].astype('string') + '-' + df['month'].astype('string')
        return df

    def generate_dates(self, df, dt_col_nm):
        df = self.generate_date_cols(df, dt_col_nm)
        years_months = []
        min_year = int(df['year'].min())
        max_year = int(df['year'].max())
        max_month = int(df['month'].max())
        year = min_year
        while year <= max_year:
            for month in range(1, 13):
                years_months.append('{0}-{1}'.format(str(year), str(month).zfill(2)))
                if month == max_month and year == max_year:
                    break
            year += 1
        df_ym = pd.DataFrame({'year_month': years_months})
        df_ym = pd.merge(df_ym, df, on='year_month', how='left')
        return df_ym