Nugroho's blog.: 3D (Polar/Cylindrical Coordinate) Animation of 2D Diffusion Equation using Python, Scipy, and Matplotlib

Friday, November 27, 2015

3D (Polar/Cylindrical Coordinate) Animation of 2D Diffusion Equation using Python, Scipy, and Matplotlib

Yup, that same code but in polar coordinate.

I use nabla operator for cylindrical coordinate but ditch the z component.

So, what's the z-axis for? It's represent the u value, in this case, temperature, as function of r and phi (I know I should use rho, but, ...)

`import scipy as spfrom mpl_toolkits.mplot3d import Axes3Dfrom matplotlib import cmfrom matplotlib.ticker import LinearLocator, FormatStrFormatterimport matplotlib.pyplot as pltimport mpl_toolkits.mplot3d.axes3d as p3import matplotlib.animation as animation#dr  = .1#dp  = .1#nr      = int(1/dr)#np      = int(2*sp.pi/dp)nr  = 10np  = 10r       = sp.linspace(0.,1.,nr)p       = sp.linspace(0.,2*sp.pi,np)dr  = r[1]-r[0]dp  = p[1]-p[0]a   = .5tmax    = 100t       = 0.dr2     = dr**2dp2     = dp**2dt      = dr2 * dp2 / (2 * a * (dr2 + dp2) )dt      /=10.print 'dr = ',dr print 'dp = ',dp print 'dt = ',dtut      = sp.zeros([nr,np])u0      = sp.zeros([nr,np])ur      = sp.zeros([nr,np])ur2     = sp.zeros([nr,np])#initialfor i in range(nr):    for j in range(np):        if ((i>4)&(i<6)):            u0[i,j] = 1.#print u0def hitung_ut(ut,u0):    for i in sp.arange (len(r)):        if r[i]!= 0.:            ur[i,:]     = u0[i,:]/r[i]            ur2[i,:]     = u0[i,:]/(r[i]**2)    ut[1:-1, 1:-1]  = u0[1:-1, 1:-1] + a*dt*(            (ur[1:-1, 1:-1] - ur[:-2, 1:-1])/dr+            (u0[2:, 1:-1] - 2*u0[1:-1, 1:-1] + u0[:-2,1:-1])/dr2+            (ur2[1:-1, 2:] - 2*ur2[1:-1, 1:-1] + ur2[1:-1, :-2])/dp2)    #calculate the edge    ut[1:-1, 0]  = u0[1:-1, 0] + a*dt*(            (ur[1:-1, 0] - ur[:-2, 0])/dr+            (u0[2:, 0] - 2*u0[1:-1, 0] + u0[:-2, 0])/dr2+            (ur2[1:-1, 1] - 2*ur2[1:-1, 0] + ur2[1:-1, np-1])/dp2)    ut[1:-1, np-1]  = u0[1:-1, np-1] + a*dt*(            (ur[1:-1, np-1] - ur[:-2, np-1])/dr+            (u0[2:, np-1] - 2*u0[1:-1, np-1] + u0[:-2,np-1])/dr2+            (ur2[1:-1, 0] - 2*ur2[1:-1, np-1] + ur2[1:-1, np-2])/dp2)#hitung_ut(ut,u0)#print utdef data_gen(framenumber, Z ,surf):    global ut    global u0    global t    hitung_ut(ut,u0)    u0[:] = ut[:]    Z = u0    t += 1    print t        ax.clear()    plotset()    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,                       linewidth=0, antialiased=False, alpha=0.7)    return surf,fig = plt.figure()#ax = fig.gca(projection='3d')ax = fig.add_subplot(111, projection='3d')P,R = sp.meshgrid(p,r)X,Y = R*sp.cos(P),R*sp.sin(P) Z = u0print len(R), len(P)def plotset():    ax.set_xlim3d(-1., 1.)    ax.set_ylim3d(-1., 1.)    ax.set_zlim3d(-1.,1.)    ax.set_autoscalez_on(False)    ax.zaxis.set_major_locator(LinearLocator(10))    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))    cset = ax.contour(X, Y, Z, zdir='x', offset=-1. , cmap=cm.coolwarm)    cset = ax.contour(X, Y, Z, zdir='y', offset=1. , cmap=cm.coolwarm)    cset = ax.contour(X, Y, Z, zdir='z', offset=-1., cmap=cm.coolwarm)plotset()surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,                       linewidth=0, antialiased=False, alpha=0.7)fig.colorbar(surf, shrink=0.5, aspect=5)ani = animation.FuncAnimation(fig, data_gen, fargs=(Z, surf),frames=4096, interval=4, blit=False)#ani.save('2dDiffusionfRadialf1024b512.mp4', bitrate=1024)plt.show()    `
.

 100x100 size