3d plotting: cylindrical coordinate function

Question

I have a to plot a function that is in polar coordinates.

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

pi=np.pi
r=np.linspace(0,4,30)
phi=np.linspace(0,2*pi,10)

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.

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.

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)

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

I'm trying something like :

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)

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).

Any ideas??

It works fine if I do :

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')

But I want a surface or wire plot, not scatter.

Below is the full 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')

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.

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'))

xnx xnx · Accepted Answer · 2016-01-18T18:50:09

Things seem a little muddled in your code, but unravelling a few things, the following code made a plot for me. Your A0 function is constant in phi, so you can assign all values for a given r with A0[:,j]. I think you had the phi and r indexes the wrong way round, which might have caused your NumPy broadcasting error.

import numpy as np
from scipy.special import binom as binom 
from mpl_toolkits.mplot3d import Axes3D 
from matplotlib import pyplot
plt = pyplot
from matplotlib import cm
import matplotlib

'''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(phi),len(r)))
    for j, rn in enumerate(r):
            A0[:,j]=np.exp(-rn**2)*Ln(p,0,2*rn**2)

    A1=np.zeros_like(A0)
    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_like(A0)

    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,40)

A0, A1, A2=A(5,1,r,phi)#

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)
ax.plot_surface(X, Y, A0)

plt.show()

3d plotting: cylindrical coordinate function

1 Answers