1 files changed, 12 insertions, 12 deletions
diff --git a/python/exploration/swaption_calendar_spread.py b/python/exploration/swaption_calendar_spread.py
index fb7d4b6a..a007bbb9 100644
--- a/python/exploration/swaption_calendar_spread.py
+++ b/python/exploration/swaption_calendar_spread.py
@@ -13,6 +13,7 @@ from mpl_toolkits.mplot3d import Axes3D
 import matplotlib.pyplot as plt
 from copy import deepcopy
 
+import os
 import numpy as np
 import matplotlib
 import matplotlib.pyplot as plt
@@ -81,29 +82,28 @@ def plot_df(df, spread_shock, vol_shock):
     ax.set_ylabel("spread")
     ax.set_zlabel("PV")
 
-def plot_color_map(df):
+def plot_color_map(df, spread_shock, vol_shock, attr="pv", path="."):
 
-    val_date = df.index[0]
-    df = df.set_index(['spread_shock','vol_shock'])['pv'].unstack(-2)
-    df = df.sort_index(ascending = False)
+    val_date = df.index[0].date()
 
     #rows are spread, columns are volatility surface shift
     fig, ax = plt.subplots()
-
+    series = df[attr]
     #Different ways to do a colormap: imshow and pcolormesh. using imshow here
-    midpoint = 1 - df.max().max()/(df.max().max() + abs(df.min().min()))
+    midpoint = 1 - series.max() / (series.max() + abs(series.min()))
     shifted_cmap = shiftedColorMap(cm.RdYlGn, midpoint=midpoint, name='shifted')
 
-    chart = ax.imshow(df, extent=(df.columns.min(), df.columns.max(), df.index.min(), df.index.max()) \
-                            ,aspect= 'auto', interpolation='bilinear', cmap=shifted_cmap)
+    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('Spread')
     ax.set_ylabel('Volatility shock')
-    ax.set_title('PV of Trade on ' + str(val_date.date()))
-
-    fig.colorbar(chart, shrink = .8)
+    ax.set_title('{} of Trade on {}'.format(attr.title(), val_date))
 
-    fig.savefig("/home/serenitas/edwin/PythonGraphs/payer_swap_" + str(val_date.date()) + ".png")
+    fig.colorbar(chart, shrink=.8)
+    fig.savefig(os.path.join(path, "payer_swap_{}.png".format(val_date)))
 
 def calc_delta_pnl(index, date_range, spread_shock):