Changeset 768 for codes/icosagcm

Timestamp:

10/31/18 14:22:54 (6 years ago)

Author:

jisesh

Message:

devel/Python : Parallel version of NH_3D_bubble with XIOS

File:

• codes/icosagcm/devel/Python/test/py/NH_3D_bubble_parallel.py (modified) (8 diffs)

codes/icosagcm/devel/Python/test/py/NH_3D_bubble_parallel.py

@@ -2 +2 @@
 from dynamico import dyn
 from dynamico import time_step
-from dynamico import DCMIP
 from dynamico import meshes
 from dynamico import xios
 from dynamico import precision as prec
-#from dynamico.meshes import Cartesian_mesh as Mesh
 import math as math
 import matplotlib.pyplot as plt
@@ -12 +10 @@
 import time
 import argparse
-
-parser = argparse.ArgumentParser()
-
-parser.add_argument("-mpi_ni", type=int,
-                    default=64, choices=None,
-                    help="number of x processors")
-parser.add_argument("-mpi_nj", type=int,
-                    default=64, choices=None,
-                    help="number of y processors")
-args = parser.parse_args()
-
 
 def thermal_bubble_3D(Lx,nx,Ly,ny,llm,ztop=1000., zc=350.,
@@ -37 +24 @@
     temp  =  lambda p: theta0*(p/p0)**(Rd/Cpd)
     T     = lambda x,y,p: deform(x,y,p) + temp(p)
-    phi0  = 45.       # Reference latitude North pi/4 (deg)
-    omega = 7.292e-5  # Angular velocity of the Earth (s^-1)
-    f0    = 2*omega*np.sin(phi0)
-    lap   = 0.005      # Lapse rate (K m^-1)
-    Rd    = 287.0      # Gas constant for dryy air (j kg^-1 K^-1)
-
-    def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels
-    #def eta(alpha) : return alpha/float(llm)
-
-    filename = 'cart_%03d_%03d.nc'%(nx,ny)
-    print 'Reading Cartesian mesh ...'
-    def coriolis(lon,lat): return f0+0.*lon
-    meshfile = meshes.DYNAMICO_Format(filename)
-    radius = None
-    #mesh = Mesh(nx,ny,llm,nqdyn,Lx,Ly,0.)
-    print('----read--------')
-    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
-    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
-    mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)
-
-    print(mesh.__dict__.keys())
 
     alpha_k = (np.arange(llm)  +.5)/llm
@@ -68 +34 @@
     y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij')
 
-    print('alpha_l=',alpha_l)
-
     thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0)
 
-    eta_il = eta(alpha_il)
-    eta_ik = eta(alpha_ik)
-    eta_ek = eta(alpha_ek)
-
-    print('eta_il=',eta_il)
-    #Phi_il = Phi(mesh.llp1/float(llm)) #llp is not defined in local mesh
-    #Phi_ik = Phi((mesh.ll+.5)/llm)
-    Phi_il = Phi(eta_il)
-    Phi_ik = Phi(eta_ik)
+    Phi_il = Phi(alpha_il)
+    Phi_ik = Phi(alpha_ik)
     p_ik = p(Phi_ik)
     T_ik = T(x_ik, y_ik, p_ik)
@@ -98 +55 @@
     params.dx_g0=dx/g
     params.g = g
-    #pbot = p(Phi_il[:,:,0])
-    #gas_bot = thermo.set_pT(pbot, temp(pbot))
+
     # define parameters for lower BC
-    pbot = p(eta_il[0])
+    pbot = p(alpha_il[0])
     print 'min p, T :', pbot.min(), temp(pbot/p0)
     gas_bot = thermo.set_pT(pbot, temp(pbot/p0))
@@ -109 +65 @@
     return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas
 
