source: codes/icosagcm/devel/Python/test/py/DCMIP2008c5_MPI.py @ 798

# apparently we should initialize MPI before doing anything else
from mpi4py import MPI
comm = MPI.COMM_WORLD
mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
print '%d/%d starting'%(mpi_rank,mpi_size)

# now start doing something useful
from dynamico import unstructured as unst
from dynamico import dyn
from dynamico import time_step
from dynamico import DCMIP
from dynamico import meshes
#import dynamico.xios as xios

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

#------------------------ initial condition -------------------------

# Parameters for the simulation
g              = 9.80616   # gravitational acceleration in meters per second squared
omega          = 7.29211e-5  # Earth's angular velocity in radians per second
f0             = 2.0*omega   # Coriolis parameter
u_0            = 20.0        # velocity in meters per second
T_0            = 288.0       # temperature in Kelvin
d2             = 1.5e6**2    # square of half width of Gaussian mountain profile in meters
h_0            = 2.0e3       # mountain height in meters
lon_c          = np.pi/2.0   # mountain peak longitudinal location in radians
lat_c          = np.pi/6.0   # mountain peak latitudinal location in radians
radius         = 6.371229e6  # mean radius of the Earth in meters
ref_press      = 100145.6    # reference pressure (mean surface pressure) in Pascals
ref_surf_press = 930.0e2     # South Pole surface pressure in Pascals
Rd             = 287.04      # ideal gas constant for dry air in joules per kilogram Kelvin
Cpd            = 1004.64     # specific heat at constant pressure in joules per kilogram Kelvin
kappa          = Rd/Cpd      # kappa=R_d/c_p
N_freq         = np.sqrt(g**2/(Cpd*T_0)) # Brunt-Vaisala buoyancy frequency
N2, g2kappa = N_freq**2, g*g*kappa

def DCMIP2008c5(grid,llm):
    def Phis(lon,lat): 
        rgrc = radius*np.arccos(np.sin(lat_c)*np.sin(lat)+np.cos(lat_c)*np.cos(lat)*np.cos(lon-lon_c))
        return g*h_0*np.exp(-rgrc**2/d2)
    def ps(lon, lat, Phis): 
        return ref_surf_press * np.exp(
            - radius*N2*u_0/(2.0*g2kappa)*(u_0/radius+f0)*(np.sin(lat)**2-1.)
            - N2/(g2kappa)*Phis )
    def ulon(lat): return u_0*np.cos(lat)
    def ulat(lat): return 0.*lat
    def f(lon,lat): return f0*np.sin(lat) # Coriolis parameter

    nqdyn, preff, Treff = 1, 1e5, 300.
    thermo = dyn.Ideal_perfect(Cpd, Rd, preff, Treff)

    meshfile = meshes.MPAS_Format('grids/x1.%d.grid.nc'%grid)
