Index: codes/icosagcm/devel/Python/test/py/NH_3D_bubble.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/NH_3D_bubble.py	(revision 641)
+++ codes/icosagcm/devel/Python/test/py/NH_3D_bubble.py	(revision 642)
@@ -53,6 +53,6 @@
     return thermo, mesh, params, (mass_ik,Sik,ujk,Phi_il,Wil), gas
 
-Lx, nx, llm, thetac, T, Nslice, courant = 2000., 100, 50, 30., 5., 10, 2.8
-#Lx, nx, llm, thetac, T, Nslice, courant = 2000., 50, 25, 1., 50., 10, 2.8
+#Lx, nx, llm, thetac, T, Nslice, courant = 2000., 100, 50, 30., 5., 10, 2.8
+Lx, nx, llm, thetac, T, Nslice, courant = 2000., 50, 25, 30, 5., 10, 2.8
 #Lx, nx, llm, thetac, T, Nslice, courant = 3000., 75, 25, -30, 5., 10, 2.8
 
@@ -61,11 +61,7 @@
 thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac)
 
-caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
-caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,params,thermo,params,params.g)
-
 # compute hybrid coefs from initial distribution of mass
 mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0])
 unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk)
-print mass_dbk
 
 dz = flow0[3].max()/(params.g*llm)
@@ -74,15 +70,34 @@
 nt = int(math.ceil(T/dt))
 dt = T/nt
-print 'Time step : %g s' % dt
+print 'Time step : %d x %g s' % (nt,dt)
 
-scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt)
+caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
+#caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag
 
-flow=flow0
+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
 w=mesh.field_mass()
 z=mesh.field_mass()
+
 for it in range(Nslice):
-    junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
-    m,S,u,Phi,W = flow
-    s=S/m ; # s=.5*(s+abs(s))
+    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]
@@ -92,7 +107,7 @@
     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))
     
-    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)')
@@ -100,7 +115,7 @@
     plt.colorbar(c,ax=ax1)
     ax1.set_title('Specific entropy at t=%g s (J/K/kg)' % (it*T,))
-#    plt.show()
+    #    plt.show()
     
-#    plt.figure(figsize=(12,5))
+    #    plt.figure(figsize=(12,5))
     c=ax2.contourf(xx,zz,ww,20)
     ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)')
@@ -108,10 +123,10 @@
     plt.colorbar(c,ax=ax2)
     ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,))
-#    plt.tight_layout()
-#    plt.show()
+    #    plt.tight_layout()
+    #    plt.show()
     plt.savefig('fig_NH_3D_bubble/%02d.png'%it)
     
     time1, elapsed1 =time.time(), unst.getvar('elapsed')
-    flow = scheme.advance(flow,nt)
+    m,S,u,Phi,W = next_flow(m,S,u,Phi,W)
     time2, elapsed2 =time.time(), unst.getvar('elapsed')
     factor = 1000./(4*nt)
Index: codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py	(revision 641)
+++ codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py	(revision 642)
@@ -38,6 +38,5 @@
 #scheme = time_step.RKn_simple(1,caldyn.bwd_fast_slow, dt) # forward Euler scheme
 print dt, scheme.csjl, scheme.cfjl
-step = unst.Caldyn_step(mesh,scheme)
-step.setup_TRSW()
+step = unst.caldyn_step_TRSW(mesh,scheme,nt)
 
 u0 = Omega*radius/12. # cf Williamson (1991), p.13
@@ -63,6 +62,5 @@
 for i in range(20):
     if True:
-        for j in range(nt):
-            step.next()
+        step.next()
         gh,u = step.mass, step.u
         print i, gh.shape, gh.min(), gh.max(), u.min(), u.max()
Index: codes/icosagcm/devel/Python/test/py/slice_GW_hydro.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/slice_GW_hydro.py	(revision 641)
+++ codes/icosagcm/devel/Python/test/py/slice_GW_hydro.py	(revision 642)
@@ -15,5 +15,4 @@
     caldyn_thermo = unst.thermo_entropy
     g = mesh.dx/metric.dx_g0
-    caldyn = unst.Caldyn_HPE(caldyn_thermo,caldyn_eta, mesh,metric,thermo,BC,g)
 
     Sik = m0ik*gas0.s
@@ -25,5 +24,5 @@
         S0ik = m0ik*gas0.s
 
-    u0=mesh.field_u() 
+    u0=mesh.field_u()
     u0[:,range(0,2*nx,2),:] = 100.*mesh.dx # Doppler shift 100 m/s
     flow0=(m0ik,S0ik,u0)
@@ -35,17 +34,30 @@
     dt=T/nt
     print 'nt,dt,dx', nt,dt,dx
-    scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt, a32=0.7)
-    #scheme = time_step.ARK3(caldyn.bwd_fast_slow, dt)
 
-    flow=flow0
+    m,S,u=flow0
+
+    if False: # time stepping in Python
+        caldyn = unst.Caldyn_HPE(caldyn_thermo,caldyn_eta, mesh,thermo,BC,g)
+        scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt, a32=0.7)
+        #scheme = time_step.ARK3(caldyn.bwd_fast_slow, dt)
+        def next_flow(m,S,u):
+            m,S,u = scheme.advance((m,S,u),nt)
+            return m,S,u,caldyn.geopot
+    else: # time stepping in Fortran
+        scheme = time_step.ARK2(None, dt, a32=0.7)
+        caldyn_step = unst.caldyn_step_HPE(mesh,scheme,nt, caldyn_thermo,caldyn_eta,thermo,BC,g)
+        def next_flow(m,S,u):
+            caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
+            caldyn_step.next()
+            return caldyn_step.mass, caldyn_step.theta_rhodz, caldyn_step.u, caldyn_step.geopot
+        
     for i in range(10):
         time1=time.time()
-        flow = scheme.advance(flow,nt)
+        m,S,u,geopot = next_flow(m,S,u)
         time2=time.time()
-        m,S,u = flow
-        print 'ms per time step : ', 1000*(time2-time1)/nt
-        print 'ptop, model top (m) :', unst.getvar('ptop'), caldyn.geopot.max()/unst.getvar('g')
-        junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
-        zz=caldyn.geopot[:,0:llm]/metric.g0/1000
+        print 'ms per time step : ', 1000*(time2-time1)/nt, 1000*unst.getvar('elapsed')/(nt*(i+1))
+        print 'ptop, model top (m) :', unst.getvar('ptop'), geopot.max()/unst.getvar('g')
+        #        junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
+        zz=geopot[:,0:llm]/metric.g0/1000
         plt.figure(figsize=(12,3))
         ucomp = mesh.ucomp(u)