+def diagnose(Phi,S,m,W):
+    s=S/m ; s=.5*(s+abs(s))
+    for l in range(llm):
+        v[:,l]=(Phi[:,l+1]-Phi[:,l])/(g*m[:,l])
+        w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l]
+        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
+    gas = thermo.set_vs(v,s)
+    return gas, w, z
+
+def reshape(data): return data.reshape((nx,ny))
+
+def plot():
+    x, y = map(reshape, (xx,yy) )
+    zz=np.zeros((nx,ny,llm))
+    ss=np.zeros((nx,ny,llm))
+    ww=np.zeros((nx,ny,llm))
+    x3d=np.zeros((nx,ny,llm))
+     
+    for l in range(llm):
+        zz[:,:,l],ss[:,:,l],ww[:,:,l] = map(reshape, (z[:,l],gas.s[:,l],w[:,l]) )
+        x3d[:,:,l]=x[:,:]
+
+    jj=ny/2
+    xp=x3d[:,jj,:]
+    zp=zz[jj,:,:]/1000.
+    sp=ss[jj,:,:]
+    wp=ww[jj,:,:]
+
+    f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4))
+
+    c=ax1.contourf(xp,zp,sp,20)
+    ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)')        
+    plt.colorbar(c,ax=ax1)
+    ax1.set_title(title_format % (it*T,))
+
+    c=ax2.contourf(xp,zp,wp,20)
+    ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)')        
+    plt.colorbar(c,ax=ax2)
+    ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,))
+    plt.savefig('fig_NH_3D_bubble/%02d.png'%it)    
+
+#------------------------- main program --------------------------
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--mpi_ni", type=int, default=64, 
+                    help="number of x processors")
+parser.add_argument("--mpi_nj", type=int, default=64,
+                    help="number of y processors")
+
+parser.add_argument("--python_stepping", type=bool, default=False,
+                    help="Time stepping in Python or Fortran")
+parser.add_argument("--dt", type=float, default=.25,
+                    help="Time step in seconds")
+parser.add_argument("--T", type=float, default=5.,
+                    help="Length of time slice in seconds")
+parser.add_argument("--Nslice", type=int, default=10,
+                    help="Number of time slices")
+
+parser.add_argument("--thetac", type=float, default=30.,
+                    help="Initial extra temperature of bubble")
+parser.add_argument("--Lx", type=float, default=2000.,
+                    help="Size of box in meters")
+parser.add_argument("--Ly", type=float, default=2000.,
+                    help="Size of box in meters")
+parser.add_argument("--nx", type=int, default=20,
+                    help="Resolution in the x direction")
+parser.add_argument("--ny", type=int, default=20,
+                    help="Resolution in the y direction")
+parser.add_argument("--llm", type=int, default=79,
+                    help="Number of vertical levels")
+
+args = parser.parse_args()
+
 with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator
     comm = client.comm
     mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
     print '%d/%d starting'%(mpi_rank,mpi_size)
-
-    #Lx, nx, llm, thetac, T, Nslice, courant = 2000., 100, 50, 30., 5., 10, 2.8
-    Lx, nx, llm, thetac = 2000., 200, 79, 30
-    #Lx, nx, llm, thetac, T, Nslice, courant = 3000., 75, 25, -30, 5., 10, 2.8
-
+    T, Nslice, dt, thetac = args.T, args.Nslice, args.dt, args.thetac
+    Lx, nx, Ly, ny, llm   = args.Lx, args.nx, args.Ly, args.ny, args.llm
+    
     nqdyn, dx = 1, Lx/nx
     Ly,ny,dy = Lx,nx,dx
@@ -124 +152 @@
     unst.setvar('g',g)
 
-    print('bubble call-----------')
+    filename = 'cart_%03d_%03d.nc'%(nx,ny)
+    print 'Reading Cartesian mesh ...'
+    def coriolis(lon,lat): return 0.*lon
+    meshfile = meshes.DYNAMICO_Format(filename)
+    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
+    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
+    mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, None, coriolis)
+
     thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac)
-    print('bubble done-----------')
-
 
     # compute hybrid coefs from initial distribution of mass
@@ -134 +167 @@
     unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk)
 
-    #dz = flow0[3].max()/(params.g*llm)
+    dz = flow0[3].max()/(params.g*llm)
+    # courant = 2.8
     #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2))
     #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2))
-    #nt = int(math.ceil(T/dt))
-    #dt = T/nt
-    #print 'Time step : %d x %g s' % (nt,dt)
-
-
-    T, nslice, dt = 3600., 1, 3600.
-    #T, nslice  = 3600., 4
-
-    with xios.Context_Curvilinear(mesh,1, dt) as context:
+    nt = int(math.ceil(T/dt))
+    dt = T/nt
+    print 'Time step : %d x %g = %g s' % (nt,dt,nt*dt)
+
+#    #caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass
+    caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
+#    #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag
+
+    if args.python_stepping: # time stepping in Python
+        caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g)
+        scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt)
+        def next_flow(m,S,u,Phi,W):
+            return scheme.advance((m,S,u,Phi,W),nt)
+
+    else: # time stepping in Fortran
+        scheme = time_step.ARK2(None, dt)
+        caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta,
+                                      thermo,params,params.g)
+        def next_flow(m,S,u,Phi,W):
+            # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
+            caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
+            caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W
+            caldyn_step.next()
+            return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(),
+                caldyn_step.geopot.copy(), caldyn_step.W.copy())
+    
+    m,S,u,Phi,W=flow0
+    if caldyn_thermo == unst.thermo_theta:
+        s=S/m
+        theta = thermo.T0*np.exp(s/thermo.Cpd)
+        S=m*theta
+        title_format = 'Potential temperature at t=%g s (K)'
+    else:
+        title_format = 'Specific entropy at t=%g s (J/K/kg)'
+    
+    w=mesh.field_mass()
+    z=mesh.field_mass()
+    xx,yy = mesh.lat_i, mesh.lon_i
+
+    # XIOS writes to disk every 24h
+    # each iteration lasts it*nt seconds but 
+    # we pretend that each iteration lasts 24h to make sure data is written to disk
+
+    with xios.Context_Curvilinear(mesh,1, 24*3600) as context:
         # now XIOS knows about the mesh and we can write to disk