#    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
    pmesh.partition_metis()
    mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, f)
    mesh.pmesh=pmesh

    lev,levp1 = np.arange(llm),np.arange(llm+1)
    lon_i, lat_i, lon_e, lat_e =  mesh.lon_i, mesh.lat_i, mesh.lon_e, mesh.lat_e
    lat_ik,k_i = np.meshgrid(mesh.lat_i,lev, indexing='ij')
    lon_ik,k_i = np.meshgrid(mesh.lon_i,lev, indexing='ij')
    lat_il,l_i = np.meshgrid(mesh.lat_i,levp1, indexing='ij')            
    lon_il,l_i = np.meshgrid(mesh.lon_i,levp1, indexing='ij')
    lat_ek,k_e = np.meshgrid(mesh.lat_e,lev, indexing='ij')

    Phis_i = Phis(lon_i, lat_i)
    ps_i = ps(lon_i, lat_i, Phis_i)

    if llm==18:
        ap_l=[0.00251499, 0.00710361, 0.01904260, 0.04607560, 0.08181860,
              0.07869805, 0.07463175, 0.06955308, 0.06339061, 0.05621774, 0.04815296,
              0.03949230, 0.03058456, 0.02193336, 0.01403670, 0.007458598, 0.002646866,
              0.0, 0.0 ]
        bp_l=[0.0, 0.0, 0.0, 0.0, 0.0, 0.03756984, 0.08652625, 0.1476709, 0.221864,
              0.308222, 0.4053179, 0.509588, 0.6168328, 0.7209891, 0.816061, 0.8952581, 
              0.953189, 0.985056, 1.0 ]
    if llm==26:
        ap_l=[0.002194067, 0.004895209, 0.009882418, 0.01805201, 0.02983724, 0.04462334, 0.06160587, 
               0.07851243, 0.07731271, 0.07590131, 0.07424086, 0.07228744, 0.06998933, 0.06728574, 0.06410509, 
               0.06036322, 0.05596111, 0.05078225, 0.04468960, 0.03752191, 0.02908949, 0.02084739, 0.01334443, 
               0.00708499, 0.00252136, 0.0, 0.0 ]
        bp_l=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01505309, 0.03276228, 0.05359622, 
                0.07810627, 0.1069411, 0.1408637, 0.1807720, 0.2277220, 0.2829562, 0.3479364, 0.4243822, 
                0.5143168, 0.6201202, 0.7235355, 0.8176768, 0.8962153, 0.9534761, 0.9851122, 1.0 ]
    if llm==49:
        ap_l=[0.002251865, 0.003983890, 0.006704364, 0.01073231, 0.01634233, 0.02367119, 
              0.03261456, 0.04274527, 0.05382610, 0.06512175, 0.07569850, 0.08454283, 
              0.08396310, 0.08334103, 0.08267352, 0.08195725, 0.08118866, 0.08036393, 
              0.07947895, 0.07852934, 0.07751036, 0.07641695, 0.07524368, 0.07398470, 
              0.07263375, 0.07118414, 0.06962863, 0.06795950, 0.06616846, 0.06424658, 
              0.06218433, 0.05997144, 0.05759690, 0.05504892, 0.05231483, 0.04938102, 
              0.04623292, 0.04285487, 0.03923006, 0.03534049, 0.03116681, 0.02668825, 
              0.02188257, 0.01676371, 0.01208171, 0.007959612, 0.004510297, 0.001831215, 
              0.0, 0.0 ]
        bp_l=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 
              0.006755112, 0.01400364, 0.02178164, 0.03012778, 0.03908356, 0.04869352, 
              0.05900542, 0.07007056, 0.08194394, 0.09468459, 0.1083559, 0.1230258, 
              0.1387673, 0.1556586, 0.1737837, 0.1932327, 0.2141024, 0.2364965, 
              0.2605264, 0.2863115, 0.3139801, 0.3436697, 0.3755280, 0.4097133, 
              0.4463958, 0.4857576, 0.5279946, 0.5733168, 0.6219495, 0.6741346, 
              0.7301315, 0.7897776, 0.8443334, 0.8923650, 0.9325572, 0.9637744, 
              0.9851122, 1.0]

    ap_l, bp_l = ref_press*np.asarray(ap_l[-1::-1]), bp_l[-1::-1] # reverse indices
    ptop = ap_l[-1] # pressure BC
    
    if mpi_rank==0: print ptop, ap_l, bp_l
    ps_il,ap_il = np.meshgrid(ps_i,ap_l, indexing='ij')
    ps_il,bp_il = np.meshgrid(ps_i,bp_l, indexing='ij')
    hybrid_coefs = meshes.mass_coefs_from_pressure_coefs(g, ap_il, bp_il)
    pb_il = ap_il + bp_il*ps_il
    mass_ik, pb_ik = mesh.field_mass(), mesh.field_mass()
    for l in range(llm):
        pb_ik[:,l]=.5*(pb_il[:,l]+pb_il[:,l+1])
        mass_ik[:,l]=(pb_il[:,l]-pb_il[:,l+1])/g
    Tb_ik = T_0 + 0.*pb_ik
    gas = thermo.set_pT(pb_ik,Tb_ik)
    Sik  = gas.s*mass_ik
# start at hydrostatic geopotential
    Phi_il = mesh.field_w()
    Phi_il[:,0]=Phis_i
    for l in range(llm):
        Phi_il[:,l+1] = Phi_il[:,l] + g*mass_ik[:,l]*gas.v[:,l]

    ujk, Wil = mesh.ucov3D(ulon(lat_ek),ulat(lat_ek)), mesh.field_w()
    
    if mpi_rank==0: 
        print 'ztop (m) = ', Phi_il[0,-1]/g
        print 'ptop (Pa) = ', gas.p[0,-1], ptop
    dx=mesh.de.min()
    params=dyn.Struct()
    params.g, params.ptop = g, ptop
    params.dx, params.dx_g0 = dx, dx/g
    params.pbot, params.rho_bot = 1e5+0.*mesh.lat_i, 1e6+0.*mesh.lat_i
    return thermo, mesh, hybrid_coefs, params, (mass_ik,Sik,ujk,Phi_il,Wil), gas

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

