source: codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py

from __future__ import print_function
from __future__ import division

from dynamico import getargs

getargs.add("--T", type=float, default=3600.,
                    help="Length of time slice in seconds")
getargs.add("--perturb", type=float, default=1.,
                    help="Amplitude of wind perturbation in m/s")
getargs.add("--N", type=int, default=48,
                    help="Number of time slices")
getargs.add("--Davies_N1", type=int, default=3)
getargs.add("--Davies_N2", type=int, default=3)
getargs.add("--nx", type=int, default=200)
getargs.add("--ny", type=int, default=30)
getargs.add("--betaplane", action='store_true')

getargs.add("--kappa_divgrad",  type=float, default=3.0e15)
getargs.add("--kappa_curlcurl", type=float, default=3.0e15)
getargs.add("--ztop", type=float, default=30000.)

getargs.defaults(dt=360., llm=50)

#getargs.config_log(filename='out.log',filemode='w') # must be done before calling Logger()
#    getargs.config_log(filename='out.log',filemode='w',
#                    format='%(processName)s:%(name)s:%(filename)s:%(module)s:%(funcName)s:%(lineno)d:%(levelname)s:%(message)s' )

args = getargs.parse()
log_master, log_world = getargs.getLogger()
INFO, DEBUG, DEBUG_ALL, ERROR = log_master.info, log_master.debug, log_world.debug, log_world.error

from dynamico.dev import unstructured as unst
from dynamico.dev import meshes
from dynamico.dev import xios
from dynamico.dev import precision as prec
from dynamico import dyn
from dynamico import time_step
from dynamico import DCMIP
from dynamico import maps
from dynamico.kernels import grad, curl, div, KE
from dynamico.LAM import Davies

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 create_pmesh(nx,ny):
    filename = 'cart_%03d_%03d.nc'%(nx,ny)
    INFO('Reading Cartesian mesh ...')
    meshfile = meshes.DYNAMICO_Format(filename)
    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
    return pmesh

def baroclinic_3D(pmesh, dx,Lx,Ly,llm,ztop):
    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)
    y0    = .5*Ly
    Cpd   = 1004.5
    p0    = 1e5
    up    = args.perturb # amplitude (m/s)
    xc,yc,lp = 0.,Ly/2.,600000.

    omega  = 7.292e-5  # Angular velocity of the Earth (s^-1)    
    phi0   = 45.*pi/180.0 # Reference latitude North pi/4 (deg)
    f0     = 2*omega*sin(phi0) 
    beta0  = 2*omega*cos(phi0)/a if args.betaplane else 0.
    fb     = f0 - beta0*y0

    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(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):
        u = -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/(b*b))) 
        # periodize the envelope of the initial perturbation
        A = exp(-(((x-xc)**2+(y-yc)**2)/lp**2))
        A = A + exp(-(((x-xc-Lx)**2+(y-yc)**2)/lp**2))
        A = A + exp(-(((x-xc+Lx)**2+(y-yc)**2)/lp**2))
        return  u + up*A

    def tmean(eta) : return T0*eta**(Rd*lap/g)
    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

    def eta(alpha) : return (1.-(lap*ztop*alpha/T0))**(g/(Rd*lap)) # roughly equispaced levels
    
    betaplane = maps.BetaPlaneMap(dx, .5*Lx, .5*Ly, f0, beta0, .5*Ly)
    nqdyn = 1
    mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, betaplane)

    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(np.shape(alpha_k),np.shape(alpha_l))
    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(y_il, eta_il)
    Phi_ik = Phi_xyeta(y_ik, eta_ik)
    p_ik, p_il = p(eta_ik), p(eta_il)
    T_ik = T(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])
#        mass_ik[:,l]=(p_il[:,l]-p_il[:,l+1])/g
    Sik, Wil  = gas.s*mass_ik, mesh.field_w()
    
    u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek)

    INFO('ztop (m) = %f %f' % (Phi_il[0,-1]/g, ztop) )
    ptop = p(eta(1.))
    INFO( 'ptop (Pa) = %f %f' % (gas.p[0,-1], ptop) )
    INFO('ztop(ptop) according to Eq. 7: %f' % (T0/lap*(1.-(ptop/p0)**(Rd*lap/g))) ) 

    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])
    INFO( 'min p, T : %f %s' % (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,u_ek,Phi_il,Wil]), gas

def diagnose(Phi,S,m,W,u):
    s=S/m
    curl_vk = curl(mesh, u)
    abs_vort_vk = mesh.field_z() # absolute vorticity
    un = mesh.field_u()
    v = mesh.field_mass() # specific volume
    w = mesh.field_mass()
    z = mesh.field_mass()
    for l in range(llm):
        v[:,l]=(Phi[:,l+1]-Phi[:,l])/(g*m[:,l])
        w[:,l]=.5*g*(W[:,l+1]+W[:,l])/m[:,l]
        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/g
        un[:,l]=u[:,l]/mesh.de
        abs_vort_vk[:,l]=curl_vk[:,l] + mesh.fv
    gas = thermo.set_vs(v,s)
    ps = gas.p[:,0]+ .5*g*m[:,0]
    return gas, w, z, ps, un, curl_vk, abs_vort_vk

