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' ) logging = getargs.Logger('main') args = getargs.parse() 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 maps from dynamico import xios from dynamico import precision as prec from dynamico.meshes import Cartesian_mesh as Mesh 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) 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) 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))) 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 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) print('ztop (m) = ', Phi_il[0,-1]/g, ztop) ptop = p(eta(1.)) print( 'ptop (Pa) = ', gas.p[0,-1], ptop) print('ztop(ptop) according to Eq. 7:', 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]) 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,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*params.g*(W[:,l+1]+W[:,l])/m[:,l] z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.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): logging.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() logging.info('%d/%d starting'%(mpi_rank,mpi_size)) g, Lx, Ly = 9.80616, 4e7, 6e6 nx, ny, llm = args.nx, args.ny, args.llm dx = Lx/nx unst.setvar('g',g) 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 logging.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() ) logging.debug('mask min/max :') logging.debug('%f'% davies.beta_i.min()) logging.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('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) 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************')