import numpy as np import matplotlib.pyplot as plt from matplotlib import cm from matplotlib.colors import LinearSegmentedColormap from scipy.interpolate import griddata from scipy.stats import norm from scipy.optimize import curve_fit def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'): ''' Function to offset the "center" of a colormap. Useful for data with a negative min and positive max and you want the middle of the colormap's dynamic range to be at zero Input ----- cmap : The matplotlib colormap to be altered start : Offset from lowest point in the colormap's range. Defaults to 0.0 (no lower ofset). Should be between 0.0 and `midpoint`. midpoint : The new center of the colormap. Defaults to 0.5 (no shift). Should be between 0.0 and 1.0. In general, this should be 1 - vmax/(vmax + abs(vmin)) For example if your data range from -15.0 to +5.0 and you want the center of the colormap at 0.0, `midpoint` should be set to 1 - 5/(5 + 15)) or 0.75 stop : Offset from highets point in the colormap's range. Defaults to 1.0 (no upper ofset). Should be between `midpoint` and 1.0. ''' cdict = { 'red': [], 'green': [], 'blue': [], 'alpha': [] } # regular index to compute the colors reg_index = np.linspace(start, stop, 257) # shifted index to match the data shift_index = np.hstack([ np.linspace(0.0, midpoint, 128, endpoint=False), np.linspace(midpoint, 1.0, 129, endpoint=True) ]) for ri, si in zip(reg_index, shift_index): r, g, b, a = cmap(ri) cdict['red'].append((si, r, r)) cdict['green'].append((si, g, g)) cdict['blue'].append((si, b, b)) cdict['alpha'].append((si, a, a)) newcmap = LinearSegmentedColormap(name, cdict) plt.register_cmap(cmap=newcmap) 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) ax.set_xlabel('Days') ax.set_ylabel('Price') if index == 'HY' else ax.set_ylabel('Spread') ax.set_title('{} of Trade'.format(attr.title())) 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))) def plot_prob_map(df, attr="pnl", path=".", color_map=cm.RdYlGn, index='IG'): val_date = df.index[0].date() df = df.reset_index() df['days'] = (df['date'] - val_date).dt.days series = df[attr] days_defined = np.linspace(df.days.min(), df.days.max(), 1000) prob_defined = np.linspace(0.001, .999, 1000) midpoint = 1 - series.max() / (series.max() + abs(series.min())) shifted_cmap = shiftedColorMap(color_map, midpoint=midpoint, name='shifted') resampled = griddata((df.days, df.prob), series, (days_defined[None, :], prob_defined[:, None]), method='linear') #plot fig, ax = plt.subplots() chart = ax.imshow(resampled.reshape(days_defined.size, prob_defined.size), extent=(df.days.min(), df.days.max(), 0, 1), aspect='auto', interpolation='bilinear', cmap=shifted_cmap) ax.set_xlabel('Days') ax.set_ylabel('Probability') ax.set_title('{} of Trade'.format(attr.title())) fig.colorbar(chart, shrink=.8) def plot_swaption_df(df, spread_shock, vol_shock, attr="pnl"): val_date = df.index[0].date() fig = plt.figure() ax = fig.gca(projection='3d') ## use smoothing spline on a finer grid series = df[attr] f = SmoothBivariateSpline(df.vol_shock.values, df.spread_shock.values, series.values) xx, yy = np.meshgrid(vol_shock, spread_shock) surf = ax.plot_surface(xx, yy, f(vol_shock, spread_shock).T, cmap=cm.viridis) ax.set_xlabel("Volatility shock") ax.set_ylabel("Spread") ax.set_zlabel("PnL") ax.set_title('{} of Trade on {}'.format(attr.title(), val_date))