[790] | 1 | from __future__ import print_function |
---|
[795] | 2 | from __future__ import division |
---|
| 3 | |
---|
[790] | 4 | from dynamico import getargs |
---|
| 5 | |
---|
| 6 | getargs.add("--T", type=float, default=3600., |
---|
| 7 | help="Length of time slice in seconds") |
---|
| 8 | getargs.add("--perturb", type=float, default=1., |
---|
| 9 | help="Amplitude of wind perturbation in m/s") |
---|
| 10 | getargs.add("--N", type=int, default=48, |
---|
| 11 | help="Number of time slices") |
---|
| 12 | getargs.add("--Davies_N1", type=int, default=3) |
---|
| 13 | getargs.add("--Davies_N2", type=int, default=3) |
---|
| 14 | getargs.add("--nx", type=int, default=200) |
---|
| 15 | getargs.add("--ny", type=int, default=30) |
---|
| 16 | getargs.add("--betaplane", action='store_true') |
---|
[795] | 17 | |
---|
| 18 | getargs.add("--kappa_divgrad", type=float, default=3.0e15) |
---|
| 19 | getargs.add("--kappa_curlcurl", type=float, default=3.0e15) |
---|
[797] | 20 | getargs.add("--ztop", type=float, default=30000.) |
---|
[795] | 21 | |
---|
[790] | 22 | getargs.defaults(dt=360., llm=50) |
---|
| 23 | |
---|
[802] | 24 | #getargs.config_log(filename='out.log',filemode='w') # must be done before calling Logger() |
---|
[795] | 25 | # getargs.config_log(filename='out.log',filemode='w', |
---|
| 26 | # format='%(processName)s:%(name)s:%(filename)s:%(module)s:%(funcName)s:%(lineno)d:%(levelname)s:%(message)s' ) |
---|
| 27 | |
---|
[790] | 28 | args = getargs.parse() |
---|
[802] | 29 | log_master, log_world = getargs.getLogger() |
---|
| 30 | INFO, DEBUG, DEBUG_ALL, ERROR = log_master.info, log_master.debug, log_world.debug, log_world.error |
---|
[790] | 31 | |
---|
[761] | 32 | from dynamico import unstructured as unst |
---|
| 33 | from dynamico import dyn |
---|
| 34 | from dynamico import time_step |
---|
| 35 | from dynamico import DCMIP |
---|
| 36 | from dynamico import meshes |
---|
[800] | 37 | from dynamico import maps |
---|
[761] | 38 | from dynamico import xios |
---|
| 39 | from dynamico import precision as prec |
---|
| 40 | from dynamico.meshes import Cartesian_mesh as Mesh |
---|
[791] | 41 | from dynamico.kernels import grad, curl, div, KE |
---|
| 42 | from dynamico.LAM import Davies |
---|
[761] | 43 | |
---|
| 44 | import math as math |
---|
| 45 | import numpy as np |
---|
| 46 | import time |
---|
| 47 | from numpy import pi, log, exp, sin, cos |
---|
| 48 | |
---|
| 49 | # Baroclinic instability test based on Ullrich et al. 2015, QJRMS |
---|
| 50 | |
---|
[794] | 51 | def create_pmesh(nx,ny): |
---|
| 52 | filename = 'cart_%03d_%03d.nc'%(nx,ny) |
---|
[802] | 53 | INFO('Reading Cartesian mesh ...') |
---|
[794] | 54 | meshfile = meshes.DYNAMICO_Format(filename) |
---|
| 55 | pmesh = meshes.Unstructured_PMesh(comm,meshfile) |
---|
| 56 | pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj) |
---|
| 57 | return pmesh |
---|
[764] | 58 | |
---|
[800] | 59 | def baroclinic_3D(pmesh, dx,Lx,Ly,llm,ztop): |
---|
[761] | 60 | Rd = 287.0 # Gas constant for dryy air (j kg^-1 K^-1) |
---|
| 61 | T0 = 288.0 # Reference temperature (K) |
---|
| 62 | lap = 0.005 # Lapse rate (K m^-1) |
---|
| 63 | b = 2. # Non dimensional vertical width parameter |
---|
| 64 | u0 = 35. # Reference zonal wind speed (m s^-1) |
---|
| 65 | a = 6.371229e6 # Radius of the Earth (m) |
---|
[787] | 66 | y0 = .5*Ly |
---|
[761] | 67 | Cpd = 1004.5 |
---|
| 68 | p0 = 1e5 |
---|
[790] | 69 | up = args.perturb # amplitude (m/s) |
---|
[774] | 70 | xc,yc,lp = 0.,Ly/2.,600000. |
---|
[761] | 71 | |
---|
| 72 | omega = 7.292e-5 # Angular velocity of the Earth (s^-1) |
---|
[790] | 73 | phi0 = 45.*pi/180.0 # Reference latitude North pi/4 (deg) |
---|
| 74 | f0 = 2*omega*sin(phi0) |
---|
| 75 | beta0 = 2*omega*cos(phi0)/a if args.betaplane else 0. |
---|
[787] | 76 | fb = f0 - beta0*y0 |
---|
[761] | 77 | |
---|
[769] | 78 | def Phi_xy(y): |
---|
[761] | 79 | fc = y*y - (Ly*y/pi)*sin(2*pi*y/Ly) |
---|
| 80 | fd = Ly*Ly/(2*pi*pi)*cos(2*pi*y/Ly) |
---|
| 81 | 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)) ) |
---|
| 82 | |
---|
[769] | 83 | def Phi_xyeta(y,eta): return T0*g/lap*(1-eta**(Rd*lap/g)) + Phi_xy(y)*log(eta)*exp(-((log(eta)/b)**2)) |
---|
[774] | 84 | def ulon(x,y,eta): |
---|
| 85 | u = -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/(b*b))) |
---|
[790] | 86 | u = u + up*exp(-(((x-xc)**2+(y-yc)**2)/lp**2)) |
---|
[774] | 87 | return u |
---|
| 88 | |
---|
[761] | 89 | def tmean(eta) : return T0*eta**(Rd*lap/g) |
---|
[769] | 90 | def T(y,eta) : return tmean(eta)+(Phi_xy(y)/Rd)*(((2/(b*b))*(log(eta))**2)-1)*exp(-((0.5*log(eta))**2)) |
---|
[761] | 91 | def p(eta): return p0*eta # eta = p/p_s |
---|
| 92 | |
---|
[787] | 93 | def eta(alpha) : return (1.-(lap*ztop*alpha/T0))**(g/(Rd*lap)) # roughly equispaced levels |
---|
[800] | 94 | |
---|
| 95 | betaplane = maps.BetaPlaneMap(dx, .5*Lx, .5*Ly, f0, beta0, .5*Ly) |
---|
| 96 | nqdyn = 1 |
---|
| 97 | mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, betaplane) |
---|
[761] | 98 | |
---|
| 99 | alpha_k = (np.arange(llm) +.5)/llm |
---|
| 100 | alpha_l = (np.arange(llm+1)+ 0.)/llm |
---|
[800] | 101 | x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij') |
---|
| 102 | y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij') |
---|
| 103 | x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij') |
---|
| 104 | y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij') |
---|
| 105 | x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij') |
---|
| 106 | y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij') |
---|
[761] | 107 | |
---|
[802] | 108 | # print(np.shape(alpha_k),np.shape(alpha_l)) |
---|
[761] | 109 | thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) |
---|
| 110 | |
---|
| 111 | eta_il = eta(alpha_il) |
---|
| 112 | eta_ik = eta(alpha_ik) |
---|
| 113 | eta_ek = eta(alpha_ek) |
---|
[790] | 114 | # print('min max eta_il', np.min(eta_il),np.max(eta_il)) |
---|
[761] | 115 | |
---|
[769] | 116 | Phi_il = Phi_xyeta(y_il, eta_il) |
---|
| 117 | Phi_ik = Phi_xyeta(y_ik, eta_ik) |
---|
[787] | 118 | p_ik, p_il = p(eta_ik), p(eta_il) |
---|
[769] | 119 | T_ik = T(y_ik, eta_ik) #ik full level(40), il(41) |
---|
[761] | 120 | |
---|
| 121 | gas = thermo.set_pT(p_ik,T_ik) |
---|
| 122 | mass_ik = mesh.field_mass() |
---|
| 123 | for l in range(llm): |
---|
| 124 | mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l]) |
---|
[787] | 125 | # mass_ik[:,l]=(p_il[:,l]-p_il[:,l+1])/g |
---|
[774] | 126 | Sik, Wil = gas.s*mass_ik, mesh.field_w() |
---|
[761] | 127 | |
---|
| 128 | u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek) |
---|
| 129 | |
---|
[802] | 130 | INFO('ztop (m) = %f %f' % (Phi_il[0,-1]/g, ztop) ) |
---|
[761] | 131 | ptop = p(eta(1.)) |
---|
[802] | 132 | INFO( 'ptop (Pa) = %f %f' % (gas.p[0,-1], ptop) ) |
---|
| 133 | INFO('ztop(ptop) according to Eq. 7: %f' % (T0/lap*(1.-(ptop/p0)**(Rd*lap/g))) ) |
---|
[761] | 134 | |
---|
| 135 | params=dyn.Struct() |
---|
| 136 | params.ptop=ptop |
---|
| 137 | params.dx=dx |
---|
| 138 | params.dx_g0=dx/g |
---|
| 139 | params.g = g |
---|
| 140 | |
---|
| 141 | # define parameters for lower BC |
---|
| 142 | pbot = p(eta_il[:,0]) |
---|
[802] | 143 | INFO( 'min p, T : %f %s' % (pbot.min(), tmean(pbot/p0)) ) |
---|
[761] | 144 | gas_bot = thermo.set_pT(pbot, tmean(pbot/p0)) |
---|
| 145 | params.pbot = gas_bot.p |
---|
| 146 | params.rho_bot = 1e6/gas_bot.v |
---|
| 147 | |
---|
[774] | 148 | return thermo, mesh, params, prec.asnum([mass_ik,Sik,u_ek,Phi_il,Wil]), gas |
---|
[761] | 149 | |
---|
[795] | 150 | def diagnose(Phi,S,m,W,u): |
---|
[780] | 151 | s=S/m |
---|
[795] | 152 | curl_vk = curl(mesh, u) |
---|
| 153 | abs_vort_vk = mesh.field_z() # absolute vorticity |
---|
| 154 | un = mesh.field_u() |
---|
| 155 | v = mesh.field_mass() # specific volume |
---|
| 156 | w = mesh.field_mass() |
---|
| 157 | z = mesh.field_mass() |
---|
[769] | 158 | for l in range(llm): |
---|
| 159 | v[:,l]=(Phi[:,l+1]-Phi[:,l])/(g*m[:,l]) |
---|
| 160 | w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l] |
---|
| 161 | z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g |
---|
[795] | 162 | un[:,l]=u[:,l]/mesh.de |
---|
| 163 | abs_vort_vk[:,l]=curl_vk[:,l] + mesh.fv |
---|
[769] | 164 | gas = thermo.set_vs(v,s) |
---|
[795] | 165 | ps = gas.p[:,0]+ .5*g*m[:,0] |
---|
| 166 | return gas, w, z, ps, un, curl_vk, abs_vort_vk |
---|
[769] | 167 | |
---|
[771] | 168 | class myDavies(Davies): |
---|
| 169 | def mask(self,x,y): |
---|
[802] | 170 | DEBUG('davies dy = %f'% dx) |
---|
[800] | 171 | return self.mask0(y,Ly,dx)*self.mask0(-y,0.,dx) |
---|
[771] | 172 | |
---|
[763] | 173 | with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator |
---|
| 174 | comm = client.comm |
---|
| 175 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
[761] | 176 | |
---|
[802] | 177 | INFO('%d/%d starting'%(mpi_rank,mpi_size)) |
---|
| 178 | if mpi_rank>0 : |
---|
| 179 | getargs.not_master() |
---|
| 180 | unst.setvar('is_mpi_master', False) |
---|
| 181 | |
---|
[795] | 182 | g, Lx, Ly = 9.80616, 4e7, 6e6 |
---|
[787] | 183 | nx, ny, llm = args.nx, args.ny, args.llm |
---|
[794] | 184 | dx = Lx/nx |
---|
| 185 | |
---|
[763] | 186 | unst.setvar('g',g) |
---|
[802] | 187 | unst.setvar('debug_hevi_solver', False) |
---|
| 188 | |
---|
[794] | 189 | pmesh = create_pmesh(nx,ny) |
---|
[800] | 190 | thermo, mesh, params, flow0, gas0 = baroclinic_3D(pmesh, dx,Lx,Ly,llm, args.ztop) |
---|
[761] | 191 | |
---|
[769] | 192 | mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0]) |
---|
| 193 | unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk) |
---|
[761] | 194 | |
---|
[774] | 195 | T = args.T |
---|
[787] | 196 | dt = args.dt |
---|
[769] | 197 | dz = flow0[3].max()/(params.g*llm) |
---|
| 198 | nt = int(math.ceil(T/dt)) |
---|
| 199 | dt = T/nt |
---|
[802] | 200 | INFO( 'Time step : %d x %g = %g s' % (nt,dt,nt*dt)) |
---|
[794] | 201 | |
---|
[769] | 202 | |
---|
| 203 | caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass |
---|
| 204 | |
---|
| 205 | if False: # time stepping in Python |
---|
| 206 | caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g) |
---|
| 207 | scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt) |
---|
| 208 | def next_flow(m,S,u,Phi,W): |
---|
| 209 | return scheme.advance((m,S,u,Phi,W),nt) |
---|
| 210 | |
---|
| 211 | else: # time stepping in Fortran |
---|
| 212 | scheme = time_step.ARK2(None, dt) |
---|
[787] | 213 | if args.hydrostatic: |
---|
[795] | 214 | caldyn_step = unst.caldyn_step_HPE(mesh,scheme,1, caldyn_thermo,caldyn_eta, |
---|
| 215 | thermo,params,params.g) |
---|
[787] | 216 | else: |
---|
[795] | 217 | caldyn_step = unst.caldyn_step_NH(mesh,scheme,1, caldyn_thermo,caldyn_eta, |
---|
| 218 | thermo,params,params.g) |
---|
| 219 | davies = myDavies(args.Davies_N1, args.Davies_N2, |
---|
| 220 | mesh.lon_i, mesh.lat_i, mesh.lon_e,mesh.lat_e) |
---|
[790] | 221 | # print('mask min/max :', davies.beta_i.min(), davies.beta_i.max() ) |
---|
[802] | 222 | DEBUG('mask min/max :') |
---|
| 223 | DEBUG('%f'% davies.beta_i.min()) |
---|
| 224 | DEBUG('%f'% davies.beta_i.max()) |
---|
[795] | 225 | |
---|
[769] | 226 | def next_flow(m,S,u,Phi,W): |
---|
| 227 | # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) |
---|
| 228 | caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u |
---|
| 229 | caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W |
---|
[771] | 230 | for i in range(nt): |
---|
| 231 | caldyn_step.next() |
---|
| 232 | davies.relax(llm, caldyn_step, flow0) |
---|
[792] | 233 | m,S,u = caldyn_step.mass, caldyn_step.theta_rhodz, caldyn_step.u |
---|
[785] | 234 | s = S/m |
---|
[792] | 235 | laps, bilaps = mesh.field_mass(), mesh.field_mass() |
---|
| 236 | lapu, bilapu = mesh.field_u(), mesh.field_u() |
---|
[785] | 237 | unst.ker.dynamico_scalar_laplacian(s,laps) |
---|
| 238 | unst.ker.dynamico_scalar_laplacian(laps,bilaps) |
---|
[792] | 239 | unst.ker.dynamico_curl_laplacian(u,lapu) |
---|
| 240 | unst.ker.dynamico_curl_laplacian(lapu,bilapu) |
---|
[795] | 241 | caldyn_step.theta_rhodz[:] = S - dt*args.kappa_divgrad*bilaps*m # Euler step |
---|
| 242 | caldyn_step.u[:] = u - dt*args.kappa_curlcurl*bilapu # Euler step |
---|
[792] | 243 | |
---|
[769] | 244 | return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(), |
---|
[771] | 245 | caldyn_step.geopot.copy(), caldyn_step.W.copy()) |
---|
[769] | 246 | |
---|
[774] | 247 | m,S,u,Phi,W = flow0 |
---|
[769] | 248 | if caldyn_thermo == unst.thermo_theta: |
---|
| 249 | s=S/m |
---|
[790] | 250 | theta = thermo.T0*exp(s/thermo.Cpd) |
---|
[769] | 251 | S=m*theta |
---|
| 252 | title_format = 'Potential temperature at t=%g s (K)' |
---|
| 253 | else: |
---|
| 254 | title_format = 'Specific entropy at t=%g s (J/K/kg)' |
---|
| 255 | |
---|
| 256 | #T, nslice, dt = 3600., 1, 3600. |
---|
[774] | 257 | Nslice = args.N |
---|
[769] | 258 | |
---|
[795] | 259 | temp_v = mesh.field_z(), |
---|
| 260 | |
---|
[769] | 261 | with xios.Context_Curvilinear(mesh,1, 24*3600) as context: |
---|
[763] | 262 | # now XIOS knows about the mesh and we can write to disk |
---|
[792] | 263 | for i in range(Nslice+1): |
---|
[785] | 264 | context.update_calendar(i+1) |
---|
[769] | 265 | |
---|
| 266 | # Diagnose quantities of interest from prognostic variables m,S,u,Phi,W |
---|
[795] | 267 | gas, w, z, ps, un, zeta_vk, zeta_abs_vk = diagnose(Phi,S,m,W,u) |
---|
[790] | 268 | KE_ik = KE(mesh,u) |
---|
[787] | 269 | du_fast, du_slow = caldyn_step.du_fast[0,:,:], caldyn_step.du_slow[0,:,:] |
---|
| 270 | div_fast, div_slow = div(mesh,du_fast), div(mesh,du_slow) |
---|
[790] | 271 | drhodz, hflux = caldyn_step.drhodz[0,:,:], caldyn_step.hflux[:,:] |
---|
| 272 | |
---|
[769] | 273 | # write to disk |
---|
[787] | 274 | context.send_field_primal('ps', ps) |
---|
[785] | 275 | context.send_field_primal('temp', gas.T) |
---|
[790] | 276 | |
---|
[769] | 277 | context.send_field_primal('p', gas.p) |
---|
[790] | 278 | # context.send_field_primal('p', KE_ik) |
---|
| 279 | # context.send_field_primal('p', drhodz) |
---|
| 280 | |
---|
[800] | 281 | context.send_field_primal('theta', thermo.T0*exp(gas.s/thermo.Cpd)) |
---|
| 282 | |
---|
[769] | 283 | context.send_field_primal('uz', w) |
---|
[780] | 284 | context.send_field_primal('z', z) |
---|
[787] | 285 | context.send_field_primal('div_fast', div_fast) |
---|
[790] | 286 | |
---|
[787] | 287 | context.send_field_primal('div_slow', div_slow) |
---|
[790] | 288 | |
---|
[795] | 289 | context.send_field_dual('curl', zeta_vk) |
---|
| 290 | context.send_field_dual('curl_abs', zeta_abs_vk) |
---|
[769] | 291 | |
---|
[802] | 292 | INFO( 'ptop, model top (m) : %f %f' % (unst.getvar('ptop'), Phi.max()/unst.getvar('g')) ) |
---|
[769] | 293 | |
---|
| 294 | time1, elapsed1 =time.time(), unst.getvar('elapsed') |
---|
| 295 | m,S,u,Phi,W = next_flow(m,S,u,Phi,W) |
---|
| 296 | time2, elapsed2 =time.time(), unst.getvar('elapsed') |
---|
| 297 | factor = 1000./nt |
---|
[802] | 298 | INFO( 'ms per full time step : %f %f' %(factor*(time2-time1), factor*(elapsed2-elapsed1)) ) |
---|
[769] | 299 | factor = 1e9/(4*nt*nx*ny*llm) |
---|
[802] | 300 | INFO( 'nanosec per gridpoint per full time step : %f %f' % |
---|
| 301 | (factor*(time2-time1), factor*(elapsed2-elapsed1)) ) |
---|
[769] | 302 | |
---|
[802] | 303 | INFO('************DONE************') |
---|