[632] | 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 | import math as math |
---|
| 7 | import matplotlib.pyplot as plt |
---|
| 8 | import numpy as np |
---|
| 9 | import time |
---|
| 10 | |
---|
| 11 | #------------------------ initial condition ------------------------- |
---|
| 12 | |
---|
| 13 | Cpd, Rd, g = 1004.5, 287., 9.81 |
---|
[678] | 14 | u0, dT, Xfactor = 20., 1., 125. |
---|
| 15 | radius = 6.37122e6/Xfactor |
---|
[632] | 16 | N, Teq, peq = 0.01, 300., 1e5 # background |
---|
| 17 | N2, g2 = N*N, g*g |
---|
| 18 | N2_g2, g2_N2 = N2/g2, g2/N2 |
---|
| 19 | G, kappa = g2/(N2*Cpd), Rd/Cpd |
---|
[678] | 20 | lonc, d2 = 2.*math.pi/3., (625e3/Xfactor)**2 # perturbation |
---|
[632] | 21 | |
---|
| 22 | def psTs(lat) : |
---|
| 23 | cos2latm1 = np.cos(2*lat)-1. |
---|
| 24 | Ts = G+(Teq-G)*np.exp(-.25*N2_g2*u0*u0*cos2latm1) |
---|
| 25 | ps = peq*np.exp(u0/(4.*G*Rd)*u0*cos2latm1)*((Ts/Teq)**(1./kappa)) |
---|
| 26 | return ps, Ts |
---|
| 27 | |
---|
| 28 | def Phi_top(lat) : |
---|
| 29 | ps, Ts = psTs(lat) |
---|
| 30 | return -g2_N2*np.log( 1+(Ts/G)*( ((ptop/ps)**kappa) - 1 ) ) |
---|
| 31 | |
---|
| 32 | def T0_p0(lon,lat,Phi): |
---|
| 33 | ps, Ts = psTs(lat) |
---|
| 34 | pb = ps*( (G/Ts)*(np.exp(-N2_g2*Phi)-1.)+1. )**(1./kappa) |
---|
| 35 | Tb = G + (Ts-G)*np.exp(N2_g2*Phi) |
---|
| 36 | r = radius*np.arccos(np.cos(lat)*np.cos(lon-lonc)) |
---|
| 37 | Theta = dT*np.sin(math.pi*Phi/(g*ztop))*d2/(d2+r*r) |
---|
| 38 | T = Tb + Theta*(pb/peq)**kappa |
---|
| 39 | return T, pb |
---|
| 40 | |
---|
| 41 | def DCMIP31_3D(grid): |
---|
| 42 | def Phi(eta, lat): return eta*Phi_top(lat) |
---|
| 43 | def ulon(lat): return u0*np.cos(lat) |
---|
| 44 | def ulat(lat): return 0.*lat |
---|
| 45 | def f(lon,lat): return 0.*np.sin(lat) # Coriolis parameter |
---|
| 46 | |
---|
| 47 | |
---|
| 48 | nqdyn, preff, Treff = 1, 1e5, 300. |
---|
| 49 | thermo = dyn.Ideal_perfect(Cpd, Rd, preff, Treff) |
---|
[680] | 50 | meshfile = meshes.MPAS_Format('grids/x1.%d.grid.nc'%grid) |
---|
| 51 | mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, f) |
---|
[632] | 52 | lev,levp1 = np.arange(llm),np.arange(llm+1) |
---|
| 53 | lat_ik,k_i = np.meshgrid(mesh.lat_i,lev, indexing='ij') |
---|
| 54 | lon_ik,k_i = np.meshgrid(mesh.lon_i,lev, indexing='ij') |
---|
| 55 | lat_il,l_i = np.meshgrid(mesh.lat_i,levp1, indexing='ij') |
---|
| 56 | lon_il,l_i = np.meshgrid(mesh.lon_i,levp1, indexing='ij') |
---|
| 57 | lat_ek,k_e = np.meshgrid(mesh.lat_e,lev, indexing='ij') |
---|
| 58 | |
---|
| 59 | Phi_ik, Phi_il = Phi(k_i/float(llm),lat_ik), Phi(l_i/float(llm),lat_il) |
---|
| 60 | Tb_ik, pb_ik = T0_p0(lon_ik, lat_ik, Phi_ik) |
---|
| 61 | Tb_il, pb_il = T0_p0(lon_il, lat_il, Phi_il) |
---|
| 62 | gas = thermo.set_pT(pb_ik,Tb_ik) |
---|
| 63 | mass_ik = mesh.field_mass() |
---|
| 64 | for l in range(llm): |
---|
| 65 | mass_ik[:,l]=(pb_il[:,l]-pb_il[:,l+1])/g |
---|
| 66 | Sik = gas.s*mass_ik |
---|
| 67 | |
---|
| 68 | # start at hydrostatic geopotential |
---|
| 69 | pb_ik = .5*(pb_il[:,1:]+pb_il[:,0:-1]) # pressure at full layers |
---|
| 70 | gas = thermo.set_ps(pb_ik,gas.s) |
---|
| 71 | for l in range(llm): |
---|
| 72 | Phi_il[:,l+1] = Phi_il[:,l] + g*mass_ik[:,l]*gas.v[:,l] |
---|
| 73 | |
---|
| 74 | ujk, Wil = mesh.ucov3D(ulon(lat_ek),ulat(lat_ek)), mesh.field_w() |
---|
| 75 | |
---|
| 76 | print 'ztop (m) = ', Phi_il[0,-1]/g, ztop |
---|
| 77 | print 'ptop (Pa) = ', gas.p[0,-1], ptop |
---|
| 78 | dx=mesh.de.min() |
---|
| 79 | params=dyn.Struct() |
---|
| 80 | params.g, params.ztop, params.ptop = g, ztop, ptop |
---|
| 81 | params.dx, params.dx_g0 = dx, dx/g |
---|
| 82 | params.pbot, params.rho_bot = 1e5+0.*mesh.lat_i, 1e6+0.*mesh.lat_i |
---|
| 83 | return thermo, mesh, params, (mass_ik,Sik,ujk,Phi_il,Wil), gas |
---|
| 84 | |
---|
| 85 | #------------------------ main program ------------------------- |
---|
| 86 | |
---|
[678] | 87 | ztop, grid, llm = 1e4, 10242, 20 |
---|
| 88 | T, Nslice, courant = 45000./Xfactor, 10, 3.0 |
---|
[632] | 89 | |
---|
| 90 | junk, ptop = T0_p0(0.,0.,g*ztop) |
---|
| 91 | |
---|
| 92 | thermo, mesh, params, flow0, gas0 = DCMIP31_3D(grid) |
---|
| 93 | llm, dx = mesh.llm, params.dx |
---|
| 94 | dz = params.ztop/llm |
---|
| 95 | print 'llm, nb_hex, dx, dz =', llm, mesh.Ai.size, dx, dz |
---|
| 96 | |
---|
[676] | 97 | #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dz**-2)) |
---|
| 98 | dt = courant*.5*dx/np.sqrt(gas0.c2.max()) |
---|
[632] | 99 | |
---|
[676] | 100 | nt = int(math.ceil(T/dt)) |
---|
| 101 | dt = T/nt |
---|
| 102 | print 'Time step : %g s' % dt |
---|
| 103 | |
---|
[632] | 104 | mesh.plot_e(mesh.le/mesh.de) ; plt.title('le/de') |
---|
| 105 | plt.savefig('fig_NH_3D_DCMIP31/le_de.png') ;plt.close() |
---|
| 106 | |
---|
| 107 | mesh.plot_i(mesh.Ai) ; plt.title('Ai') |
---|
| 108 | plt.savefig('fig_NH_3D_DCMIP31/Ai.png'); plt.close() |
---|
| 109 | |
---|
[676] | 110 | caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag |
---|
[632] | 111 | |
---|
[676] | 112 | if False: # time stepping in Python |
---|
| 113 | caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g) |
---|
| 114 | scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt, a32=0.7) |
---|
| 115 | def next_flow(m,S,u,Phi,W): return scheme.advance((m,S,u,Phi,W),nt) |
---|
| 116 | else: # time stepping in Fortran |
---|
| 117 | scheme = time_step.ARK2(None, dt, a32=0.7) |
---|
| 118 | caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta, thermo,params,params.g) |
---|
| 119 | def next_flow(m,S,u,Phi,W): |
---|
| 120 | caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u |
---|
| 121 | caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W |
---|
| 122 | caldyn_step.next() |
---|
| 123 | return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(), |
---|
| 124 | caldyn_step.geopot.copy(), caldyn_step.W.copy()) |
---|
[632] | 125 | |
---|
| 126 | def plots(it): |
---|
| 127 | s=S/m ; s=.5*(s+abs(s)) |
---|
| 128 | for l in range(llm): |
---|
| 129 | w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l] |
---|
| 130 | z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g |
---|
| 131 | |
---|
| 132 | print( 'ptop, model top (m), max(w), max(W) (m/s) :', |
---|
| 133 | unst.getvar('ptop'), Phi.max()/unst.getvar('g'), w.max(), W.max() ) |
---|
| 134 | |
---|
| 135 | mesh.plot_i(w[:,llm/2]) |
---|
[678] | 136 | plt.title('Vertical velocity at t=%d'%(it*T)) |
---|
[632] | 137 | plt.savefig('fig_NH_3D_DCMIP31/w%02d.png'%it) |
---|
| 138 | plt.close() |
---|
[678] | 139 | mesh.plot_i(z[:,llm/2]) |
---|
| 140 | plt.title('Altitude at t=%d'%(it*T)) |
---|
| 141 | plt.savefig('fig_NH_3D_DCMIP31/z%02d.png'%it) |
---|
| 142 | plt.close() |
---|
[632] | 143 | |
---|
[676] | 144 | w=mesh.field_mass() |
---|
| 145 | z=mesh.field_mass() |
---|
| 146 | m,S,u,Phi,W=flow0 |
---|
[632] | 147 | plots(0) |
---|
| 148 | |
---|
| 149 | for it in range(Nslice): |
---|
[676] | 150 | unst.setvar('debug_hevi_solver',False) |
---|
[632] | 151 | time1, elapsed1 =time.time(), unst.getvar('elapsed') |
---|
[676] | 152 | m,S,u,Phi,W = next_flow(m,S,u,Phi,W) |
---|
[632] | 153 | time2, elapsed2 =time.time(), unst.getvar('elapsed') |
---|
[676] | 154 | factor = 1000./nt |
---|
| 155 | print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) |
---|
[632] | 156 | factor = 1e9/(4*nt*w.size) |
---|
| 157 | print 'nanosec per gridpoint per call to caldyn_hevi : ', factor*(time2-time1), factor*(elapsed2-elapsed1) |
---|
| 158 | plots(it+1) |
---|