diff options
Diffstat (limited to 'python/graphics.py')
| -rw-r--r-- | python/graphics.py | 28 |
1 files changed, 27 insertions, 1 deletions
diff --git a/python/graphics.py b/python/graphics.py index 0f348a04..3a88147a 100644 --- a/python/graphics.py +++ b/python/graphics.py @@ -1,6 +1,7 @@ import numpy as np import matplotlib.pyplot as plt from matplotlib import cm +from matplotlib.colors import LinearSegmentedColormap def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'): ''' @@ -48,7 +49,7 @@ def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'): cdict['blue'].append((si, b, b)) cdict['alpha'].append((si, a, a)) - newcmap = matplotlib.colors.LinearSegmentedColormap(name, cdict) + newcmap = LinearSegmentedColormap(name, cdict) plt.register_cmap(cmap=newcmap) return newcmap @@ -82,3 +83,28 @@ def plot_time_color_map(df, spread_shock, attr="pnl", path=".", color_map=cm.RdY fig.colorbar(chart, shrink=.8) #fig.savefig(os.path.join(path, "spread_time_color_map_"+ attr+ "_{}.png".format(val_date))) + +def plot_color_map(df, spread_shock, vol_shock, attr="pnl", path=".", index='IG'): + # TODO: merge with plot_time_color_map + val_date = df.index[0].date() + #rows are spread, columns are volatility surface shift + fig, ax = plt.subplots() + #We are plotting an image, so we have to sort from high to low on the Y axis + ascending = [False,False] if index == 'HY' else [True,False] + df.sort_values(by=['spread','vol_shock'], ascending=ascending, inplace=True) + series = df[attr] + + midpoint = 1 - series.max() / (series.max() + abs(series.min())) + shifted_cmap = shiftedColorMap(cm.RdYlGn, midpoint=midpoint, name='shifted') + + chart = ax.imshow(series.values.reshape(spread_shock.size, vol_shock.size).T, + extent=(spread_shock.min(), spread_shock.max(), + vol_shock.min(), vol_shock.max()), + aspect='auto', interpolation='bilinear', cmap=shifted_cmap) + + ax.set_xlabel('Price') if index == 'HY' else ax.set_xlabel('Spread') + ax.set_ylabel('Volatility shock') + ax.set_title('{} of Trade on {}'.format(attr.title(), val_date)) + + fig.colorbar(chart, shrink=.8) + #fig.savefig(os.path.join(path, "vol_spread_color_map"+ attr+ "_{}.png".format(val_date))) |