unst.setvar('dynamico_mpi_rank', mpi_rank)

parser = argparse.ArgumentParser()
parser.add_argument("-r", "--refinement", type=int, 
                    default=5, choices=[4, 5, 6, 7],
                    help="grid refinement level")
parser.add_argument("-l", "--llm", type=int, 
                    default=49, choices=[18, 26, 49],
                    help="number of vertical levels")
args = parser.parse_args()
nqtot, llm, grid = 1, args.llm, 2+10*(4**args.refinement)
#nqtot, llm, grid = 1,26,40962

T, Nslice, courant = 14400, 24, 3.0
caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag
#caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass

thermo, mesh, hybrid_coefs, params, flow0, gas0 = DCMIP2008c5(grid,llm)
llm, dx = mesh.llm, params.dx
if mpi_rank==0: 
    print 'grid, llm, local_gridsize, dx =', grid, llm, mesh.Ai.size, dx
    if caldyn_eta == unst.eta_lag:
        print 'Lagrangian coordinate.'
    else:
        print 'Mass-based coordinate.'

unst.ker.dynamico_init_hybrid(*hybrid_coefs)
    
dt = courant*.5*dx/np.sqrt(gas0.c2.max())

if False:
    nt = int(math.ceil(T/dt))
    dt = T/nt
else:
    nt=100
if mpi_rank==0: print 'Time step : %g s' % dt

#mesh.plot_e(mesh.le/mesh.de) ; plt.title('le/de')
#plt.savefig('fig_DCMIP2008c5/le_de.png'); plt.close()

#mesh.plot_i(mesh.Ai) ; plt.title('Ai')
#plt.savefig('fig_DCMIP2008c5/Ai.png'); plt.close()

scheme = time_step.ARK2(None, dt, a32=0.7)
caldyn_step = unst.caldyn_step_HPE(mesh,scheme,nt, caldyn_thermo,caldyn_eta, thermo,params,params.g)

def next_flow(m,S,u):
    caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
#    caldyn_step.remap()
    caldyn_step.next()
    return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), 
            caldyn_step.u.copy(), caldyn_step.geopot.copy())

def plots(it):
    s=S/m
    for l in range(llm):
        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
        vol[:,l]=(Phi[:,l+1]-Phi[:,l])/params.g/m[:,l] # specific volume
    gas = thermo.set_vs(vol, s)
    s=.5*(s+abs(s))
    t = (it*T)/3600
    print( 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g') )
    mesh.plot_i(gas.T[:,llm/2])
    plt.title('T at t=%dh'%(t))
    plt.savefig('fig_DCMIP2008c5/T%02d.png'%it)
    plt.close()

    mesh.plot_i(m[:,llm/2])
    plt.title('mass at t=%dh'%(t))
    plt.savefig('fig_DCMIP2008c5/m%02d.png'%it)
    plt.close()

    mesh.plot_i(Phi[:,0])
    plt.title('Surface geopotential at t=%dh'%(t))
    plt.savefig('fig_DCMIP2008c5/Phis%02d.png'%it)
    plt.close()

z, vol = mesh.field_mass(), mesh.field_mass()
m,S,u,Phi,W = flow0 
caldyn_step.geopot[:,0]=Phi[:,0]
#plots(0)

#context=xios.XIOS_Context(mesh.pmesh,mesh,nqtot, T)

for it in range(Nslice):
#    context.update_calendar(it)
    unst.setvar('debug_hevi_solver',False)
    time1, elapsed1 =time.time(), unst.getvar('elapsed')
    m,S,u,Phi = next_flow(m,S,u)
    time2, elapsed2 = time.time(), unst.getvar('elapsed')
    if mpi_rank==0: 
        factor = 1000./nt
        print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
        factor = 1e9/(4*nt*m.size)
        print 'nanosec per gridpoint per call to caldyn_hevi : ', factor*(time2-time1), factor*(elapsed2-elapsed1)
    
    if False:
        s=S/m
        s=.5*(s+abs(s))
        for l in range(llm):
            z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g
            vol[:,l]=(Phi[:,l+1]-Phi[:,l])/params.g/m[:,l] # specific volume
        gas = thermo.set_vs(vol, s)
        ss = np.asarray(gas.T, dtype=np.double)
        #    context.send_field_primal('theta', ss)
        #plots(it+1)

unst.ker.dynamico_print_trace()
 
#print 'xios.context_finalize()'
#context.finalize()
#print 'xios.finalize()'
#xios.finalize()
print 'MPI Rank %d Done'%mpi_rank