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 | from mpi4py import MPI |
---|
11 | comm = MPI.COMM_WORLD |
---|
12 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
13 | print '%d/%d starting'%(mpi_rank,mpi_size) |
---|
14 | |
---|
15 | import math as math |
---|
16 | import numpy as np |
---|
17 | import time |
---|
18 | |
---|
19 | from numpy import pi, log, exp, sin, cos |
---|
20 | |
---|
21 | # Baroclinic instability test based on Ullrich et al. 2015, QJRMS |
---|
22 | |
---|
23 | def baroclinic_3D(Lx,nx,Ly,ny,llm,ztop=25000.): |
---|
24 | Rd = 287.0 # Gas constant for dryy air (j kg^-1 K^-1) |
---|
25 | T0 = 288.0 # Reference temperature (K) |
---|
26 | lap = 0.005 # Lapse rate (K m^-1) |
---|
27 | b = 2. # Non dimensional vertical width parameter |
---|
28 | u0 = 35. # Reference zonal wind speed (m s^-1) |
---|
29 | a = 6.371229e6 # Radius of the Earth (m) |
---|
30 | ptop = 2000. |
---|
31 | y0 = Ly*0.5 |
---|
32 | Cpd = 1004.5 |
---|
33 | p0 = 1e5 |
---|
34 | |
---|
35 | omega = 7.292e-5 # Angular velocity of the Earth (s^-1) |
---|
36 | phi0 = 45. # Reference latitude North pi/4 (deg) |
---|
37 | f0 = 2*omega*np.sin(phi0) |
---|
38 | beta0 = 2*omega*np.cos(phi0)/a |
---|
39 | fb = 2*omega*np.sin(phi0) - y0*2*omega*np.cos(phi0)/a |
---|
40 | |
---|
41 | def Phi_xy(x,y): |
---|
42 | fc = y*y - (Ly*y/pi)*sin(2*pi*y/Ly) |
---|
43 | fd = Ly*Ly/(2*pi*pi)*cos(2*pi*y/Ly) |
---|
44 | 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)) ) |
---|
45 | |
---|
46 | 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)) |
---|
47 | def ulon(x,y,eta): return -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/b/b)) |
---|
48 | def tmean(eta) : return T0*eta**(Rd*lap/g) |
---|
49 | 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)) |
---|
50 | def p(eta): return p0*eta # eta = p/p_s |
---|
51 | |
---|
52 | def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels |
---|
53 | |
---|
54 | filename = 'cart_%03d_%03d.nc'%(nx,ny) |
---|
55 | print 'Reading Cartesian mesh ...' |
---|
56 | def coriolis(lon,lat): return f0+0.*lon |
---|
57 | meshfile = meshes.DYNAMICO_Format(filename) |
---|
58 | nqdyn, radius = 1, None |
---|
59 | # mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, coriolis) |
---|
60 | pmesh = meshes.Unstructured_PMesh(comm,meshfile) |
---|
61 | mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis) |
---|
62 | |
---|
63 | alpha_k = (np.arange(llm) +.5)/llm |
---|
64 | alpha_l = (np.arange(llm+1)+ 0.)/llm |
---|
65 | x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij') |
---|
66 | y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij') |
---|
67 | x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij') |
---|
68 | y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij') |
---|
69 | x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij') |
---|
70 | y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij') |
---|
71 | |
---|
72 | print('----------------') |
---|
73 | print 'ztop(ptop) according to Eq. 7:', T0/lap*(1.-(ptop/p0)**(Rd*lap/g)) |
---|
74 | print(np.shape(alpha_k),np.shape(alpha_l)) |
---|
75 | print(mesh.__dict__.keys()) |
---|
76 | thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) |
---|
77 | |
---|
78 | eta_il = eta(alpha_il) |
---|
79 | eta_ik = eta(alpha_ik) |
---|
80 | eta_ek = eta(alpha_ek) |
---|
81 | print('min max eta_il', np.min(eta_il),np.max(eta_il)) |
---|
82 | |
---|
83 | Phi_il = Phi_xyeta(x_il, y_il, eta_il) |
---|
84 | Phi_ik = Phi_xyeta(x_ik, y_ik, eta_ik) |
---|
85 | p_ik = p(eta_ik) |
---|
86 | T_ik = T(x_ik, y_ik, eta_ik) #ik full level(40), il(41) |
---|
87 | |
---|
88 | gas = thermo.set_pT(p_ik,T_ik) |
---|
89 | mass_ik = mesh.field_mass() |
---|
90 | for l in range(llm): |
---|
91 | mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l]) |
---|
92 | Sik, ujk, Wil = gas.s*mass_ik, mesh.field_u(), mesh.field_w() |
---|
93 | |
---|
94 | print(np.shape(ujk)) |
---|
95 | print('P_ik',p_ik[0,:]) |
---|
96 | |
---|
97 | u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek) |
---|
98 | |
---|
99 | print(np.shape(u_ek)) |
---|
100 | |
---|
101 | print 'ztop (m) = ', Phi_il[0,-1]/g, ztop |
---|
102 | ptop = p(eta(1.)) |
---|
103 | print 'ptop (Pa) = ', gas.p[0,-1], ptop |
---|
104 | |
---|
105 | params=dyn.Struct() |
---|
106 | params.ptop=ptop |
---|
107 | params.dx=dx |
---|
108 | params.dx_g0=dx/g |
---|
109 | params.g = g |
---|
110 | |
---|
111 | # define parameters for lower BC |
---|
112 | pbot = p(eta_il[:,0]) |
---|
113 | print 'min p, T :', pbot.min(), tmean(pbot/p0) |
---|
114 | gas_bot = thermo.set_pT(pbot, tmean(pbot/p0)) |
---|
115 | params.pbot = gas_bot.p |
---|
116 | params.rho_bot = 1e6/gas_bot.v |
---|
117 | |
---|
118 | return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas |
---|
119 | |
---|
120 | |
---|
121 | g, Lx, Ly = 9.81, 4e7, 6e6 |
---|
122 | nx, ny, llm = 200, 30, 22 |
---|
123 | dx,dy=Lx/nx,Ly/ny |
---|
124 | |
---|
125 | unst.setvar('g',g) |
---|
126 | |
---|
127 | thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm) |
---|
128 | |
---|
129 | T, nslice, dt = 3600., 1, 3600. |
---|
130 | |
---|
131 | context = xios.XIOS_Context_Curvilinear(mesh,1, dt) |
---|
132 | |
---|
133 | for i in range(48): |
---|
134 | context.update_calendar(i) |
---|
135 | print 'send_field', i, gas0.T.shape |
---|
136 | # context.send_field_primal('ps', lat_i) |
---|
137 | context.send_field_primal('temp', gas0.T) |
---|
138 | context.send_field_primal('p', gas0.p) |
---|
139 | |
---|
140 | # to do at the end |
---|
141 | print 'xios.context_finalize()' |
---|
142 | context.finalize() |
---|
143 | print 'xios.finalize()' |
---|
144 | xios.finalize() |
---|
145 | print 'Done' |
---|
146 | |
---|
147 | |
---|
148 | |
---|
149 | # compute hybrid coefs from initial distribution of mass |
---|
150 | mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0]) |
---|
151 | print 'Type of mass_bl, mass_dak, mass_dbk : ', [x.dtype for x in mass_bl, mass_dak, mass_dbk] |
---|
152 | unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk) |
---|
153 | |
---|
154 | dz = flow0[3].max()/(params.g*llm) |
---|
155 | #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2)) |
---|
156 | #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2)) |
---|
157 | nt = int(math.ceil(T/dt)) |
---|
158 | dt = T/nt |
---|
159 | print 'Time step : %d x %g s' % (nt,dt) |
---|
160 | |
---|
161 | #caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass |
---|
162 | caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass |
---|
163 | #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag |
---|
164 | |
---|
165 | if False: # time stepping in Python |
---|
166 | caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g) |
---|
167 | scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt) |
---|
168 | def next_flow(m,S,u,Phi,W): |
---|
169 | # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) |
---|
170 | return scheme.advance((m,S,u,Phi,W),nt) |
---|
171 | |
---|
172 | else: # time stepping in Fortran |
---|
173 | scheme = time_step.ARK2(None, dt) |
---|
174 | caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta, |
---|
175 | thermo,params,params.g) |
---|
176 | def next_flow(m,S,u,Phi,W): |
---|
177 | # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) |
---|
178 | caldyn_step.mass[:,:,:], caldyn_step.theta_rhodz[:,:,:], caldyn_step.u[:,:,:] = m,S,u |
---|
179 | caldyn_step.geopot[:,:,:], caldyn_step.W[:,:,:] = Phi,W |
---|
180 | caldyn_step.next() |
---|
181 | return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(), |
---|
182 | caldyn_step.geopot.copy(), caldyn_step.W.copy()) |
---|
183 | |
---|
184 | m,S,u,Phi,W=flow0 |
---|
185 | if caldyn_thermo == unst.thermo_theta: |
---|
186 | s=S/m |
---|
187 | theta = thermo.T0*np.exp(s/thermo.Cpd) |
---|
188 | S=m*theta |
---|
189 | title_format = 'Potential temperature at t=%g s (K)' |
---|
190 | else: |
---|
191 | title_format = 'Specific entropy at t=%g s (J/K/kg)' |
---|
192 | |
---|
193 | w=mesh.field_mass() |
---|
194 | z=mesh.field_mass() |
---|
195 | |
---|
196 | Nslice=0 |
---|
197 | for it in range(Nslice): |
---|
198 | s=S/m ; s=.5*(s+abs(s)) |
---|
199 | for l in range(llm): |
---|
200 | w[:,:,l]=.5*params.g*(W[:,:,l+1]+W[:,:,l])/m[:,:,l] |
---|
201 | z[:,:,l]=.5*(Phi[:,:,l+1]+Phi[:,:,l])/params.g |
---|
202 | |
---|
203 | print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g') |
---|
204 | jj=ny/2 |
---|
205 | xx,zz,ss,ww = mesh.xx[:,jj,:]/1000., z[:,jj,:]/1000., s[:,jj,:], w[:,jj,:] |
---|
206 | |
---|
207 | f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4)) |
---|
208 | |
---|
209 | c=ax1.contourf(xx,zz,ss,20) |
---|
210 | ax1.set_xlim((-.5,.5)), ax1.set_xlabel('x (km)') |
---|
211 | ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)') |
---|
212 | plt.colorbar(c,ax=ax1) |
---|
213 | ax1.set_title(title_format % (it*T,)) |
---|
214 | |
---|
215 | # plt.show() |
---|
216 | |
---|
217 | # plt.figure(figsize=(12,5)) |
---|
218 | c=ax2.contourf(xx,zz,ww,20) |
---|
219 | ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)') |
---|
220 | ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)') |
---|
221 | plt.colorbar(c,ax=ax2) |
---|
222 | ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,)) |
---|
223 | # plt.tight_layout() |
---|
224 | # plt.show() |
---|
225 | plt.savefig('fig_baroclinic_3D/%02d.png'%it) |
---|
226 | |
---|
227 | time1, elapsed1 =time.time(), unst.getvar('elapsed') |
---|
228 | m,S,u,Phi,W = next_flow(m,S,u,Phi,W) |
---|
229 | time2, elapsed2 =time.time(), unst.getvar('elapsed') |
---|
230 | factor = 1000./nt |
---|
231 | print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) |
---|
232 | factor = 1e9/(4*nt*nx*ny*llm) |
---|
233 | print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) |
---|
234 | |
---|