diff options
Diffstat (limited to 'python/graphics.py')
| -rw-r--r-- | python/graphics.py | 74 |
1 files changed, 25 insertions, 49 deletions
diff --git a/python/graphics.py b/python/graphics.py index 75740081..9fa5a1ac 100644 --- a/python/graphics.py +++ b/python/graphics.py @@ -58,62 +58,38 @@ def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'): return newcmap -def plot_time_color_map(df, spread_shock, attr="pnl", path=".", color_map=cm.RdYlGn, index='IG', centered = True): - - val_date = df.index[0].date() - df = df.reset_index() - df['days'] = (df['date'].dt.date - val_date).dt.days - ascending = [True,True] if index == 'HY' else [True,False] - df.sort_values(by=['date','spread'], ascending = ascending, inplace = True) - date_range = df.days.unique() - - #plt.style.use('seaborn-whitegrid') - fig, ax = plt.subplots() - series = df[attr] - if centered is True: - midpoint = 1 - series.max() / (series.max() + abs(series.min())) - shifted_cmap = shiftedColorMap(color_map, midpoint=midpoint, name='shifted') - else: - shifted_cmap = color_map - - chart = ax.imshow(series.values.reshape(date_range.size, spread_shock.size).T, - extent=(date_range.min(), date_range.max(), - spread_shock.min(), spread_shock.max()), - aspect='auto', interpolation='bilinear', cmap=shifted_cmap) - - #chart = ax.contour(date_range, spread_shock, series.values.reshape(date_range.size, spread_shock.size).T) +def plot_color_map(series, sort_order = [True,True], color_map=cm.RdYlGn, centered = True): + ''' + 2D heat map - if x-axis is time translate to days instead - ax.set_xlabel('Days') - ax.set_ylabel('Price') if index == 'HY' else ax.set_ylabel('Spread') - ax.set_title('{} of Trade'.format(attr.title())) + Parameters + ----- + series: Series with multilevel index (x: first level, y: second level) + sort_order: sorting in the x,y axis + color_map: default Red-Yellow-Green + centered: center yellow as 0 of the series. + ''' - fig.colorbar(chart, shrink=.8) - #fig.savefig(os.path.join(path, "spread_time_color_map_"+ attr+ "_{}.png".format(val_date))) + x = series.index.get_level_values(0) + y = series.index.get_level_values(1) + if x.dtype == '<M8[ns]': + x = (x - x[0]).days + x.name = 'Days' + series.sort_index(ascending = sort_order, inplace = True) -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) + midpoint = 1 - series.max() / (series.max() \ + + abs(series.min())) if centered is True else 0.5 + color_map = shiftedColorMap(color_map, midpoint=midpoint) - 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)) + chart = ax.imshow(series.values.reshape(x.unique().size, y.unique().size).T, + extent=(x.min(), x.max(), y.min(), y.max()), + aspect='auto', interpolation='bilinear', cmap=color_map) + ax.set_xlabel(x.name) + ax.set_ylabel(y.name) + ax.set_title('{} of Trade'.format(series.name)) fig.colorbar(chart, shrink=.8) - #fig.savefig(os.path.join(path, "vol_spread_color_map"+ attr+ "_{}.png".format(val_date))) def plot_prob_map(df, attr="pnl", path=".", color_map=cm.RdYlGn, index='IG'): |
