1 files changed, 84 insertions, 0 deletions
diff --git a/python/graphics.py b/python/graphics.py
new file mode 100644
index 00000000..0f348a04
--- /dev/null
+++ b/python/graphics.py
@@ -0,0 +1,84 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib import cm
+
+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 = matplotlib.colors.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'):
+
+    val_date = df.index[0].date()
+    df = df.reset_index()
+    df['days'] = (df['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]
+    midpoint = 1 - series.max() / (series.max() + abs(series.min()))
+    shifted_cmap = shiftedColorMap(color_map, midpoint=midpoint, name='shifted')
+
+    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)))