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

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

import math as math
import numpy as np
import time
import argparse
import logging
from numpy import pi, log, exp, sin, cos

# Baroclinic instability test based on Ullrich et al. 2015, QJRMS

parser = argparse.ArgumentParser()

parser.add_argument("--mpi_ni", type=int, default=1,
                    help="number of x processors")
parser.add_argument("--mpi_nj", type=int, default=1,
                    help="number of y processors")
parser.add_argument("--T", type=float, default=3600.,
                    help="Length of time slice in seconds")
parser.add_argument("--N", type=int, default=48,
                    help="Number of time slices")
parser.add_argument("--log",type=str,default='DEBUG',choices=('DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'),
                    help="Set verbosity level of logging")
parser.add_argument("--Davies_N1", type=int, default=3)
parser.add_argument("--Davies_N2", type=int, default=3)
parser.add_argument("--llm", type=int, default=50)
parser.add_argument("--nx", type=int, default=200)
parser.add_argument("--ny", type=int, default=30)
parser.add_argument("--dt", type=float, default=360., help='Time step in seconds')
parser.add_argument("--hydrostatic", action='store_true')
parser.add_argument("--betaplane", action='store_true')

args = parser.parse_args()


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    = .5*Ly
    Cpd   = 1004.5
    p0    = 1e5
    up    = 1. # amplitude (m/s)
    xc,yc,lp = 0.,Ly/2.,600000.

    omega  = 7.292e-5  # Angular velocity of the Earth (s^-1)    
    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 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))) 
#        u = u + up*exp(-(((x-xc)**2+(y-yc)**2)/lp**2))
        return  u

    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
    def coriolis(x,y): return f0+beta0*y # here y is 0 at domain center

    filename = 'cart_%03d_%03d.nc'%(nx,ny)
    logging.info('Reading Cartesian mesh ...')
    meshfile = meshes.DYNAMICO_Format(filename)
    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
    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+y0, 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+y0, 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+y0, alpha_k, indexing='ij')

    print('ztop(ptop) according to Eq. 7:', T0/lap*(1.-(ptop/p0)**(Rd*lap/g))) 
    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)

    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,u_ek,Phi_il,Wil]), gas

def diagnose(Phi,S,m,W):
    s=S/m
    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

def grad(mesh, s_ik):
    llm=s_ik.shape[1]
    edge_num, left, right = mesh.edge_num, mesh.left, mesh.right
    grad_ek = prec.zeros((edge_num, llm))
    for e in range(edge_num):
        left_e = left[e]-1 # Fortran => Python indexing
        right_e = right[e]-1 # Fortran => Python indexing
        grad_ek[e,:] = s_ik[right_e,:]-s_ik[left_e,:]
    return grad_ek

def div(mesh, u_ek):
    llm=u_ek.shape[1]
    primal_num, primal_deg, primal_edge = mesh.primal_num, mesh.primal_deg, mesh.primal_edge
    primal_ne, le_de, Ai = mesh.primal_ne, mesh.le_de, mesh.Ai
    div_ik = prec.zeros((primal_num, llm))
    for i in range(primal_num):
        div_i = 0.
        for iedge in range(primal_deg[i]):
            e = primal_edge[i,iedge]-1 # Fortran => Python indexing
            div_i = div_i + u_ek[e,:]*primal_ne[i,iedge]*le_de[e]
        div_ik[i,:] = div_i / Ai[i]
    return div_ik

def curl(mesh, u_ek):
    llm=u_ek.shape[1]
    dual_num, dual_deg, dual_edge, dual_ne, Av = mesh.dual_num, mesh.dual_deg, mesh.dual_edge, mesh.dual_ne, mesh.Av
    curl_vk = prec.zeros((dual_num, llm))
    for v in range(dual_num):
        curl_v = 0.
        for iedge in range(dual_deg[v]):
            e = dual_edge[v,iedge]-1 # Fortran => Python indexing
            curl_v = curl_v + u_ek[e,:]*dual_ne[v,iedge]
        curl_vk[v,:] = curl_v / Av[v]
    return curl_vk

