Index: /codes/icosagcm/devel/Python/test/py/DCMIP2008c5_MPI.py
===================================================================
--- /codes/icosagcm/devel/Python/test/py/DCMIP2008c5_MPI.py	(revision 697)
+++ /codes/icosagcm/devel/Python/test/py/DCMIP2008c5_MPI.py	(revision 697)
@@ -0,0 +1,215 @@
+from dynamico import unstructured as unst
+from dynamico import dyn
+from dynamico import time_step
+from dynamico import DCMIP
+from dynamico import meshes
+import dynamico.xios as xios
+
+import math as math
+import matplotlib.pyplot as plt
+import numpy as np
+import time
+
+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)
+
+#------------------------ initial condition -------------------------
+
+# Parameters for the simulation
+g              = 9.80616   # gravitational acceleration in meters per second squared
+omega          = 7.29211e-5  # Earth's angular velocity in radians per second
+f0             = 2.0*omega   # Coriolis parameter
+u_0            = 20.0        # velocity in meters per second
+T_0            = 288.0       # temperature in Kelvin
+d2             = 1.5e6**2    # square of half width of Gaussian mountain profile in meters
+h_0            = 2.0e3       # mountain height in meters
+lon_c          = np.pi/2.0   # mountain peak longitudinal location in radians
+lat_c          = np.pi/6.0   # mountain peak latitudinal location in radians
+radius         = 6.371229e6  # mean radius of the Earth in meters
+ref_press      = 100145.6    # reference pressure (mean surface pressure) in Pascals
+ref_surf_press = 930.0e2     # South Pole surface pressure in Pascals
+Rd             = 287.04      # ideal gas constant for dry air in joules per kilogram Kelvin
+Cpd            = 1004.64     # specific heat at constant pressure in joules per kilogram Kelvin
+kappa          = Rd/Cpd      # kappa=R_d/c_p
+N_freq         = np.sqrt(g**2/(Cpd*T_0)) # Brunt-Vaisala buoyancy frequency
+N2, g2kappa = N_freq**2, g*g*kappa
+
+def DCMIP2008c5(grid,llm):
+    def Phis(lon,lat): 
+        rgrc = radius*np.arccos(np.sin(lat_c)*np.sin(lat)+np.cos(lat_c)*np.cos(lat)*np.cos(lon-lon_c))
+        return g*h_0*np.exp(-rgrc**2/d2)
+    def ps(lon, lat, Phis): 
+        return ref_surf_press * np.exp(
+            - radius*N2*u_0/(2.0*g2kappa)*(u_0/radius+f0)*(np.sin(lat)**2-1.)
+            - N2/(g2kappa)*Phis )
+    def ulon(lat): return u_0*np.cos(lat)
+    def ulat(lat): return 0.*lat
+    def f(lon,lat): return f0*np.sin(lat) # Coriolis parameter
+
+    nqdyn, preff, Treff = 1, 1e5, 300.
+    thermo = dyn.Ideal_perfect(Cpd, Rd, preff, Treff)
+
+    meshfile = meshes.MPAS_Format('grids/x1.%d.grid.nc'%grid)
+#    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
+    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
+    mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, f)
+    mesh.pmesh=pmesh
+
+    lev,levp1 = np.arange(llm),np.arange(llm+1)
+    lon_i, lat_i, lon_e, lat_e =  mesh.lon_i, mesh.lat_i, mesh.lon_e, mesh.lat_e
+    lat_ik,k_i = np.meshgrid(mesh.lat_i,lev, indexing='ij')
+    lon_ik,k_i = np.meshgrid(mesh.lon_i,lev, indexing='ij')
+    lat_il,l_i = np.meshgrid(mesh.lat_i,levp1, indexing='ij')            
+    lon_il,l_i = np.meshgrid(mesh.lon_i,levp1, indexing='ij')
+    lat_ek,k_e = np.meshgrid(mesh.lat_e,lev, indexing='ij')
+
+    Phis_i = Phis(lon_i, lat_i)
+    ps_i = ps(lon_i, lat_i, Phis_i)
+
+    if llm==18:
+        ap_l=[0.00251499, 0.00710361, 0.01904260, 0.04607560, 0.08181860,
+              0.07869805, 0.07463175, 0.06955308, 0.06339061, 0.05621774, 0.04815296,
+              0.03949230, 0.03058456, 0.02193336, 0.01403670, 0.007458598, 0.002646866,
+              0.0, 0.0 ]
+        bp_l=[0.0, 0.0, 0.0, 0.0, 0.0, 0.03756984, 0.08652625, 0.1476709, 0.221864,
+              0.308222, 0.4053179, 0.509588, 0.6168328, 0.7209891, 0.816061, 0.8952581, 
+              0.953189, 0.985056, 1.0 ]
+    if llm==26:
+        ap_l=[0.002194067, 0.004895209, 0.009882418, 0.01805201, 0.02983724, 0.04462334, 0.06160587, 
+               0.07851243, 0.07731271, 0.07590131, 0.07424086, 0.07228744, 0.06998933, 0.06728574, 0.06410509, 
+               0.06036322, 0.05596111, 0.05078225, 0.04468960, 0.03752191, 0.02908949, 0.02084739, 0.01334443, 
+               0.00708499, 0.00252136, 0.0, 0.0 ]
+        bp_l=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01505309, 0.03276228, 0.05359622, 
+                0.07810627, 0.1069411, 0.1408637, 0.1807720, 0.2277220, 0.2829562, 0.3479364, 0.4243822, 
+                0.5143168, 0.6201202, 0.7235355, 0.8176768, 0.8962153, 0.9534761, 0.9851122, 1.0 ]
+
+    ap_l, bp_l = ref_press*np.asarray(ap_l[-1::-1]), bp_l[-1::-1] # reverse indices
+    ptop = ap_l[-1] # pressure BC
+    
+    print ptop, ap_l, bp_l
+    ps_il,ap_il = np.meshgrid(ps_i,ap_l, indexing='ij')
+    ps_il,bp_il = np.meshgrid(ps_i,bp_l, indexing='ij')
+    hybrid_coefs = meshes.mass_coefs_from_pressure_coefs(g, ap_il, bp_il)
+    pb_il = ap_il + bp_il*ps_il
+    mass_ik, pb_ik = mesh.field_mass(), mesh.field_mass()
+    for l in range(llm):
+        pb_ik[:,l]=.5*(pb_il[:,l]+pb_il[:,l+1])
+        mass_ik[:,l]=(pb_il[:,l]-pb_il[:,l+1])/g
+    Tb_ik = T_0 + 0.*pb_ik
+    gas = thermo.set_pT(pb_ik,Tb_ik)
+    Sik  = gas.s*mass_ik
+# start at hydrostatic geopotential
+    Phi_il = mesh.field_w()
+    Phi_il[:,0]=Phis_i
+    for l in range(llm):
+        Phi_il[:,l+1] = Phi_il[:,l] + g*mass_ik[:,l]*gas.v[:,l]
+
+    ujk, Wil = mesh.ucov3D(ulon(lat_ek),ulat(lat_ek)), mesh.field_w()
+    
+    print 'ztop (m) = ', Phi_il[0,-1]/g
+    print 'ptop (Pa) = ', gas.p[0,-1], ptop
+    dx=mesh.de.min()
+    params=dyn.Struct()
+    params.g, params.ptop = g, ptop
+    params.dx, params.dx_g0 = dx, dx/g
+    params.pbot, params.rho_bot = 1e5+0.*mesh.lat_i, 1e6+0.*mesh.lat_i
+    return thermo, mesh, hybrid_coefs, params, (mass_ik,Sik,ujk,Phi_il,Wil), gas
+
+
+#------------------------ main program -------------------------
+
+#grid, llm = 40962, 26
+nqtot, llm, grid = 1, 26, 2562
+T, Nslice, courant = 14400, 24, 3.0
+caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag
+#caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
+thermo, mesh, hybrid_coefs, params, flow0, gas0 = DCMIP2008c5(grid,llm)
+llm, dx = mesh.llm, params.dx
+print 'llm, nb_hex, dx =', llm, mesh.Ai.size, dx
+if caldyn_eta == unst.eta_lag:
+    print 'Lagrangian coordinate.'
+else:
+    print 'Mass-based coordinate.'
+
+unst.ker.dynamico_init_hybrid(*hybrid_coefs)
+    
+dt = courant*.5*dx/np.sqrt(gas0.c2.max())
+
+nt = int(math.ceil(T/dt))
+dt = T/nt
+print 'Time step : %g s' % dt
+
+#mesh.plot_e(mesh.le/mesh.de) ; plt.title('le/de')
+#plt.savefig('fig_DCMIP2008c5/le_de.png'); plt.close()
+
+#mesh.plot_i(mesh.Ai) ; plt.title('Ai')
+#plt.savefig('fig_DCMIP2008c5/Ai.png'); plt.close()
+
+scheme = time_step.ARK2(None, dt, a32=0.7)
+caldyn_step = unst.caldyn_step_HPE(mesh,scheme,nt, caldyn_thermo,caldyn_eta, thermo,params,params.g)
+def next_flow(m,S,u):
+    caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
+    caldyn_step.remap()
+    caldyn_step.next()
+    return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), 
+            caldyn_step.u.copy(), caldyn_step.geopot.copy())
+
+def plots(it):
+    s=S/m
+    for l in range(llm):
+        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
+        vol[:,l]=(Phi[:,l+1]-Phi[:,l])/params.g/m[:,l] # specific volume
+    gas = thermo.set_vs(vol, s)
+    s=.5*(s+abs(s))
+    t = (it*T)/3600
+    print( 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g') )
+    mesh.plot_i(gas.T[:,llm/2])
+    plt.title('T at t=%dh'%(t))
+    plt.savefig('fig_DCMIP2008c5/T%02d.png'%it)
+    plt.close()
+
+    mesh.plot_i(m[:,llm/2])
+    plt.title('mass at t=%dh'%(t))
+    plt.savefig('fig_DCMIP2008c5/m%02d.png'%it)
+    plt.close()
+
+    mesh.plot_i(Phi[:,0])
+    plt.title('Surface geopotential at t=%dh'%(t))
+    plt.savefig('fig_DCMIP2008c5/Phis%02d.png'%it)
+    plt.close()
+
+z, vol = mesh.field_mass(), mesh.field_mass()
+m,S,u,Phi,W = flow0 
+caldyn_step.geopot[:,0]=Phi[:,0]
+#plots(0)
+
+context=xios.XIOS_Context(mesh.pmesh,mesh,nqtot, T)
+
+for it in range(Nslice):
+    context.update_calendar(it)
+    unst.setvar('debug_hevi_solver',False)
+    time1, elapsed1 =time.time(), unst.getvar('elapsed')
+    m,S,u,Phi = next_flow(m,S,u)
+    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*m.size)
+    print 'nanosec per gridpoint per call to caldyn_hevi : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
+    
+    s=S/m
+    s=.5*(s+abs(s))
+    for l in range(llm):
+        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
+        vol[:,l]=(Phi[:,l+1]-Phi[:,l])/params.g/m[:,l] # specific volume
+    gas = thermo.set_vs(vol, s)
+    ss = np.asarray(gas.T, dtype=np.double)
+    context.send_field_primal('theta', ss)
+    #plots(it+1)
+
+print 'xios.context_finalize()'
+context.finalize()
+print 'xios.finalize()'
+xios.finalize()
+print 'Done'
