yield_curve.py

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D
from datetime import datetime
from datetime import timedelta
from scipy.interpolate import CubicSpline

df = pd.read_csv('yields.csv')

x = []
y = []
z = []

header = []
for name in df.columns[1:]:
    maturity = float(name.split(" ")[0])
    if name.split(" ")[1] == 'Mo':
        maturity = maturity / 12
    header.append(maturity)

xs = np.arange(0, len(header))
for idx, row in df.iterrows():
	# x.append([(datetime.strptime(row['Date'], '%m/%d/%Y') - datetime(2019, 1, 1)).days for x in range(0, 1200)])
	x.append([(datetime.strptime(row['Date'], '%m/%d/%Y') - datetime(2019, 1, 1)).days for x in range(0, len(xs))])

	y.append(header)
	z.append(CubicSpline(np.arange(1, 13), row[1:])(xs))

x = np.array(x)
y = np.array(y)
z = np.array(z)


fig = plt.figure()
ax = fig.gca(projection='3d')

print(x.shape, y.shape, z.shape)
surf = ax.plot_surface(x, y, z, cmap='viridis',
                       linewidth=0, antialiased=True)
ax.set_xlabel('Date')
ax.set_ylabel('Time to maturity (months)')
ax.set_zlabel('Yield')

def format_date(x, pos=None):
     return (datetime(2019, 1, 1) + timedelta(days = x)).strftime("%Y-%m-%d")
ax.w_xaxis.set_major_formatter(ticker.FuncFormatter(format_date))

fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()