diff --git a/fanchart/plot.py b/fanchart/plot.py index 5d1a49c..c4403b4 100644 --- a/fanchart/plot.py +++ b/fanchart/plot.py @@ -1,5 +1,6 @@ import matplotlib.dates as mdates import matplotlib.pyplot as plt +import matplotlib as mpl import matplotlib.ticker as ticker import numpy as np import pandas as pd @@ -7,9 +8,24 @@ from twopiece.scale import tpnorm import pkg_resources -DARK_BG = "#212946" register_matplotlib_converters() +DARK_BG = "#212946" +dark_mpl_rc_context = { + "lines.color": "white", + "patch.edgecolor": "white", + "text.color": 'white', + "axes.facecolor": DARK_BG, + "axes.edgecolor": "lightgray", + "axes.labelcolor": "white", + "xtick.color": "white", + "ytick.color": "white", + "grid.color": "lightgray", + "figure.facecolor": DARK_BG, + "figure.edgecolor": DARK_BG, + "savefig.facecolor": DARK_BG, + "savefig.edgecolor": DARK_BG} + def load_boe_parameters(): """ @@ -46,7 +62,7 @@ def _get_alphas(p): def _color_setup(darkmode): - plt.rcParams.update(plt.rcParamsDefault) + # plt.rcParams.update(plt.rcParamsDefault) if darkmode: plt.rcParams.update({ "lines.color": "white", @@ -64,18 +80,13 @@ def _color_setup(darkmode): "savefig.edgecolor": DARK_BG}) -def _fan_single_customised(p, loc, sigma, gamma, kind='pdf', color='xkcd:tomato red', grid=False, fig_size=None, - darkmode=False): +def _fan_single_customised(p, loc, sigma, gamma, kind='pdf', color='xkcd:tomato red', hide_xgrid=True, + **fig_kw): dist = tpnorm(loc=loc, sigma=sigma, gamma=gamma, kind='boe') q = dist.ppf(p) x = np.linspace(loc - 4 * sigma, loc + 4 * sigma, 500) - _color_setup(darkmode) - if fig_size: - fs = fig_size - else: - fs = plt.rcParamsDefault['figure.figsize'] - fig = plt.figure(figsize=fs) + fig = plt.figure(**fig_kw) ax = fig.add_subplot(111) ax.spines[['top', 'right', 'left']].set_visible(False) ax.xaxis.set_ticks_position('bottom') @@ -102,10 +113,11 @@ def _fan_single_customised(p, loc, sigma, gamma, kind='pdf', color='xkcd:tomato ax.set_yticks(y_values) ax.set_yticklabels(y_labs) - ax.set_title(title, loc='right', fontdict={'fontsize': 10, 'fontweight': 'medium'}) + ax.set_title(title, loc='right') - if grid: + if hide_xgrid: ax.grid(axis='y', alpha=0.25) + ax.grid(axis='x', alpha=0.0) ax.set_axisbelow(True) alpha_fill = _get_alphas(p) @@ -120,11 +132,11 @@ def _fan_single_customised(p, loc, sigma, gamma, kind='pdf', color='xkcd:tomato plt.fill_between(x, y, alpha=alpha_fill[i], facecolor=color, edgecolor="none", linewidth=0.0) plt.show() - assert isinstance(fig, object) + return fig -def fan_single(loc, sigma, gamma, probs, kind='pdf', color='xkcd:tomato red'): +def fan_single(loc, sigma, gamma, probs, kind='pdf', color='xkcd:tomato red', **fig_kw): """ Returns a fan_single chart :param loc: Mode parameter (double) @@ -135,10 +147,11 @@ def fan_single(loc, sigma, gamma, probs, kind='pdf', color='xkcd:tomato red'): :param color: A color to make the chart :return: fanchart single (figure) """ - return _fan_single_customised(probs, loc, sigma, gamma, kind=kind, grid=True, color=color) + return _fan_single_customised(probs, loc, sigma, gamma, kind=kind, color=color, **fig_kw) -def fan_single_dark(loc, sigma, gamma, probs, kind='pdf', color='orange'): +@mpl.rc_context(dark_mpl_rc_context) +def fan_single_dark(loc, sigma, gamma, probs, kind='pdf', color='orange', **fig_kw): """ Returns a fan_single chart :param loc: Mode parameter (double) @@ -149,50 +162,35 @@ def fan_single_dark(loc, sigma, gamma, probs, kind='pdf', color='orange'): :param color: A color to make the chart :return: fanchart single (figure) """ - return _fan_single_customised(probs, loc, sigma, gamma, kind=kind, grid=True, color=color, darkmode=True) + + return _fan_single_customised(probs, loc, sigma, gamma, kind=kind, color=color, **fig_kw) -def _fan_customised(data, p, historic=None, color='cornflowerblue', grid=False, fig_size=None, title=None, - title_loc=None, - darkmode=False): +def _fan_customised(data, p, historic=None, color='cornflowerblue', title='CPI Inflation Projection', + title_loc='left', darkmode=False, hide_xgrid=True, + **fig_kw): dat = data.copy() dat['Date'] = pd.to_datetime(dat['Date']) dat = dat.reset_index() max_historical_value = 0 min_historical_value = 0 - _color_setup(darkmode) - if fig_size: - fs = fig_size - else: - fs = (12, 6) - fig = plt.figure(figsize=fs) + fig = plt.figure(**fig_kw) ax = fig.add_subplot(111) ax.spines[['top', 'right', 'left']].set_visible(False) - if title: - fig_title = title - else: - fig_title = 'CPI Inflation Projection' - if title_loc: - ax.set_title(fig_title, loc=title_loc, fontdict={'fontsize': 12, 'fontweight': 'medium'}) - else: - ax.set_title(fig_title, loc='left', fontdict={'fontsize': 12, 'fontweight': 'medium'}) + ax.set_title(title, loc=title_loc) ax.xaxis.set_ticks_position('bottom') ax.xaxis.set_major_locator(mdates.YearLocator()) ax.xaxis.set_minor_locator(mdates.MonthLocator((1, 4, 7, 10))) - if grid: - ax.grid(axis='y', alpha=0.25) - ax.set_axisbelow(True) - results = np.zeros((dat.shape[0], len(p))) - for index, _ in enumerate(results): - mode = dat['Mode'][index] - sigma = dat['Uncertainty'][index] - gamma = dat['Skewness'][index] + for idx, _ in enumerate(results): + mode = dat['Mode'][idx] + sigma = dat['Uncertainty'][idx] + gamma = dat['Skewness'][idx] dist = tpnorm(loc=mode, sigma=sigma, gamma=gamma, kind='boe') - results[index] = dist.ppf(p) + results[idx] = dist.ppf(p) results = pd.DataFrame(results, columns=p, index=dat['Date']) if historic is not None: @@ -213,10 +211,10 @@ def _fan_customised(data, p, historic=None, color='cornflowerblue', grid=False, results = results.sort_index() alpha_fill = _get_alphas(p) - for index, column in enumerate(results.columns): - if index < results.shape[1] - 1: - plt.fill_between(results.index, results.iloc[:, index], results.iloc[:, index + 1], - alpha=alpha_fill[index], facecolor=color, edgecolor="none", linewidth=0.0) + for idx, column in enumerate(results.columns): + if idx < results.shape[1] - 1: + plt.fill_between(results.index, results.iloc[:, idx], results.iloc[:, idx + 1], + alpha=alpha_fill[idx], facecolor=color, edgecolor="none", linewidth=0.0) if darkmode: ax.axhline(y=2.0, color='white', linestyle='--', linewidth=1.5) @@ -230,6 +228,11 @@ def _fan_customised(data, p, historic=None, color='cornflowerblue', grid=False, ax.set_ylim([round(min_inflation) - 2., round(max_inflation) + 2.]) ax.legend(loc='lower left', frameon=False) + if hide_xgrid: + ax.grid(axis='y', alpha=0.25) + ax.grid(axis='x', alpha=0.0) + ax.set_axisbelow(True) + ax1 = ax.twinx() ax1.set_yticks(ax.get_yticks()) ax1.set_ylim(ax.get_ylim()) @@ -237,11 +240,10 @@ def _fan_customised(data, p, historic=None, color='cornflowerblue', grid=False, plt.show() - assert isinstance(fig, object) return fig -def fan(pars, probs, historic=None, color='xkcd:tomato red', title=None): +def fan(pars, probs, historic=None, color='xkcd:tomato red', title='CPI Inflation Projection', **fig_kw): """ Returns a fanchart :param pars: dataframe - Parameters table with columns Mode, Uncertainty, Skewness @@ -251,11 +253,11 @@ def fan(pars, probs, historic=None, color='xkcd:tomato red', title=None): :param title: string - title for the chart :return: fanchart (figure) """ - return _fan_customised(pars, probs, historic=historic, color=color, grid=True, title=title, title_loc=None, - darkmode=False) + return _fan_customised(pars, probs, historic=historic, color=color, darkmode=False, title=title, **fig_kw) -def fan_dark(pars, probs, historic=None, color='orange', title=None): +@mpl.rc_context(dark_mpl_rc_context) +def fan_dark(pars, probs, historic=None, color='orange', title='CPI Inflation Projection', **fig_kw): """ Returns a fanchart in dark mode :param pars: dataframe - Parameters table with columns Mode, Uncertainty, Skewness @@ -265,5 +267,4 @@ def fan_dark(pars, probs, historic=None, color='orange', title=None): :param title: string - title for the chart :return: fanchart (figure) """ - return _fan_customised(pars, probs, historic=historic, color=color, grid=True, title=title, title_loc=None, - darkmode=True) + return _fan_customised(pars, probs, historic=historic, color=color, darkmode=True, title=title, **fig_kw)