3d plotting: cylindrical coordinate function
<p>I have a to plot a function that is in polar coordinates.</p>

<p>first I set the values for the radial (r), and polar (phi) variables.</p>

<pre><code>pi=np.pi
r=np.linspace(0,4,30)
phi=np.linspace(0,2*pi,10)
</code></pre>

<p>Then I do A0(r,phi), where A0 is a function I wrote, that has a real value for each value of r and phi.</p>

<p>The way I store it is A0[r, phi]. witch means that the 'row' indicates the radial coordinate and the 'column' values the polar coordinate.</p>

<p>Ao, if I use n values in r, and m values in phi, A0 is a (n,m) Matrix , or (n,m) shaped numpy array. (in my example, it is a (30,10) numpy array)</p>

<p>'A0' seems to work as it should, but the problem is that I don't know how to plot it. </p>

<p>I'm trying something like :</p>

<pre><code>fig=pyplot.figure()
ax = pyplot.subplot(111, projection='3d')

R, P = np.meshgrid(r, phi)
X, Y = R*np.cos(P), R*np.sin(P)

ax.plot_wireframe(X,Y, A0, rstride=10, cstride=10)
</code></pre>

<p>Trying to follow the matplotlib examples, but it says that the shapes don't agree, because it has a problem with the dimensions(between X, Y and A0).</p>

<p>Any ideas??</p>

<p>It works fine if I do :</p>

<pre><code>fig = pyplot.figure(figsize=(11,7), dpi=100)
ax = pyplot.subplot(111, projection='3d')

for j, rn in enumerate(r):
    for k, phin in enumerate(phi):
        ax.scatter(rn*np.cos(phin), rn*np.sin(phin), A0[j,k],',b')
</code></pre>

<p>But I want a surface or wire plot, not scatter.</p>

<p>Below is the full code: </p>

<pre><code>import numpy as np
from scipy.special import binom as binom 
from mpl_toolkits.mplot3d import Axes3D 
from matplotlib import pyplot
from matplotlib import cm
import matplotlib
%matplotlib inline

'''Define A_(pm)  (Gauss-laguerre polinomials) '''

def Ln(p,m,xj):
    Ln=0.
    for i in range(p+1):#el p+1 es por la indexacion
        Ln=Ln+((-1)**i)*binom(p+m,p-i)*(xj**i)/np.math.factorial(i)
    return Ln

def A(p, m, r, phi):#defino G-L

    A0=np.zeros((len(r),len(phi)))
    for j, rn in enumerate(r):
        A0[j]=np.exp(-rn**2)*Ln(p,0,2*rn**2)

    A1=np.zeros((len(r),len(phi)))
    for j, rn in enumerate(r):
         for k, phin in enumerate(phi):
            A1[k,j]=2*((2*rn**2)**(m/2))*np.sqrt(np.math.factorial(p)/np.math.factorial(p+m))*np.exp(-rn**2)*Ln(p,m,2*(rn**2))*np.sin(m*phin)

    A2=np.zeros((len(r),len(phi)))

    for j, rn in enumerate(r):
        for k, phin in enumerate(phi):
            A2[k,j]=2*(2*rn**2)**(m/2)*np.sqrt(np.math.factorial(p)/np.math.factorial(p+m))*np.exp(-rn**2)*Ln(p,m,2*rn**2)*np.cos(m*phin)


    return A0, A1, A2

   pi=np.pi
   r=np.linspace(0,6,50)
   phi=np.linspace(0,2*pi,50)

   A0, A1, A2=A(5,1,r,phi)#i think a have a bugs with the values of m .

   fig = pyplot.figure(figsize=(11,7), dpi=100)
   ax = pyplot.subplot(111, projection='3d')

for j, rn in enumerate(r):
    for k, phin in enumerate(phi):
        #ax.scatter(rn*np.cos(phin), rn*np.sin(phin), A0[j,k],'.b')
        ax.scatter(rn*np.cos(phin), rn*np.sin(phin), A1[j,k],'.r')
        #ax.scatter(rn*np.cos(phin), rn*np.sin(phin), A2[j,k],'.r')
</code></pre>

<p>This works, but a scatter is not what I want.
Here is what I tried to do, but it doesn't work if r and phi have different dimensions.</p>

<pre><code>fig = pyplot.figure(figsize=(11,9), dpi=100)
ax = pyplot.subplot(111, projection='3d')
R, P = np.meshgrid(r, phi)
X, Y = R*np.cos(P), R*np.sin(P)
print 'shape r=', np.shape(r)
print 'shape R=', np.shape(R)
print 'shape P=', np.shape(P)
A0, A1, A2=A(5,0,R[1],R[0])
ax.plot_surface(X, Y, np.transpose(A0), rstride=3, cstride=3,   cmap=pyplot.get_cmap('viridis'))
</code></pre>