- Timestamp:
- 10/31/18 14:22:54 (6 years ago)
- File:
-
- 1 edited
Legend:
- Unmodified
- Added
- Removed
-
codes/icosagcm/devel/Python/test/py/NH_3D_bubble_parallel.py
r766 r768 2 2 from dynamico import dyn 3 3 from dynamico import time_step 4 from dynamico import DCMIP5 4 from dynamico import meshes 6 5 from dynamico import xios 7 6 from dynamico import precision as prec 8 #from dynamico.meshes import Cartesian_mesh as Mesh9 7 import math as math 10 8 import matplotlib.pyplot as plt … … 12 10 import time 13 11 import argparse 14 15 parser = argparse.ArgumentParser()16 17 parser.add_argument("-mpi_ni", type=int,18 default=64, choices=None,19 help="number of x processors")20 parser.add_argument("-mpi_nj", type=int,21 default=64, choices=None,22 help="number of y processors")23 args = parser.parse_args()24 25 12 26 13 def thermal_bubble_3D(Lx,nx,Ly,ny,llm,ztop=1000., zc=350., … … 37 24 temp = lambda p: theta0*(p/p0)**(Rd/Cpd) 38 25 T = lambda x,y,p: deform(x,y,p) + temp(p) 39 phi0 = 45. # Reference latitude North pi/4 (deg)40 omega = 7.292e-5 # Angular velocity of the Earth (s^-1)41 f0 = 2*omega*np.sin(phi0)42 lap = 0.005 # Lapse rate (K m^-1)43 Rd = 287.0 # Gas constant for dryy air (j kg^-1 K^-1)44 45 def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels46 #def eta(alpha) : return alpha/float(llm)47 48 filename = 'cart_%03d_%03d.nc'%(nx,ny)49 print 'Reading Cartesian mesh ...'50 def coriolis(lon,lat): return f0+0.*lon51 meshfile = meshes.DYNAMICO_Format(filename)52 radius = None53 #mesh = Mesh(nx,ny,llm,nqdyn,Lx,Ly,0.)54 print('----read--------')55 pmesh = meshes.Unstructured_PMesh(comm,meshfile)56 pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)57 mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)58 59 print(mesh.__dict__.keys())60 26 61 27 alpha_k = (np.arange(llm) +.5)/llm … … 68 34 y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij') 69 35 70 print('alpha_l=',alpha_l)71 72 36 thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) 73 37 74 eta_il = eta(alpha_il) 75 eta_ik = eta(alpha_ik) 76 eta_ek = eta(alpha_ek) 77 78 print('eta_il=',eta_il) 79 #Phi_il = Phi(mesh.llp1/float(llm)) #llp is not defined in local mesh 80 #Phi_ik = Phi((mesh.ll+.5)/llm) 81 Phi_il = Phi(eta_il) 82 Phi_ik = Phi(eta_ik) 38 Phi_il = Phi(alpha_il) 39 Phi_ik = Phi(alpha_ik) 83 40 p_ik = p(Phi_ik) 84 41 T_ik = T(x_ik, y_ik, p_ik) … … 98 55 params.dx_g0=dx/g 99 56 params.g = g 100 #pbot = p(Phi_il[:,:,0]) 101 #gas_bot = thermo.set_pT(pbot, temp(pbot)) 57 102 58 # define parameters for lower BC 103 pbot = p( eta_il[0])59 pbot = p(alpha_il[0]) 104 60 print 'min p, T :', pbot.min(), temp(pbot/p0) 105 61 gas_bot = thermo.set_pT(pbot, temp(pbot/p0)) … … 109 65 return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas 110 66 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 111 141 with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator 112 142 comm = client.comm 113 143 mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() 114 144 print '%d/%d starting'%(mpi_rank,mpi_size) 115 116 #Lx, nx, llm, thetac, T, Nslice, courant = 2000., 100, 50, 30., 5., 10, 2.8 117 Lx, nx, llm, thetac = 2000., 200, 79, 30 118 #Lx, nx, llm, thetac, T, Nslice, courant = 3000., 75, 25, -30, 5., 10, 2.8 119 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 120 148 nqdyn, dx = 1, Lx/nx 121 149 Ly,ny,dy = Lx,nx,dx … … 124 152 unst.setvar('g',g) 125 153 126 print('bubble call-----------') 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 127 162 thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac) 128 print('bubble done-----------')129 130 163 131 164 # compute hybrid coefs from initial distribution of mass … … 134 167 unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk) 135 168 136 #dz = flow0[3].max()/(params.g*llm) 169 dz = flow0[3].max()/(params.g*llm) 170 # courant = 2.8 137 171 #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2)) 138 172 #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2)) 139 #nt = int(math.ceil(T/dt)) 140 #dt = T/nt 141 #print 'Time step : %d x %g s' % (nt,dt) 142 143 144 T, nslice, dt = 3600., 1, 3600. 145 #T, nslice = 3600., 4 146 147 with xios.Context_Curvilinear(mesh,1, dt) as context: 173 nt = int(math.ceil(T/dt)) 174 dt = T/nt 175 print 'Time step : %d x %g = %g s' % (nt,dt,nt*dt) 176 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 180 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) 186 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() 210 xx,yy = mesh.lat_i, mesh.lon_i 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: 148 217 # now XIOS knows about the mesh and we can write to disk 149 for i in range(48): # 2 days 150 context.update_calendar(i) 151 152 print 'send_field', i, gas0.T.shape 153 # context.send_field_primal('ps', lat_i) 154 context.send_field_primal('temp', gas0.T) 155 context.send_field_primal('p', gas0.p) 156 print(gas0.__dict__.keys()) 157 print('size of T, p',np.shape(gas0.T),np.shape(gas0.p)) 158 print('************DONE************') 159 160 # #caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass 161 # caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass 162 # #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag 163 # 164 # if False: # time stepping in Python 165 # caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g) 166 # scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt) 167 # def next_flow(m,S,u,Phi,W): 168 # # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) 169 # return scheme.advance((m,S,u,Phi,W),nt) 170 # 171 # else: # time stepping in Fortran 172 # scheme = time_step.ARK2(None, dt) 173 # caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta, 174 # thermo,params,params.g) 175 # def next_flow(m,S,u,Phi,W): 176 # # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) 177 # caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u 178 # caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W 179 # caldyn_step.next() 180 # return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(), 181 # caldyn_step.geopot.copy(), caldyn_step.W.copy()) 182 # 183 # m,S,u,Phi,W=flow0 184 # if caldyn_thermo == unst.thermo_theta: 185 # s=S/m 186 # theta = thermo.T0*np.exp(s/thermo.Cpd) 187 # S=m*theta 188 # title_format = 'Potential temperature at t=%g s (K)' 189 # else: 190 # title_format = 'Specific entropy at t=%g s (J/K/kg)' 191 # 192 # w=mesh.field_mass() 193 # z=mesh.field_mass() 194 # 195 # for it in range(Nslice): 196 # s=S/m ; s=.5*(s+abs(s)) 197 # for l in range(llm): 198 # #w[:,:,l]=.5*params.g*(W[:,:,l+1]+W[:,:,l])/m[:,:,l] 199 # w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l] 200 # #z[:,:,l]=.5*(Phi[:,:,l+1]+Phi[:,:,l])/params.g 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 # #xx,zz,ss,ww = x_ik[jj,:]/1000., z[jj,:]/1000., s[jj,:], w[jj,:] 207 # 208 # #f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4)) 209 # 210 # #c=ax1.contourf(xx,zz,ss,20) 211 # #ax1.set_xlim((-.5,.5)), ax1.set_xlabel('x (km)') 212 # #ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)') 213 # #plt.colorbar(c,ax=ax1) 214 # #ax1.set_title(title_format % (it*T,)) 215 # 216 # # plt.show() 217 # 218 # # plt.figure(figsize=(12,5)) 219 # #c=ax2.contourf(xx,zz,ww,20) 220 # #ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)') 221 # #ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)') 222 # #plt.colorbar(c,ax=ax2) 223 # #ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,)) 224 # # plt.tight_layout() 225 # # plt.show() 226 # #plt.savefig('fig_NH_3D_bubble/%02d.png'%it) 227 # 228 # time1, elapsed1 =time.time(), unst.getvar('elapsed') 229 # m,S,u,Phi,W = next_flow(m,S,u,Phi,W) 230 # time2, elapsed2 =time.time(), unst.getvar('elapsed') 231 # factor = 1000./nt 232 # print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) 233 # factor = 1e9/(4*nt*nx*ny*llm) 234 # print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) 235 # 218 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************')
Note: See TracChangeset
for help on using the changeset viewer.