from pylab import *
from scipy.special import *
from scipy.optimize import *
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
%config InlineBackend.figure_format = 'retina'

x,m_max = linspace(0,15),5
for m in range(m_max):
    y = jv(m,x)
    plot(x,y,label="$J_%d(x)$" %m)
legend(),grid()
show()

n_max=6
for m in range(m_max):
    hlines(jn_zeros(m,n_max),m-0.5,m+0.5,colors='green')
    hlines(jn_zeros(m,n_max),m_max+0.5,m_max+2,colors='blue')
ylim([0,20]),xlabel("$m$"),ylabel("Raíz de $j_{m\,n}$")
show()

raices = np.zeros((m_max,n_max))
for m in range(m_max):
    raices[m,:] = jn_zeros(m,n_max)
    
def modo_circ(rho,theta, m):
    return jv(m,rho)*cos(m*theta)

def modo_circ_fig(X,Y,Z,m_n):
    # comandos para graficar
    fig = figure(figsize=(10, 7))
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=plt.cm.coolwarm,linewidth=0, antialiased=False)
    ax.set_zlim(-1., 1.),ax.set_xlabel('$x$'),ax.set_ylabel('$y$')
    ax.set_zlabel(r'$\Psi_{%d %d}(\rho,\theta) $' %(m_n[0],m_n[1]) )
    ax.set_title('Modo m= %d , n= %d ' %(m_n[0],m_n[1]))
    fig.colorbar(surf, shrink=0.5, aspect=7, cmap=plt.cm.coolwarm)
    ax.view_init(30,40)
    show()

rho = linspace(0,1,150)
theta = linspace(0,2*pi,300)
RHO,T = meshgrid(rho,theta)
X,Y = RHO*cos(T),RHO*sin(T)

psi = modo_circ(RHO*raices[0,0],T,0)
modo_circ_fig(X,Y,psi,[0,1])

psi = modo_circ(RHO*raices[0,1],T,0)
modo_circ_fig(X,Y,psi,[0,2])

psi = modo_circ(RHO*raices[0,2],T,0)
modo_circ_fig(X,Y,psi,[0,3])

psi = modo_circ(RHO*raices[1,0],T,1)
modo_circ_fig(X,Y,psi,[1,1])

psi = modo_circ(RHO*raices[2,0],T,2)
modo_circ_fig(X,Y,psi,[2,1])

fig = figure(figsize=(15, 18))
num_fig = 1
for m in range(4):
    for n in range(3):
        psi = modo_circ(RHO*raices[m,n],T,m)
        ax1 = fig.add_subplot(4,3,num_fig,aspect='equal',xlabel='$x$',ylabel='$y$',
                      title='Modo m= %d , n= %d ' %(m,n+1))
        psi_g = ax1.contourf(X,Y,psi,100,cmap=plt.cm.coolwarm)
        psi_cont = ax1.contour(X,Y,psi,levels=0,linewidths=2,colors='k')
        num_fig +=1

#fig.colorbar(psi00, shrink=0.5, aspect=7, cmap=plt.cm.coolwarm)


show()