from dynamico import unstructured as unst from dynamico import dyn from dynamico import time_step from dynamico import meshes from dynamico import xios from dynamico import precision as prec import math as math import matplotlib.pyplot as plt import numpy as np import time import argparse def thermal_bubble_3D(Lx,nx,Ly,ny,llm,ztop=1000., zc=350., rc=250, thetac=0.5, x0=0., y0=0.): Cpd, Rd, g, p0,theta0, T0 = 1004.5, 287.,9.81, 1e5, 300., 300. nqdyn = 1 Phi = lambda eta : g*ztop*eta p = lambda Phi : p0*np.exp(-Phi/(Rd*T0)) zz = lambda p: -(Rd*T0*np.log(p/p0))/g rr = lambda x,y,p: np.sqrt((x-x0)**2 + (y-y0)**2 + (zz(p)-zc)**2) sa = lambda x,y,p: rr(x,y,p) < rc deform = lambda x,y,p: (0.5*thetac*(1+np.cos(np.pi*rr(x,y,p)/rc)))*sa(x,y,p) temp = lambda p: theta0*(p/p0)**(Rd/Cpd) T = lambda x,y,p: deform(x,y,p) + temp(p) alpha_k = (np.arange(llm) +.5)/llm alpha_l = (np.arange(llm+1)+ 0.)/llm x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij') y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij') x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij') y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij') x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij') y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij') thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) Phi_il = Phi(alpha_il) Phi_ik = Phi(alpha_ik) p_ik = p(Phi_ik) T_ik = T(x_ik, y_ik, p_ik) gas = thermo.set_pT(p_ik,T_ik) mass_ik = mesh.field_mass() for l in range(llm): mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l]) Sik, ujk, Wil = gas.s*mass_ik, mesh.field_u(), mesh.field_w() print 'ztop (m) = ', Phi_il[0,-1]/g, ztop ptop = p(g*ztop) print 'ptop (Pa) = ', gas.p[0,-1], ptop params=dyn.Struct() params.ptop=ptop params.dx=dx params.dx_g0=dx/g params.g = g # define parameters for lower BC pbot = p(alpha_il[0]) print 'min p, T :', pbot.min(), temp(pbot/p0) gas_bot = thermo.set_pT(pbot, temp(pbot/p0)) params.pbot = gas_bot.p params.rho_bot = 1e6/gas_bot.v return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas def diagnose(Phi,S,m,W): s=S/m ; s=.5*(s+abs(s)) for l in range(llm): v[:,l]=(Phi[:,l+1]-Phi[:,l])/(g*m[:,l]) w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l] z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/params.g gas = thermo.set_vs(v,s) return gas, w, z def reshape(data): return data.reshape((nx,ny)) def plot(): x, y = map(reshape, (xx,yy) ) zz=np.zeros((nx,ny,llm)) ss=np.zeros((nx,ny,llm)) ww=np.zeros((nx,ny,llm)) x3d=np.zeros((nx,ny,llm)) for l in range(llm): zz[:,:,l],ss[:,:,l],ww[:,:,l] = map(reshape, (z[:,l],gas.s[:,l],w[:,l]) ) x3d[:,:,l]=x[:,:] jj=ny/2 xp=x3d[:,jj,:] zp=zz[jj,:,:]/1000. sp=ss[jj,:,:] wp=ww[jj,:,:] f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4)) c=ax1.contourf(xp,zp,sp,20) ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)') plt.colorbar(c,ax=ax1) ax1.set_title(title_format % (it*T,)) c=ax2.contourf(xp,zp,wp,20) ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)') plt.colorbar(c,ax=ax2) ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,)) plt.savefig('fig_NH_3D_bubble/%02d.png'%it) #------------------------- main program -------------------------- parser = argparse.ArgumentParser() parser.add_argument("--mpi_ni", type=int, default=64, help="number of x processors") parser.add_argument("--mpi_nj", type=int, default=64, help="number of y processors") parser.add_argument("--python_stepping", type=bool, default=False, help="Time stepping in Python or Fortran") parser.add_argument("--dt", type=float, default=.25, help="Time step in seconds") parser.add_argument("--T", type=float, default=5., help="Length of time slice in seconds") parser.add_argument("--Nslice", type=int, default=10, help="Number of time slices") parser.add_argument("--thetac", type=float, default=30., help="Initial extra temperature of bubble") parser.add_argument("--Lx", type=float, default=2000., help="Size of box in meters") parser.add_argument("--Ly", type=float, default=2000., help="Size of box in meters") parser.add_argument("--nx", type=int, default=20, help="Resolution in the x direction") parser.add_argument("--ny", type=int, default=20, help="Resolution in the y direction") parser.add_argument("--llm", type=int, default=79, help="Number of vertical levels") args = parser.parse_args() with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator comm = client.comm mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() print '%d/%d starting'%(mpi_rank,mpi_size) T, Nslice, dt, thetac = args.T, args.Nslice, args.dt, args.thetac Lx, nx, Ly, ny, llm = args.Lx, args.nx, args.Ly, args.ny, args.llm nqdyn, dx = 1, Lx/nx Ly,ny,dy = Lx,nx,dx g=9.81 unst.setvar('g',g) filename = 'cart_%03d_%03d.nc'%(nx,ny) print 'Reading Cartesian mesh ...' def coriolis(lon,lat): return 0.*lon meshfile = meshes.DYNAMICO_Format(filename) pmesh = meshes.Unstructured_PMesh(comm,meshfile) pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj) mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, None, coriolis) thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac) # compute hybrid coefs from initial distribution of mass mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0]) print 'Type of mass_bl, mass_dak, mass_dbk : ', [x.dtype for x in mass_bl, mass_dak, mass_dbk] unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk) dz = flow0[3].max()/(params.g*llm) # courant = 2.8 #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2)) #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2)) nt = int(math.ceil(T/dt)) dt = T/nt print 'Time step : %d x %g = %g s' % (nt,dt,nt*dt) # #caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass # #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag if args.python_stepping: # time stepping in Python caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g) scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt) def next_flow(m,S,u,Phi,W): return scheme.advance((m,S,u,Phi,W),nt) else: # time stepping in Fortran scheme = time_step.ARK2(None, dt) caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta, thermo,params,params.g) def next_flow(m,S,u,Phi,W): # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W caldyn_step.next() return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(), caldyn_step.geopot.copy(), caldyn_step.W.copy()) m,S,u,Phi,W=flow0 if caldyn_thermo == unst.thermo_theta: s=S/m theta = thermo.T0*np.exp(s/thermo.Cpd) S=m*theta title_format = 'Potential temperature at t=%g s (K)' else: title_format = 'Specific entropy at t=%g s (J/K/kg)' w=mesh.field_mass() z=mesh.field_mass() xx,yy = mesh.lon_i, mesh.lat_i # XIOS writes to disk every 24h # each iteration lasts it*nt seconds but # we pretend that each iteration lasts 24h to make sure data is written to disk with xios.Context_Curvilinear(mesh,1, 24*3600) as context: # now XIOS knows about the mesh and we can write to disk v = mesh.field_mass() # specific volume (diagnosed) for it in range(Nslice): context.update_calendar(it+1) # Diagnose quantities of interest from prognostic variables m,S,u,Phi,W gas, w, z = diagnose(Phi,S,m,W) # write to disk context.send_field_primal('temp', gas.T) context.send_field_primal('p', gas.p) context.send_field_primal('theta', gas.s) context.send_field_primal('uz', w) print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g') if args.mpi_ni*args.mpi_nj==1: plot() time1, elapsed1 =time.time(), unst.getvar('elapsed') m,S,u,Phi,W = next_flow(m,S,u,Phi,W) time2, elapsed2 =time.time(), unst.getvar('elapsed') factor = 1000./nt print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) factor = 1e9/(4*nt*nx*ny*llm) print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) context.update_calendar(Nslice+1) # make sure XIOS writes last iteration print('************DONE************')