[761] | 1 | from dynamico import unstructured as unst |
---|
| 2 | from dynamico import dyn |
---|
| 3 | from dynamico import time_step |
---|
| 4 | from dynamico import DCMIP |
---|
| 5 | from dynamico import meshes |
---|
| 6 | from dynamico import xios |
---|
| 7 | from dynamico import precision as prec |
---|
| 8 | from dynamico.meshes import Cartesian_mesh as Mesh |
---|
| 9 | |
---|
| 10 | import math as math |
---|
| 11 | import numpy as np |
---|
| 12 | import time |
---|
[764] | 13 | import argparse |
---|
[761] | 14 | from numpy import pi, log, exp, sin, cos |
---|
| 15 | |
---|
| 16 | # Baroclinic instability test based on Ullrich et al. 2015, QJRMS |
---|
| 17 | |
---|
[764] | 18 | parser = argparse.ArgumentParser() |
---|
| 19 | |
---|
| 20 | parser.add_argument("-mpi_ni", type=int, |
---|
| 21 | default=64, choices=None, |
---|
| 22 | help="number of x processors") |
---|
| 23 | parser.add_argument("-mpi_nj", type=int, |
---|
| 24 | default=64, choices=None, |
---|
| 25 | help="number of y processors") |
---|
| 26 | args = parser.parse_args() |
---|
| 27 | |
---|
| 28 | |
---|
[761] | 29 | def baroclinic_3D(Lx,nx,Ly,ny,llm,ztop=25000.): |
---|
| 30 | Rd = 287.0 # Gas constant for dryy air (j kg^-1 K^-1) |
---|
| 31 | T0 = 288.0 # Reference temperature (K) |
---|
| 32 | lap = 0.005 # Lapse rate (K m^-1) |
---|
| 33 | b = 2. # Non dimensional vertical width parameter |
---|
| 34 | u0 = 35. # Reference zonal wind speed (m s^-1) |
---|
| 35 | a = 6.371229e6 # Radius of the Earth (m) |
---|
| 36 | ptop = 2000. |
---|
| 37 | y0 = Ly*0.5 |
---|
| 38 | Cpd = 1004.5 |
---|
| 39 | p0 = 1e5 |
---|
| 40 | |
---|
| 41 | omega = 7.292e-5 # Angular velocity of the Earth (s^-1) |
---|
| 42 | phi0 = 45. # Reference latitude North pi/4 (deg) |
---|
| 43 | f0 = 2*omega*np.sin(phi0) |
---|
| 44 | beta0 = 2*omega*np.cos(phi0)/a |
---|
| 45 | fb = 2*omega*np.sin(phi0) - y0*2*omega*np.cos(phi0)/a |
---|
| 46 | |
---|
| 47 | def Phi_xy(x,y): |
---|
| 48 | fc = y*y - (Ly*y/pi)*sin(2*pi*y/Ly) |
---|
| 49 | fd = Ly*Ly/(2*pi*pi)*cos(2*pi*y/Ly) |
---|
| 50 | 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)) ) |
---|
| 51 | |
---|
| 52 | def Phi_xyeta(x,y,eta): return T0*g/lap*(1-eta**(Rd*lap/g)) + Phi_xy(x,y)*log(eta)*exp(-((log(eta)/b)**2)) |
---|
| 53 | def ulon(x,y,eta): return -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/b/b)) |
---|
| 54 | def tmean(eta) : return T0*eta**(Rd*lap/g) |
---|
| 55 | def T(x,y,eta) : return tmean(eta)+(Phi_xy(x,y)/Rd)*(((2/(b*b))*(log(eta))**2)-1)*exp(-((0.5*log(eta))**2)) |
---|
| 56 | def p(eta): return p0*eta # eta = p/p_s |
---|
| 57 | |
---|
| 58 | def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels |
---|
| 59 | |
---|
| 60 | filename = 'cart_%03d_%03d.nc'%(nx,ny) |
---|
| 61 | print 'Reading Cartesian mesh ...' |
---|
| 62 | def coriolis(lon,lat): return f0+0.*lon |
---|
| 63 | meshfile = meshes.DYNAMICO_Format(filename) |
---|
| 64 | nqdyn, radius = 1, None |
---|
| 65 | pmesh = meshes.Unstructured_PMesh(comm,meshfile) |
---|
[763] | 66 | pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj) |
---|
[761] | 67 | mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis) |
---|
| 68 | |
---|
| 69 | alpha_k = (np.arange(llm) +.5)/llm |
---|
| 70 | alpha_l = (np.arange(llm+1)+ 0.)/llm |
---|
| 71 | x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij') |
---|
| 72 | y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij') |
---|
| 73 | x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij') |
---|
| 74 | y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij') |
---|
| 75 | x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij') |
---|
| 76 | y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij') |
---|
| 77 | |
---|
| 78 | print('----------------') |
---|
| 79 | print 'ztop(ptop) according to Eq. 7:', T0/lap*(1.-(ptop/p0)**(Rd*lap/g)) |
---|
| 80 | print(np.shape(alpha_k),np.shape(alpha_l)) |
---|
| 81 | print(mesh.__dict__.keys()) |
---|
| 82 | thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) |
---|
| 83 | |
---|
| 84 | eta_il = eta(alpha_il) |
---|
| 85 | eta_ik = eta(alpha_ik) |
---|
| 86 | eta_ek = eta(alpha_ek) |
---|
| 87 | print('min max eta_il', np.min(eta_il),np.max(eta_il)) |
---|
| 88 | |
---|
| 89 | Phi_il = Phi_xyeta(x_il, y_il, eta_il) |
---|
| 90 | Phi_ik = Phi_xyeta(x_ik, y_ik, eta_ik) |
---|
| 91 | p_ik = p(eta_ik) |
---|
| 92 | T_ik = T(x_ik, y_ik, eta_ik) #ik full level(40), il(41) |
---|
| 93 | |
---|
| 94 | gas = thermo.set_pT(p_ik,T_ik) |
---|
| 95 | mass_ik = mesh.field_mass() |
---|
| 96 | for l in range(llm): |
---|
| 97 | mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l]) |
---|
| 98 | Sik, ujk, Wil = gas.s*mass_ik, mesh.field_u(), mesh.field_w() |
---|
| 99 | |
---|
| 100 | print(np.shape(ujk)) |
---|
| 101 | print('P_ik',p_ik[0,:]) |
---|
| 102 | |
---|
| 103 | u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek) |
---|
| 104 | |
---|
| 105 | print(np.shape(u_ek)) |
---|
| 106 | |
---|
| 107 | print 'ztop (m) = ', Phi_il[0,-1]/g, ztop |
---|
| 108 | ptop = p(eta(1.)) |
---|
| 109 | print 'ptop (Pa) = ', gas.p[0,-1], ptop |
---|
| 110 | |
---|
| 111 | params=dyn.Struct() |
---|
| 112 | params.ptop=ptop |
---|
| 113 | params.dx=dx |
---|
| 114 | params.dx_g0=dx/g |
---|
| 115 | params.g = g |
---|
| 116 | |
---|
| 117 | # define parameters for lower BC |
---|
| 118 | pbot = p(eta_il[:,0]) |
---|
| 119 | print 'min p, T :', pbot.min(), tmean(pbot/p0) |
---|
| 120 | gas_bot = thermo.set_pT(pbot, tmean(pbot/p0)) |
---|
| 121 | params.pbot = gas_bot.p |
---|
| 122 | params.rho_bot = 1e6/gas_bot.v |
---|
| 123 | |
---|
| 124 | return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas |
---|
| 125 | |
---|
[763] | 126 | with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator |
---|
| 127 | comm = client.comm |
---|
| 128 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
| 129 | print '%d/%d starting'%(mpi_rank,mpi_size) |
---|
[761] | 130 | |
---|
[763] | 131 | g, Lx, Ly = 9.81, 4e7, 6e6 |
---|
| 132 | nx, ny, llm = 200, 30, 22 |
---|
| 133 | dx,dy=Lx/nx,Ly/ny |
---|
[761] | 134 | |
---|
[763] | 135 | unst.setvar('g',g) |
---|
[761] | 136 | |
---|
[763] | 137 | thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm) |
---|
[761] | 138 | |
---|
[763] | 139 | T, nslice, dt = 3600., 1, 3600. |
---|
[761] | 140 | |
---|
[763] | 141 | with xios.Context_Curvilinear(mesh,1, dt) as context: |
---|
| 142 | # now XIOS knows about the mesh and we can write to disk |
---|
| 143 | for i in range(48): |
---|
| 144 | context.update_calendar(i) |
---|
| 145 | print 'send_field', i, gas0.T.shape |
---|
| 146 | # context.send_field_primal('ps', lat_i) |
---|
| 147 | context.send_field_primal('temp', gas0.T) |
---|
| 148 | context.send_field_primal('p', gas0.p) |
---|