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

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 matplotlib.pyplot as plt
import numpy as np
import time
import argparse

parser = argparse.ArgumentParser()

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


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)
    phi0  = 45.       # Reference latitude North pi/4 (deg)
    omega = 7.292e-5  # Angular velocity of the Earth (s^-1)
    f0    = 2*omega*np.sin(phi0)
    lap   = 0.005      # Lapse rate (K m^-1)
    Rd    = 287.0      # Gas constant for dryy air (j kg^-1 K^-1)

    def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels
    #def eta(alpha) : return alpha/float(llm)

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

    print(mesh.__dict__.keys())

    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('alpha_l=',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('eta_il=',eta_il)
    #Phi_il = Phi(mesh.llp1/float(llm)) #llp is not defined in local mesh
    #Phi_ik = Phi((mesh.ll+.5)/llm)
    Phi_il = Phi(eta_il)
    Phi_ik = Phi(eta_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
    #pbot = p(Phi_il[:,:,0])
    #gas_bot = thermo.set_pT(pbot, temp(pbot))
    # define parameters for lower BC
    pbot = p(eta_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

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)

    #Lx, nx, llm, thetac, T, Nslice, courant = 2000., 100, 50, 30., 5., 10, 2.8
    Lx, nx, llm, thetac = 2000., 200, 79, 30
    #Lx, nx, llm, thetac, T, Nslice, courant = 3000., 75, 25, -30, 5., 10, 2.8

    nqdyn, dx = 1, Lx/nx
    Ly,ny,dy = Lx,nx,dx

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

    print('bubble call-----------')
    thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac)
    print('bubble done-----------')


    # 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)
    #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 s' % (nt,dt)


    T, nslice, dt = 3600., 1, 3600.
    #T, nslice  = 3600., 4

    with xios.Context_Curvilinear(mesh,1, dt) as context:
        # now XIOS knows about the mesh and we can write to disk
        for i in range(48): # 2 days
            context.update_calendar(i)
            
            print 'send_field', i, gas0.T.shape
    #    context.send_field_primal('ps', lat_i)
            context.send_field_primal('temp', gas0.T)
            context.send_field_primal('p', gas0.p)
    print(gas0.__dict__.keys())
    print('size of T, p',np.shape(gas0.T),np.shape(gas0.p))
    print('************DONE************')

#    #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 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
#    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()
#    
#    for it in range(Nslice):
#        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]
#            w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l]
#            #z[:,:,l]=.5*(Phi[:,:,l+1]+Phi[:,:,l])/params.g
#            z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
#            
#        print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g')
#        jj=ny/2
#        #xx,zz,ss,ww = mesh.xx[jj,:,:]/1000., z[jj,:,:]/1000., s[jj,:,:], w[jj,:,:]
#        #xx,zz,ss,ww = x_ik[jj,:]/1000., z[jj,:]/1000., s[jj,:], w[jj,:]
#    
#        #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)')
#        #ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)')        
#        #plt.colorbar(c,ax=ax1)
#        #ax1.set_title(title_format % (it*T,))
#    
#        #    plt.show()
#        
#        #    plt.figure(figsize=(12,5))
#        #c=ax2.contourf(xx,zz,ww,20)
#        #ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)')
#        #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.tight_layout()
#        #    plt.show()
#        #plt.savefig('fig_NH_3D_bubble/%02d.png'%it)
#        
#        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)
#