Skip to content

Commit

Permalink
Update plot.py
Browse files Browse the repository at this point in the history 
Refactor and improvements
  • Loading branch information
@quantgirluk
quantgirluk committed Jan 11, 2024
1 parent fa457ef commit 4d984bb
Showing 1 changed file with 55 additions and 54 deletions.
109 changes: 55 additions & 54 deletions fanchart/plot.py
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,31 @@
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
from pandas.plotting import register_matplotlib_converters
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():
"""
Expand Down Expand Up @@ -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",
Expand All @@ -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')
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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:
Expand All @@ -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)
Expand All @@ -230,18 +228,22 @@ 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())
ax1.spines[['top', 'right', 'left']].set_visible(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
Expand All @@ -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
Expand All @@ -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)

0 comments on commit 4d984bb

Please sign in to comment.