source: codes/icosagcm/devel/Python/test/py/slice_GW_hydro.py @ 632

from dynamico.meshes import Cartesian_mesh as Mesh
from dynamico import unstructured as unst
from dynamico import dyn
from dynamico import time_step
from dynamico import DCMIP

import math as math
import matplotlib.pyplot as plt
import numpy as np
import time

unst.setvar('nb_threads', 1)

def run(mesh,metric,thermo,BC,caldyn_eta, m0ik,gas0):
    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
    m0ik = m0ik/dx # NB : different definitions of mass field in DYNAMICO / Slice
    if caldyn_thermo == unst.thermo_theta:
        theta0 = thermo.T0*np.exp(gas0.s/thermo.Cpd)
        S0ik = m0ik*theta0
    else:
        S0ik = m0ik*gas0.s

    u0=mesh.field_u() 
    u0[:,range(0,2*nx,2),:] = 100.*mesh.dx # Doppler shift 100 m/s
    flow0=(m0ik,S0ik,u0)

    #T, courant = 60., 1.
    T, courant = 60., 2.9
    dt = courant*.5*dx/np.sqrt(gas0.c2.max())
    nt=int(T/dt)+1
    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
    for i in range(10):
        time1=time.time()
        flow = scheme.advance(flow,nt)
        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
        plt.figure(figsize=(12,3))
        ucomp = mesh.ucomp(u)
        plt.pcolor(xx,zz,ucomp[0,:,:]/dx)
        plt.ylim(0,10.)
        plt.colorbar()
        plt.title("u(t = %f s)" % ((i+1)*dt*nt))
        plt.savefig('fig_slice_GW_hydro/%02d.png'%i)
        plt.close()

Lx, nx, llm = 3e5, 300, 20
nqdyn, ny, dx = 1, 1, Lx/nx
mesh = Mesh(nx,ny,llm,nqdyn,nx*dx,ny*dx,0.)
xx,ll = mesh.xx[:,0,:]/1000, mesh.ll[:,0,:]
metric, thermo, BC, m0ik, gas0, Phi0_il, u0_jk = DCMIP.DCMIP31(Lx, nx, llm)

run(mesh,metric,thermo,BC,unst.eta_lag,  m0ik,gas0)