class myDavies(Davies):
    def mask(self,x,y):
        DEBUG('davies dy = %f'% dx)
        return self.mask0(y,Ly,dx)*self.mask0(-y,0.,dx)
        
with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator
    comm = client.comm
    mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()

    INFO('%d/%d starting'%(mpi_rank,mpi_size))
    if mpi_rank>0 : 
        getargs.not_master()
        unst.setvar('is_mpi_master', False)
            
    g, Lx, Ly = 9.80616, 4e7, 6e6
    nx, ny, llm = args.nx, args.ny, args.llm
    dx = Lx/nx
    
    unst.setvar('g',g)
    unst.setvar('debug_hevi_solver', False)

    pmesh = create_pmesh(nx,ny)
    thermo, mesh, params, flow0, gas0 = baroclinic_3D(pmesh, dx,Lx,Ly,llm, args.ztop)

    mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0])
    unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk)

    T  = args.T
    dt = args.dt
    dz = flow0[3].max()/(params.g*llm)
    nt = int(math.ceil(T/dt))
    dt = T/nt
    INFO( '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)
        if args.hydrostatic:
            caldyn_step = unst.caldyn_step_HPE(mesh,scheme,1, caldyn_thermo,caldyn_eta,
                                               thermo,params,params.g)
        else:
            caldyn_step = unst.caldyn_step_NH(mesh,scheme,1, caldyn_thermo,caldyn_eta,
                                              thermo,params,params.g)
        davies = myDavies(args.Davies_N1, args.Davies_N2, 
                          mesh.lon_i, mesh.lat_i, mesh.lon_e,mesh.lat_e)
#        print('mask min/max :', davies.beta_i.min(), davies.beta_i.max() )
        DEBUG('mask min/max :')
        DEBUG('%f'% davies.beta_i.min())
        DEBUG('%f'% davies.beta_i.max())

        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
            for i in range(nt):
                caldyn_step.next()
                davies.relax(llm, caldyn_step, flow0)
                m,S,u = caldyn_step.mass, caldyn_step.theta_rhodz, caldyn_step.u
                s = S/m
                laps, bilaps = mesh.field_mass(), mesh.field_mass()
                lapu, bilapu = mesh.field_u(), mesh.field_u()
                unst.ker.dynamico_scalar_laplacian(s,laps)
                unst.ker.dynamico_scalar_laplacian(laps,bilaps)
                unst.ker.dynamico_curl_laplacian(u,lapu)
                unst.ker.dynamico_curl_laplacian(lapu,bilapu)
                caldyn_step.theta_rhodz[:] = S - dt*args.kappa_divgrad*bilaps*m # Euler step
                caldyn_step.u[:] = u - dt*args.kappa_curlcurl*bilapu # Euler step

            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*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)'

    #T, nslice, dt = 3600., 1, 3600.
    Nslice = args.N

    temp_v = mesh.field_z(),     

    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(Nslice+1):
            context.update_calendar(i+1)

            # Diagnose quantities of interest from prognostic variables m,S,u,Phi,W
            gas, w, z, ps, un, zeta_vk, zeta_abs_vk = diagnose(Phi,S,m,W,u)
            KE_ik = KE(mesh,u)
            du_fast, du_slow = caldyn_step.du_fast[0,:,:], caldyn_step.du_slow[0,:,:] 
            div_fast, div_slow = div(mesh,du_fast), div(mesh,du_slow)
            drhodz, hflux = caldyn_step.drhodz[0,:,:], caldyn_step.hflux[:,:]

            # write to disk
            context.send_field_primal('ps', ps)
            context.send_field_primal('temp', gas.T)

            context.send_field_primal('p', gas.p)
#            context.send_field_primal('p', KE_ik)
#            context.send_field_primal('p', drhodz)

            context.send_field_primal('theta', thermo.T0*exp(gas.s/thermo.Cpd))

            context.send_field_primal('uz', w)
            context.send_field_primal('Phi', Phi)
            context.send_field_primal('z', z)
            context.send_field_primal('div_fast', div_fast)

            context.send_field_primal('div_slow', div_slow)

            context.send_field_dual('curl', zeta_vk)
            context.send_field_dual('curl_abs', zeta_abs_vk)

            INFO( 'ptop, model top (m) : %f %f' % (unst.getvar('ptop'), Phi.max()/unst.getvar('g')) )

            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
            INFO( 'ms per full time step : %f %f' %(factor*(time2-time1), factor*(elapsed2-elapsed1)) )
            factor = 1e9/(4*nt*nx*ny*llm)
            INFO( 'nanosec per gridpoint per full time step : %f %f' %
                  (factor*(time2-time1), factor*(elapsed2-elapsed1)) )

INFO('************DONE************')