class Davies:
    def __init__(self,N1,N2,x_i,y_i,x_e,y_e):
        self.N1, self.N2 = N1, N2
        self.beta_i = self.mask(x_i,y_i)
        self.beta_e = self.mask(x_e,y_e)
        self.x_e,self.y_e = x_e,y_e
    def mask0(self,c,c0,delta): # 1D building block for Davies relaxation
        N1, N2 = self.N1, self.N2
        m = np.zeros(c.size)
        c1,c2 = c0-N1*delta, c0-(N1+N2)*delta
        for i in range(c.size):
            ci=c[i]
            m[i] = (1.+np.cos((ci-c2)*np.pi/(N2*delta)))/2.0
            if ci < c2 : m[i]=1.
            if ci > c1 : m[i]=0.
        return m
    def relax(self, llm, step, flow):
        beta_i, beta_e = self.beta_i, self.beta_e
        m,S,u,Phi,W=flow
        for l in range(llm):
            step.mass[:,l]        = beta_i*step.mass[:,l]        + (1.-beta_i)*m[:,l]
            step.theta_rhodz[:,l] = beta_i*step.theta_rhodz[:,l] + (1.-beta_i)*S[:,l]
            step.u[:,l]           = beta_e*step.u[:,l]           + (1.-beta_e)*u[:,l]
        for l in range(llm+1):
            step.geopot[:,l]      = beta_i*step.geopot[:,l]      + (1.-beta_i)*Phi[:,l]
            step.W[:,l]           = beta_i*step.W[:,l]           + (1.-beta_i)*W[:,l]

class myDavies(Davies):
    def mask(self,x,y):
        logging.debug('davies dy = %f' % dy)
        return self.mask0(y,.5*Ly,dy)*self.mask0(-y,.5*Ly,dy)


        
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()

    #-------------Logging verbosity configuration---------------------------
    myloglevel = args.log
    myloglevel = getattr(logging, myloglevel.upper())
    logging.basicConfig(filename='out.log',filemode='w',level=myloglevel, 
        format='%(processName)s:%(name)s:%(filename)s:%(module)s:%(funcName)s:%(lineno)d:%(levelname)s:%(message)s' )

    logging.debug('%d/%d starting'%(mpi_rank,mpi_size))

    g, Lx, Ly = 9.81, 4e7, 6e6
    nx, ny, llm = args.nx, args.ny, args.llm
    dx,dy = Lx/nx,Ly/ny

    unst.setvar('g',g)

    thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm)

    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  = args.T
    dt = args.dt
    dz = flow0[3].max()/(params.g*llm)
    nt = int(math.ceil(T/dt))
    dt = T/nt
    logging.info( 'Time step : %d x %g = %g s' % (nt,dt,nt*dt))
    
    dt=0. # FIXME

    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() )
        logging.debug('mask min/max :')
        logging.debug(davies.beta_i.min())
        logging.debug(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
            kappa = dx**4/43200.
            for i in range(nt):
                caldyn_step.next()
                davies.relax(llm, caldyn_step, flow0)
                m,S = caldyn_step.mass, caldyn_step.theta_rhodz
                s = S/m
                laps = mesh.field_mass()
                bilaps = mesh.field_mass()
                unst.ker.dynamico_scalar_laplacian(s,laps)
                unst.ker.dynamico_scalar_laplacian(laps,bilaps)
#                grads = grad(mesh,s)
#                laps = div(mesh,grads)
#                gradlaps = grad(mesh,laps) # bilaplacian
#                bilaps = div(mesh,gradlaps)                
                caldyn_step.theta_rhodz[:] = S - dt*kappa*bilaps*m # 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*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()

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

    with xios.Context_Curvilinear(mesh,1, 24*3600) as context:
        # now XIOS knows about the mesh and we can write to disk
        v = mesh.field_mass() # specific volume (diagnosed)
        
        for i in range(Nslice):
            context.update_calendar(i+1)

            # Diagnose quantities of interest from prognostic variables m,S,u,Phi,W
            gas, w, z = diagnose(Phi,S,m,W)
            ps = caldyn_step.ps
            curl_vk = curl(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)
            # 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('theta', gas.s)
            context.send_field_primal('uz', w)
            context.send_field_primal('z', z)
            context.send_field_primal('div_fast', div_fast)
            context.send_field_primal('div_slow', div_slow)
            print curl_vk.size
            context.send_field_dual('curl', curl_vk)

            print( 'ptop, model top (m) :', 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
            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))

logging.info('************DONE************')