source: codes/icosagcm/devel/Python/test/py/NH_3D_bubble_parallel.py @ 776

from dynamico import unstructured as unst
from dynamico import dyn
from dynamico import time_step
from dynamico import meshes
from dynamico import xios
from dynamico import precision as prec
import math as math
import matplotlib.pyplot as plt
import numpy as np
import time
import argparse

def thermal_bubble_3D(Lx,nx,Ly,ny,llm,ztop=1000., zc=350.,
                      rc=250, thetac=0.5, x0=0., y0=0.):
    Cpd, Rd, g, p0,theta0, T0 = 1004.5, 287.,9.81, 1e5, 300., 300.
    nqdyn = 1
    
    Phi   = lambda eta : g*ztop*eta
    p     = lambda Phi : p0*np.exp(-Phi/(Rd*T0))
    zz    = lambda p: -(Rd*T0*np.log(p/p0))/g
    rr    = lambda x,y,p: np.sqrt((x-x0)**2 + (y-y0)**2 + (zz(p)-zc)**2)
    sa    = lambda x,y,p: rr(x,y,p) < rc
    deform = lambda x,y,p: (0.5*thetac*(1+np.cos(np.pi*rr(x,y,p)/rc)))*sa(x,y,p)
    temp  =  lambda p: theta0*(p/p0)**(Rd/Cpd)
    T     = lambda x,y,p: deform(x,y,p) + temp(p)

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

    thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0)

    Phi_il = Phi(alpha_il)
    Phi_ik = Phi(alpha_ik)
    p_ik = p(Phi_ik)
    T_ik = T(x_ik, y_ik, p_ik)

    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])
    Sik, ujk, Wil  = gas.s*mass_ik, mesh.field_u(), mesh.field_w()
    
    print 'ztop (m) = ', Phi_il[0,-1]/g, ztop
    ptop = p(g*ztop)
    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(alpha_il[0])
    print 'min p, T :', pbot.min(), temp(pbot/p0)
    gas_bot = thermo.set_pT(pbot, temp(pbot/p0))
    params.pbot = gas_bot.p
    params.rho_bot = 1e6/gas_bot.v
    
    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

def reshape(data): return data.reshape((nx,ny))

def plot():
    x, y = map(reshape, (xx,yy) )
    zz=np.zeros((nx,ny,llm))
    ss=np.zeros((nx,ny,llm))
    ww=np.zeros((nx,ny,llm))
    x3d=np.zeros((nx,ny,llm))
     
    for l in range(llm):
        zz[:,:,l],ss[:,:,l],ww[:,:,l] = map(reshape, (z[:,l],gas.s[:,l],w[:,l]) )
        x3d[:,:,l]=x[:,:]

    jj=ny/2
    xp=x3d[:,jj,:]
    zp=zz[jj,:,:]/1000.
    sp=ss[jj,:,:]
    wp=ww[jj,:,:]

    f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4))

    c=ax1.contourf(xp,zp,sp,20)
    ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)')        
    plt.colorbar(c,ax=ax1)
    ax1.set_title(title_format % (it*T,))

    c=ax2.contourf(xp,zp,wp,20)
    ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)')        
    plt.colorbar(c,ax=ax2)
    ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,))
    plt.savefig('fig_NH_3D_bubble/%02d.png'%it)    

#------------------------- main program --------------------------

parser = argparse.ArgumentParser()

parser.add_argument("--mpi_ni", type=int, default=64, 
                    help="number of x processors")
parser.add_argument("--mpi_nj", type=int, default=64,
                    help="number of y processors")

parser.add_argument("--python_stepping", type=bool, default=False,
                    help="Time stepping in Python or Fortran")
parser.add_argument("--dt", type=float, default=.25,
                    help="Time step in seconds")
parser.add_argument("--T", type=float, default=5.,
                    help="Length of time slice in seconds")
parser.add_argument("--Nslice", type=int, default=10,
                    help="Number of time slices")

parser.add_argument("--thetac", type=float, default=30.,
                    help="Initial extra temperature of bubble")
parser.add_argument("--Lx", type=float, default=2000.,
                    help="Size of box in meters")
parser.add_argument("--Ly", type=float, default=2000.,
                    help="Size of box in meters")
parser.add_argument("--nx", type=int, default=20,
                    help="Resolution in the x direction")
parser.add_argument("--ny", type=int, default=20,
                    help="Resolution in the y direction")
parser.add_argument("--llm", type=int, default=79,
                    help="Number of vertical levels")

args = parser.parse_args()

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()
    print '%d/%d starting'%(mpi_rank,mpi_size)
    T, Nslice, dt, thetac = args.T, args.Nslice, args.dt, args.thetac
    Lx, nx, Ly, ny, llm   = args.Lx, args.nx, args.Ly, args.ny, args.llm
    
    nqdyn, dx = 1, Lx/nx
    Ly,ny,dy = Lx,nx,dx

    g=9.81
    unst.setvar('g',g)

    filename = 'cart_%03d_%03d.nc'%(nx,ny)
    print 'Reading Cartesian mesh ...'
    def coriolis(lon,lat): return 0.*lon
    meshfile = meshes.DYNAMICO_Format(filename)
    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
    mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, None, coriolis)

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

    # compute hybrid coefs from initial distribution of mass
    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)

    dz = flow0[3].max()/(params.g*llm)
    # courant = 2.8
    #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2))
    #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2))
    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_theta, unst.eta_mass
    caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
#    #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag

    if args.python_stepping: # 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.lon_i, mesh.lat_i

    # XIOS writes to disk every 24h
    # each iteration lasts it*nt seconds but 
    # we pretend that each iteration lasts 24h to make sure data is written to disk

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

            # 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', gas.T)
            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************')