diff options
Diffstat (limited to 'python/graphics.py')
| -rw-r--r-- | python/graphics.py | 114 |
1 files changed, 66 insertions, 48 deletions
diff --git a/python/graphics.py b/python/graphics.py index 9fa5a1ac..82ac113a 100644 --- a/python/graphics.py +++ b/python/graphics.py @@ -3,12 +3,13 @@ import matplotlib.pyplot as plt from matplotlib import cm from matplotlib.colors import LinearSegmentedColormap -from scipy.interpolate import griddata +from scipy.interpolate import griddata, SmoothBivariateSpline from scipy.stats import norm from scipy.optimize import curve_fit -def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'): - ''' + +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 @@ -28,38 +29,36 @@ def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'): 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': [] - } + """ + 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) - ]) + 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)) + 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_color_map(series, sort_order = [True,True], color_map=cm.RdYlGn, centered = True): - ''' + +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 Parameters @@ -68,68 +67,87 @@ def plot_color_map(series, sort_order = [True,True], color_map=cm.RdYlGn, center sort_order: sorting in the x,y axis color_map: default Red-Yellow-Green centered: center yellow as 0 of the series. - ''' + """ x = series.index.get_level_values(0) y = series.index.get_level_values(1) - if x.dtype == '<M8[ns]': + if x.dtype == "<M8[ns]": x = (x - x[0]).days - x.name = 'Days' - series.sort_index(ascending = sort_order, inplace = True) + x.name = "Days" + series.sort_index(ascending=sort_order, inplace=True) fig, ax = plt.subplots() - midpoint = 1 - series.max() / (series.max() \ - + abs(series.min())) if centered is True else 0.5 + midpoint = ( + 1 - series.max() / (series.max() + abs(series.min())) + if centered is True + else 0.5 + ) color_map = shiftedColorMap(color_map, midpoint=midpoint) - 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) + 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) + ax.set_title("{} of Trade".format(series.name)) + fig.colorbar(chart, shrink=0.8) -def plot_prob_map(df, attr="pnl", path=".", color_map=cm.RdYlGn, index='IG'): + +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 + 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) + prob_defined = np.linspace(0.001, 0.999, 1000) midpoint = 1 - series.max() / (series.max() + abs(series.min())) - shifted_cmap = shiftedColorMap(color_map, midpoint=midpoint, name='shifted') + shifted_cmap = shiftedColorMap(color_map, midpoint=midpoint, name="shifted") - resampled = griddata((df.days, df.prob), series, (days_defined[None, :], - prob_defined[:, None]), method='linear') + resampled = griddata( + (df.days, df.prob), + series, + (days_defined[None, :], prob_defined[:, None]), + method="linear", + ) - #plot + # 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) + 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())) - ax.set_xlabel('Days') - ax.set_ylabel('Probability') - ax.set_title('{} of Trade'.format(attr.title())) + fig.colorbar(chart, shrink=0.8) - 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') + 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) + 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))
\ No newline at end of file + ax.set_title("{} of Trade on {}".format(attr.title(), val_date)) |