Surface plot in python

Question

i have googled this question and found: surface plots in matplotlib but got incredibly confused applying it to my problem. I looked on matplotlib website too but it was tough to figure out the arguments.

I currently have a matrix U, where each row i is the value of a put option due to a certain stock price at time i . Eg U[0,0] is the value of an option at time 0 when the stock price is 0. U[1,0] is the value of an option at time 1 when the stock price is 0. U[1,1] is the value of an option at time 1 when the stock price is the first nodal point.

I am wanting to have a 3d plot of this, ie showing the evolution of the option value against stock price through time.

X axis: Stock price Y axis: Option value Z axis: Time.

Could anyone help me please? Thanks, James

Answer 1

First, you'll need to make your Z and X coordinates, which are the indices of your matrix. You can get the size of the matrix with matrix.shape , then use meshgrid to create a grid of coordinates. Then you can get the values from the matrix corresponding to each set of (z,x) coordinates and put them into Y.

For example:

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure()
ax = fig.gca(projection='3d')
matrix = np.array([[0, 1, 2, 3, 4], [.5, 1.5, 2.5, 3.5, 4.5], [1, 2, 3, 4, 5]])
z, x = matrix.shape
Z = np.arange(0, z, 1)
X = np.arange(0, x, 1)
Z, X = np.meshgrid(Z, X)
Y = matrix[Z, X]
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, 
        linewidth=0, antialiased=False, cmap = cm.coolwarm)

plt.show()