Index: codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py	(revision 761)
+++ codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py	(revision 761)
@@ -0,0 +1,234 @@
+from dynamico import unstructured as unst
+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
+
+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
+
+# Baroclinic instability test based on Ullrich et al. 2015, QJRMS
+
+def baroclinic_3D(Lx,nx,Ly,ny,llm,ztop=25000.):
+    Rd    = 287.0      # Gas constant for dryy air (j kg^-1 K^-1)
+    T0    = 288.0      # Reference temperature (K) 
+    lap   = 0.005      # Lapse rate (K m^-1)
+    b     = 2.         # Non dimensional vertical width parameter
+    u0    = 35.        # Reference zonal wind speed (m s^-1)
+    a     = 6.371229e6 # Radius of the Earth (m)
+    ptop  = 2000.
+    y0    = Ly*0.5
+    Cpd   = 1004.5
+    p0    = 1e5
+
+    omega  = 7.292e-5  # Angular velocity of the Earth (s^-1)    
+    phi0   = 45.       # Reference latitude North pi/4 (deg)
+    f0     = 2*omega*np.sin(phi0) 
+    beta0  = 2*omega*np.cos(phi0)/a
+    fb     = 2*omega*np.sin(phi0) - y0*2*omega*np.cos(phi0)/a
+
+    def Phi_xy(x,y):
+        fc = y*y - (Ly*y/pi)*sin(2*pi*y/Ly)
+        fd = Ly*Ly/(2*pi*pi)*cos(2*pi*y/Ly)
+        return .5*u0*( fb*(y-y0-Ly/(2*pi)*sin(2*pi*y/Ly)) + .5*beta0*(fc-fd-(Ly*Ly/3.)- Ly*Ly/(2*pi*pi)) )
+
+    def Phi_xyeta(x,y,eta): return T0*g/lap*(1-eta**(Rd*lap/g)) + Phi_xy(x,y)*log(eta)*exp(-((log(eta)/b)**2))
+    def ulon(x,y,eta): return -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/b/b))
+    def tmean(eta) : return T0*eta**(Rd*lap/g)
+    def T(x,y,eta) : return tmean(eta)+(Phi_xy(x,y)/Rd)*(((2/(b*b))*(log(eta))**2)-1)*exp(-((0.5*log(eta))**2)) 
+    def p(eta): return p0*eta     # eta  = p/p_s
+
+    def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels
+
+    filename = 'cart_%03d_%03d.nc'%(nx,ny)
+    print 'Reading Cartesian mesh ...'
+    def coriolis(lon,lat): return f0+0.*lon
+    meshfile = meshes.DYNAMICO_Format(filename)
+    nqdyn, radius = 1, None
+#    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, coriolis)
+    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
+    mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)
+
+    alpha_k = (np.arange(llm)  +.5)/llm
+    alpha_l = (np.arange(llm+1)+ 0.)/llm
+    x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij')
+    y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij')
+    x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij')
+    y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij')
+    x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij')
+    y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij')
+
+    print('----------------')
+    print 'ztop(ptop) according to Eq. 7:', T0/lap*(1.-(ptop/p0)**(Rd*lap/g)) 
+    print(np.shape(alpha_k),np.shape(alpha_l))
+    print(mesh.__dict__.keys())
+    thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0)
+
+    eta_il = eta(alpha_il)
+    eta_ik = eta(alpha_ik)
+    eta_ek = eta(alpha_ek)
+    print('min max eta_il', np.min(eta_il),np.max(eta_il))
+
+    Phi_il = Phi_xyeta(x_il, y_il, eta_il)
+    Phi_ik = Phi_xyeta(x_ik, y_ik, eta_ik)
+    p_ik = p(eta_ik)
+    T_ik = T(x_ik, y_ik, eta_ik)  #ik full level(40), il(41)
+
+    gas = thermo.set_pT(p_ik,T_ik)
+    mass_ik = mesh.field_mass()
+    for l in range(llm): 
+        mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l])
+    Sik, ujk, Wil  = gas.s*mass_ik, mesh.field_u(), mesh.field_w()
+    
+    print(np.shape(ujk))
+    print('P_ik',p_ik[0,:])
+
+    u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek)
+
+    print(np.shape(u_ek))
+
+    print 'ztop (m) = ', Phi_il[0,-1]/g, ztop
+    ptop = p(eta(1.))
+    print 'ptop (Pa) = ', gas.p[0,-1], ptop
+
+    params=dyn.Struct()
+    params.ptop=ptop
+    params.dx=dx
+    params.dx_g0=dx/g
+    params.g = g
+
+    # define parameters for lower BC
+    pbot = p(eta_il[:,0])
+    print 'min p, T :', pbot.min(), tmean(pbot/p0)
+    gas_bot = thermo.set_pT(pbot, tmean(pbot/p0))
+    params.pbot = gas_bot.p
+    params.rho_bot = 1e6/gas_bot.v
+    
+    return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas
+
+
+g, Lx, Ly = 9.81, 4e7, 6e6
+nx, ny, llm = 200, 30, 22
+dx,dy=Lx/nx,Ly/ny
+
+unst.setvar('g',g)
+
+thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm)
+
+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)
+        