-        for i in range(48): # 2 days
-            context.update_calendar(i)
-            
-            print 'send_field', i, gas0.T.shape
-    #    context.send_field_primal('ps', lat_i)
-            context.send_field_primal('temp', gas0.T)
-            context.send_field_primal('p', gas0.p)
-    print(gas0.__dict__.keys())
-    print('size of T, p',np.shape(gas0.T),np.shape(gas0.p))
-    print('************DONE************')
-
-#    #caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass
-#    caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
-#    #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag
-#
-#    if False: # time stepping in Python
-#        caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g)
-#        scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt)
-#        def next_flow(m,S,u,Phi,W):
-#            # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
-#            return scheme.advance((m,S,u,Phi,W),nt)
-#
-#    else: # time stepping in Fortran
-#        scheme = time_step.ARK2(None, dt)
-#        caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta,
-#                                      thermo,params,params.g)
-#        def next_flow(m,S,u,Phi,W):
-#            # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
-#            caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
-#            caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W
-#            caldyn_step.next()
-#            return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(),
-#                caldyn_step.geopot.copy(), caldyn_step.W.copy())
-#    
-#    m,S,u,Phi,W=flow0
-#    if caldyn_thermo == unst.thermo_theta:
-#        s=S/m
-#        theta = thermo.T0*np.exp(s/thermo.Cpd)
-#        S=m*theta
-#        title_format = 'Potential temperature at t=%g s (K)'
-#    else:
-#        title_format = 'Specific entropy at t=%g s (J/K/kg)'
-#    
-#    w=mesh.field_mass()
-#    z=mesh.field_mass()
-#    
-#    for it in range(Nslice):
-#        s=S/m ; s=.5*(s+abs(s))
-#        for l in range(llm):
-#            #w[:,:,l]=.5*params.g*(W[:,:,l+1]+W[:,:,l])/m[:,:,l]
-#            w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l]
-#            #z[:,:,l]=.5*(Phi[:,:,l+1]+Phi[:,:,l])/params.g
-#            z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
-#            
-#        print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g')
-#        jj=ny/2
-#        #xx,zz,ss,ww = mesh.xx[jj,:,:]/1000., z[jj,:,:]/1000., s[jj,:,:], w[jj,:,:]
-#        #xx,zz,ss,ww = x_ik[jj,:]/1000., z[jj,:]/1000., s[jj,:], w[jj,:]
-#    
-#        #f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4))
-#        
-#        #c=ax1.contourf(xx,zz,ss,20)
-#        #ax1.set_xlim((-.5,.5)), ax1.set_xlabel('x (km)')
-#        #ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)')        
-#        #plt.colorbar(c,ax=ax1)
-#        #ax1.set_title(title_format % (it*T,))
-#    
-#        #    plt.show()
-#        
-#        #    plt.figure(figsize=(12,5))
-#        #c=ax2.contourf(xx,zz,ww,20)
-#        #ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)')
-#        #ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)')        
-#        #plt.colorbar(c,ax=ax2)
-#        #ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,))
-#        #    plt.tight_layout()
-#        #    plt.show()
-#        #plt.savefig('fig_NH_3D_bubble/%02d.png'%it)
-#        
-#        time1, elapsed1 =time.time(), unst.getvar('elapsed')
-#        m,S,u,Phi,W = next_flow(m,S,u,Phi,W)
-#        time2, elapsed2 =time.time(), unst.getvar('elapsed')
-#        factor = 1000./nt
-#        print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
-#        factor = 1e9/(4*nt*nx*ny*llm)
-#        print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
-#            
+
+        v = mesh.field_mass() # specific volume (diagnosed)
+
+        for it in range(Nslice):
+            context.update_calendar(it+1)
+
+            # Diagnose quantities of interest from prognostic variables m,S,u,Phi,W
+            gas, w, z = diagnose(Phi,S,m,W)
+
+            # write to disk
+            context.send_field_primal('temp', gas.T)
+            context.send_field_primal('p', gas.p)
+            context.send_field_primal('theta', gas.s)
+            context.send_field_primal('uz', w)
+
+            print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g')
+
+            if args.mpi_ni*args.mpi_nj==1: plot()
+
+            time1, elapsed1 =time.time(), unst.getvar('elapsed')
+            m,S,u,Phi,W = next_flow(m,S,u,Phi,W)
+            time2, elapsed2 =time.time(), unst.getvar('elapsed')
+            factor = 1000./nt
+            print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
+            factor = 1e9/(4*nt*nx*ny*llm)
+            print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
+
+        context.update_calendar(Nslice+1) # make sure XIOS writes last iteration
+
+print('************DONE************')