Changeset 763 for codes/icosagcm

Timestamp:

10/18/18 15:13:21 (6 years ago)

Author:

jisesh

Message:

devel/Python : writing data in parallel

File:

• codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py (modified) (3 diffs)

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

@@ -8 +8 @@
 from dynamico.meshes import Cartesian_mesh as Mesh
 
-from mpi4py import MPI
-comm = MPI.COMM_WORLD
-mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
-print '%d/%d starting'%(mpi_rank,mpi_size)
-
 import math as math
 import numpy as np
 import time
-
 from numpy import pi, log, exp, sin, cos
 
@@ -57 +51 @@
     meshfile = meshes.DYNAMICO_Format(filename)
     nqdyn, radius = 1, None
-#    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, coriolis)
     pmesh = meshes.Unstructured_PMesh(comm,meshfile)
+    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
     mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)
 
@@ -118 +112 @@
     return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas
 
+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)
 
-g, Lx, Ly = 9.81, 4e7, 6e6
-nx, ny, llm = 200, 30, 22
-dx,dy=Lx/nx,Ly/ny
+    g, Lx, Ly = 9.81, 4e7, 6e6
+    nx, ny, llm = 200, 30, 22
+    dx,dy=Lx/nx,Ly/ny
 
-unst.setvar('g',g)
+    unst.setvar('g',g)
 
-thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm)
+    thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm)
 
-T, nslice, dt = 3600., 1, 3600.
+    T, nslice, dt = 3600., 1, 3600.
 
-context = xios.XIOS_Context_Curvilinear(mesh,1, dt)
-
-for i in range(48):
-    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)
-
-# to do at the end
-print 'xios.context_finalize()'
-context.finalize()
-print 'xios.finalize()'
-xios.finalize()
-print 'Done'
-
-
-
-# compute hybrid coefs from initial distribution of mass
-mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0])
-print 'Type of mass_bl, mass_dak, mass_dbk : ', [x.dtype for x in mass_bl, mass_dak, mass_dbk]
-unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk)
-
-dz = flow0[3].max()/(params.g*llm)
-#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)
-
-#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()
-
-Nslice=0
-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]
-        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,:]
-
-    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_baroclinic_3D/%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)
-        
+    with xios.Context_Curvilinear(mesh,1, dt) as context:
+        # now XIOS knows about the mesh and we can write to disk
+        for i in range(48):
+            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)