how do you radially 'sweep out' a 1D array to plot 3d figure in python? (to represent a wavefunction)

Question

effectively I have a large 1D array of heights. As a small example consider:

u=array([0,1,2,1,0,2,4,6,4,2,1])

and a 1D array, the same size as

Answer 1

#!/usr/bin/python

from mpl_toolkits.mplot3d import Axes3D
import matplotlib
import numpy as np
from scipy.interpolate import interp1d
from matplotlib import cm
from matplotlib import pyplot as plt
step = 0.04
maxval = 1.0
fig = plt.figure()
ax = Axes3D(fig)  



u=np.array([0,1,2,1,0,2,4,6,4,2,1])
r=np.array([0,1,2,3,4,5,6,7,8,9,10])
f=interp1d(r,u)

# walk along the circle
p = np.linspace(0,2*np.pi,50)
R,P = np.meshgrid(r,p)
# transform them to cartesian system
X,Y = R*np.cos(P),R*np.sin(P)

Z=f(R)

ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.jet)
ax.set_xticks([])
plt.show()