[766] | 1 | from dynamico import unstructured as unst |
---|
| 2 | from dynamico import dyn |
---|
| 3 | from dynamico import time_step |
---|
| 4 | from dynamico import meshes |
---|
| 5 | from dynamico import xios |
---|
| 6 | from dynamico import precision as prec |
---|
| 7 | import math as math |
---|
| 8 | import matplotlib.pyplot as plt |
---|
| 9 | import numpy as np |
---|
| 10 | import time |
---|
| 11 | import argparse |
---|
| 12 | |
---|
| 13 | def thermal_bubble_3D(Lx,nx,Ly,ny,llm,ztop=1000., zc=350., |
---|
| 14 | rc=250, thetac=0.5, x0=0., y0=0.): |
---|
| 15 | Cpd, Rd, g, p0,theta0, T0 = 1004.5, 287.,9.81, 1e5, 300., 300. |
---|
| 16 | nqdyn = 1 |
---|
| 17 | |
---|
| 18 | Phi = lambda eta : g*ztop*eta |
---|
| 19 | p = lambda Phi : p0*np.exp(-Phi/(Rd*T0)) |
---|
| 20 | zz = lambda p: -(Rd*T0*np.log(p/p0))/g |
---|
| 21 | rr = lambda x,y,p: np.sqrt((x-x0)**2 + (y-y0)**2 + (zz(p)-zc)**2) |
---|
| 22 | sa = lambda x,y,p: rr(x,y,p) < rc |
---|
| 23 | deform = lambda x,y,p: (0.5*thetac*(1+np.cos(np.pi*rr(x,y,p)/rc)))*sa(x,y,p) |
---|
| 24 | temp = lambda p: theta0*(p/p0)**(Rd/Cpd) |
---|
| 25 | T = lambda x,y,p: deform(x,y,p) + temp(p) |
---|
| 26 | |
---|
| 27 | alpha_k = (np.arange(llm) +.5)/llm |
---|
| 28 | alpha_l = (np.arange(llm+1)+ 0.)/llm |
---|
| 29 | x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij') |
---|
| 30 | y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij') |
---|
| 31 | x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij') |
---|
| 32 | y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij') |
---|
| 33 | x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij') |
---|
| 34 | y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij') |
---|
| 35 | |
---|
| 36 | thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) |
---|
| 37 | |
---|
[768] | 38 | Phi_il = Phi(alpha_il) |
---|
| 39 | Phi_ik = Phi(alpha_ik) |
---|
[766] | 40 | p_ik = p(Phi_ik) |
---|
| 41 | T_ik = T(x_ik, y_ik, p_ik) |
---|
| 42 | |
---|
| 43 | gas = thermo.set_pT(p_ik,T_ik) |
---|
| 44 | mass_ik = mesh.field_mass() |
---|
| 45 | for l in range(llm): |
---|
| 46 | mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l]) |
---|
| 47 | Sik, ujk, Wil = gas.s*mass_ik, mesh.field_u(), mesh.field_w() |
---|
| 48 | |
---|
| 49 | print 'ztop (m) = ', Phi_il[0,-1]/g, ztop |
---|
| 50 | ptop = p(g*ztop) |
---|
| 51 | print 'ptop (Pa) = ', gas.p[0,-1], ptop |
---|
| 52 | params=dyn.Struct() |
---|
| 53 | params.ptop=ptop |
---|
| 54 | params.dx=dx |
---|
| 55 | params.dx_g0=dx/g |
---|
| 56 | params.g = g |
---|
[768] | 57 | |
---|
[766] | 58 | # define parameters for lower BC |
---|
[768] | 59 | pbot = p(alpha_il[0]) |
---|
[766] | 60 | print 'min p, T :', pbot.min(), temp(pbot/p0) |
---|
| 61 | gas_bot = thermo.set_pT(pbot, temp(pbot/p0)) |
---|
| 62 | params.pbot = gas_bot.p |
---|
| 63 | params.rho_bot = 1e6/gas_bot.v |
---|
| 64 | |
---|
| 65 | return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas |
---|
| 66 | |
---|
[768] | 67 | def diagnose(Phi,S,m,W): |
---|
| 68 | s=S/m ; s=.5*(s+abs(s)) |
---|
| 69 | for l in range(llm): |
---|
| 70 | v[:,l]=(Phi[:,l+1]-Phi[:,l])/(g*m[:,l]) |
---|
| 71 | w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l] |
---|
| 72 | z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g |
---|
| 73 | gas = thermo.set_vs(v,s) |
---|
| 74 | return gas, w, z |
---|
| 75 | |
---|
| 76 | def reshape(data): return data.reshape((nx,ny)) |
---|
| 77 | |
---|
| 78 | def plot(): |
---|
| 79 | x, y = map(reshape, (xx,yy) ) |
---|
| 80 | zz=np.zeros((nx,ny,llm)) |
---|
| 81 | ss=np.zeros((nx,ny,llm)) |
---|
| 82 | ww=np.zeros((nx,ny,llm)) |
---|
| 83 | x3d=np.zeros((nx,ny,llm)) |
---|
| 84 | |
---|
| 85 | for l in range(llm): |
---|
| 86 | zz[:,:,l],ss[:,:,l],ww[:,:,l] = map(reshape, (z[:,l],gas.s[:,l],w[:,l]) ) |
---|
| 87 | x3d[:,:,l]=x[:,:] |
---|
| 88 | |
---|
| 89 | jj=ny/2 |
---|
| 90 | xp=x3d[:,jj,:] |
---|
| 91 | zp=zz[jj,:,:]/1000. |
---|
| 92 | sp=ss[jj,:,:] |
---|
| 93 | wp=ww[jj,:,:] |
---|
| 94 | |
---|
| 95 | f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4)) |
---|
| 96 | |
---|
| 97 | c=ax1.contourf(xp,zp,sp,20) |
---|
| 98 | ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)') |
---|
| 99 | plt.colorbar(c,ax=ax1) |
---|
| 100 | ax1.set_title(title_format % (it*T,)) |
---|
| 101 | |
---|
| 102 | c=ax2.contourf(xp,zp,wp,20) |
---|
| 103 | ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)') |
---|
| 104 | plt.colorbar(c,ax=ax2) |
---|
| 105 | ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,)) |
---|
| 106 | plt.savefig('fig_NH_3D_bubble/%02d.png'%it) |
---|
| 107 | |
---|
| 108 | #------------------------- main program -------------------------- |
---|
| 109 | |
---|
| 110 | parser = argparse.ArgumentParser() |
---|
| 111 | |
---|
| 112 | parser.add_argument("--mpi_ni", type=int, default=64, |
---|
| 113 | help="number of x processors") |
---|
| 114 | parser.add_argument("--mpi_nj", type=int, default=64, |
---|
| 115 | help="number of y processors") |
---|
| 116 | |
---|
| 117 | parser.add_argument("--python_stepping", type=bool, default=False, |
---|
| 118 | help="Time stepping in Python or Fortran") |
---|
| 119 | parser.add_argument("--dt", type=float, default=.25, |
---|
| 120 | help="Time step in seconds") |
---|
| 121 | parser.add_argument("--T", type=float, default=5., |
---|
| 122 | help="Length of time slice in seconds") |
---|
| 123 | parser.add_argument("--Nslice", type=int, default=10, |
---|
| 124 | help="Number of time slices") |
---|
| 125 | |
---|
| 126 | parser.add_argument("--thetac", type=float, default=30., |
---|
| 127 | help="Initial extra temperature of bubble") |
---|
| 128 | parser.add_argument("--Lx", type=float, default=2000., |
---|
| 129 | help="Size of box in meters") |
---|
| 130 | parser.add_argument("--Ly", type=float, default=2000., |
---|
| 131 | help="Size of box in meters") |
---|
| 132 | parser.add_argument("--nx", type=int, default=20, |
---|
| 133 | help="Resolution in the x direction") |
---|
| 134 | parser.add_argument("--ny", type=int, default=20, |
---|
| 135 | help="Resolution in the y direction") |
---|
| 136 | parser.add_argument("--llm", type=int, default=79, |
---|
| 137 | help="Number of vertical levels") |
---|
| 138 | |
---|
| 139 | args = parser.parse_args() |
---|
| 140 | |
---|
[766] | 141 | with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator |
---|
| 142 | comm = client.comm |
---|
| 143 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
| 144 | print '%d/%d starting'%(mpi_rank,mpi_size) |
---|
[768] | 145 | T, Nslice, dt, thetac = args.T, args.Nslice, args.dt, args.thetac |
---|
| 146 | Lx, nx, Ly, ny, llm = args.Lx, args.nx, args.Ly, args.ny, args.llm |
---|
| 147 | |
---|
[766] | 148 | nqdyn, dx = 1, Lx/nx |
---|
| 149 | Ly,ny,dy = Lx,nx,dx |
---|
| 150 | |
---|
| 151 | g=9.81 |
---|
| 152 | unst.setvar('g',g) |
---|
| 153 | |
---|
[768] | 154 | filename = 'cart_%03d_%03d.nc'%(nx,ny) |
---|
| 155 | print 'Reading Cartesian mesh ...' |
---|
| 156 | def coriolis(lon,lat): return 0.*lon |
---|
| 157 | meshfile = meshes.DYNAMICO_Format(filename) |
---|
| 158 | pmesh = meshes.Unstructured_PMesh(comm,meshfile) |
---|
| 159 | pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj) |
---|
| 160 | mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, None, coriolis) |
---|
| 161 | |
---|
[766] | 162 | thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac) |
---|
| 163 | |
---|
| 164 | # compute hybrid coefs from initial distribution of mass |
---|
| 165 | mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0]) |
---|
| 166 | print 'Type of mass_bl, mass_dak, mass_dbk : ', [x.dtype for x in mass_bl, mass_dak, mass_dbk] |
---|
| 167 | unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk) |
---|
| 168 | |
---|
[768] | 169 | dz = flow0[3].max()/(params.g*llm) |
---|
| 170 | # courant = 2.8 |
---|
[766] | 171 | #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2)) |
---|
| 172 | #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2)) |
---|
[768] | 173 | nt = int(math.ceil(T/dt)) |
---|
| 174 | dt = T/nt |
---|
| 175 | print 'Time step : %d x %g = %g s' % (nt,dt,nt*dt) |
---|
[766] | 176 | |
---|
[768] | 177 | # #caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass |
---|
| 178 | caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass |
---|
| 179 | # #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag |
---|
[766] | 180 | |
---|
[768] | 181 | if args.python_stepping: # time stepping in Python |
---|
| 182 | caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g) |
---|
| 183 | scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt) |
---|
| 184 | def next_flow(m,S,u,Phi,W): |
---|
| 185 | return scheme.advance((m,S,u,Phi,W),nt) |
---|
[766] | 186 | |
---|
[768] | 187 | else: # time stepping in Fortran |
---|
| 188 | scheme = time_step.ARK2(None, dt) |
---|
| 189 | caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta, |
---|
| 190 | thermo,params,params.g) |
---|
| 191 | def next_flow(m,S,u,Phi,W): |
---|
| 192 | # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) |
---|
| 193 | caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u |
---|
| 194 | caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W |
---|
| 195 | caldyn_step.next() |
---|
| 196 | return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(), |
---|
| 197 | caldyn_step.geopot.copy(), caldyn_step.W.copy()) |
---|
| 198 | |
---|
| 199 | m,S,u,Phi,W=flow0 |
---|
| 200 | if caldyn_thermo == unst.thermo_theta: |
---|
| 201 | s=S/m |
---|
| 202 | theta = thermo.T0*np.exp(s/thermo.Cpd) |
---|
| 203 | S=m*theta |
---|
| 204 | title_format = 'Potential temperature at t=%g s (K)' |
---|
| 205 | else: |
---|
| 206 | title_format = 'Specific entropy at t=%g s (J/K/kg)' |
---|
| 207 | |
---|
| 208 | w=mesh.field_mass() |
---|
| 209 | z=mesh.field_mass() |
---|
[776] | 210 | xx,yy = mesh.lon_i, mesh.lat_i |
---|
[768] | 211 | |
---|
| 212 | # XIOS writes to disk every 24h |
---|
| 213 | # each iteration lasts it*nt seconds but |
---|
| 214 | # we pretend that each iteration lasts 24h to make sure data is written to disk |
---|
| 215 | |
---|
| 216 | with xios.Context_Curvilinear(mesh,1, 24*3600) as context: |
---|
[766] | 217 | # now XIOS knows about the mesh and we can write to disk |
---|
| 218 | |
---|
[768] | 219 | v = mesh.field_mass() # specific volume (diagnosed) |
---|
| 220 | |
---|
| 221 | for it in range(Nslice): |
---|
| 222 | context.update_calendar(it+1) |
---|
| 223 | |
---|
| 224 | # Diagnose quantities of interest from prognostic variables m,S,u,Phi,W |
---|
| 225 | gas, w, z = diagnose(Phi,S,m,W) |
---|
| 226 | |
---|
| 227 | # write to disk |
---|
| 228 | context.send_field_primal('temp', gas.T) |
---|
| 229 | context.send_field_primal('p', gas.p) |
---|
| 230 | context.send_field_primal('theta', gas.s) |
---|
| 231 | context.send_field_primal('uz', w) |
---|
| 232 | |
---|
| 233 | print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g') |
---|
| 234 | |
---|
| 235 | if args.mpi_ni*args.mpi_nj==1: plot() |
---|
| 236 | |
---|
| 237 | time1, elapsed1 =time.time(), unst.getvar('elapsed') |
---|
| 238 | m,S,u,Phi,W = next_flow(m,S,u,Phi,W) |
---|
| 239 | time2, elapsed2 =time.time(), unst.getvar('elapsed') |
---|
| 240 | factor = 1000./nt |
---|
| 241 | print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) |
---|
| 242 | factor = 1e9/(4*nt*nx*ny*llm) |
---|
| 243 | print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) |
---|
| 244 | |
---|
| 245 | context.update_calendar(Nslice+1) # make sure XIOS writes last iteration |
---|
| 246 | |
---|
| 247 | print('************DONE************') |
---|