Index: codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py	(revision 768)
+++ codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py	(revision 769)
@@ -18,10 +18,13 @@
 parser = argparse.ArgumentParser()
 
-parser.add_argument("-mpi_ni", type=int,
+parser.add_argument("--mpi_ni", type=int,
                     default=64, choices=None,
                     help="number of x processors")
-parser.add_argument("-mpi_nj", type=int,
+parser.add_argument("--mpi_nj", type=int,
                     default=64, choices=None,
                     help="number of y processors")
+parser.add_argument("--T", type=float, default=5.,
+                    help="Length of time slice in seconds")
+
 args = parser.parse_args()
 
@@ -40,18 +43,18 @@
 
     omega  = 7.292e-5  # Angular velocity of the Earth (s^-1)    
-    phi0   = 45.       # Reference latitude North pi/4 (deg)
+    phi0   = 45.*np.pi/180.0 # 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):
+    def Phi_xy(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 Phi_xyeta(y,eta): return T0*g/lap*(1-eta**(Rd*lap/g)) + Phi_xy(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 T(y,eta) : return tmean(eta)+(Phi_xy(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
 
@@ -62,8 +65,9 @@
     def coriolis(lon,lat): return f0+0.*lon
     meshfile = meshes.DYNAMICO_Format(filename)
-    nqdyn, radius = 1, None
     pmesh = meshes.Unstructured_PMesh(comm,meshfile)
     pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
+    nqdyn, radius = 1, None
     mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)
+
 
     alpha_k = (np.arange(llm)  +.5)/llm
@@ -86,8 +90,8 @@
     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)
+    Phi_il = Phi_xyeta(y_il, eta_il)
+    Phi_ik = Phi_xyeta(y_ik, eta_ik)
     p_ik = p(eta_ik)
-    T_ik = T(x_ik, y_ik, eta_ik)  #ik full level(40), il(41)
+    T_ik = T(y_ik, eta_ik)  #ik full level(40), il(41)
 
     gas = thermo.set_pT(p_ik,T_ik)
@@ -118,4 +122,13 @@
     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
+
 with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator
     comm = client.comm
@@ -124,5 +137,5 @@
 
     g, Lx, Ly = 9.81, 4e7, 6e6
-    nx, ny, llm = 200, 30, 22
+    nx, ny, llm = 200, 30, 25
     dx,dy=Lx/nx,Ly/ny
 
@@ -131,12 +144,79 @@
     thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm)
 
-    T, nslice, dt = 3600., 1, 3600.
-
-    with xios.Context_Curvilinear(mesh,1, dt) as context:
+    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)
+
+    T=3600.
+    dt=360.
+    dz = flow0[3].max()/(params.g*llm)
+    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_entropy, unst.eta_mass
+
+    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):
+            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
+
+
+    #T, nslice, dt = 3600., 1, 3600.
+    Nslice=24
+
+    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):
+        v = mesh.field_mass() # specific volume (diagnosed)
+        
+        for i in range(Nslice):
             context.update_calendar(i)
-            print 'send_field', i, gas0.T.shape
-    #    context.send_field_primal('ps', lat_i)
+
+            # 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', gas0.T)
-            context.send_field_primal('p', gas0.p)
+            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************')
