from dynamico import unstructured as unst from dynamico import dyn from dynamico import time_step from dynamico import DCMIP from dynamico import meshes from dynamico import xios from dynamico import precision as prec from dynamico.meshes import Cartesian_mesh as Mesh from mpi4py import MPI comm = MPI.COMM_WORLD mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() print '%d/%d starting'%(mpi_rank,mpi_size) import math as math import numpy as np import time from numpy import pi, log, exp, sin, cos # Baroclinic instability test based on Ullrich et al. 2015, QJRMS def baroclinic_3D(Lx,nx,Ly,ny,llm,ztop=25000.): Rd = 287.0 # Gas constant for dryy air (j kg^-1 K^-1) T0 = 288.0 # Reference temperature (K) lap = 0.005 # Lapse rate (K m^-1) b = 2. # Non dimensional vertical width parameter u0 = 35. # Reference zonal wind speed (m s^-1) a = 6.371229e6 # Radius of the Earth (m) ptop = 2000. y0 = Ly*0.5 Cpd = 1004.5 p0 = 1e5 omega = 7.292e-5 # Angular velocity of the Earth (s^-1) phi0 = 45. # Reference latitude North pi/4 (deg) f0 = 2*omega*np.sin(phi0) beta0 = 2*omega*np.cos(phi0)/a fb = 2*omega*np.sin(phi0) - y0*2*omega*np.cos(phi0)/a def Phi_xy(x,y): fc = y*y - (Ly*y/pi)*sin(2*pi*y/Ly) fd = Ly*Ly/(2*pi*pi)*cos(2*pi*y/Ly) 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)) ) 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)) def ulon(x,y,eta): return -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/b/b)) def tmean(eta) : return T0*eta**(Rd*lap/g) 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)) def p(eta): return p0*eta # eta = p/p_s def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels filename = 'cart_%03d_%03d.nc'%(nx,ny) print 'Reading Cartesian mesh ...' def coriolis(lon,lat): return f0+0.*lon meshfile = meshes.DYNAMICO_Format(filename) nqdyn, radius = 1, None # mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, coriolis) pmesh = meshes.Unstructured_PMesh(comm,meshfile) mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis) 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') print('----------------') print 'ztop(ptop) according to Eq. 7:', T0/lap*(1.-(ptop/p0)**(Rd*lap/g)) print(np.shape(alpha_k),np.shape(alpha_l)) print(mesh.__dict__.keys()) thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) eta_il = eta(alpha_il) eta_ik = eta(alpha_ik) eta_ek = eta(alpha_ek) print('min max eta_il', np.min(eta_il),np.max(eta_il)) Phi_il = Phi_xyeta(x_il, y_il, eta_il) Phi_ik = Phi_xyeta(x_ik, y_ik, eta_ik) p_ik = p(eta_ik) T_ik = T(x_ik, y_ik, eta_ik) #ik full level(40), il(41) 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(np.shape(ujk)) print('P_ik',p_ik[0,:]) u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek) print(np.shape(u_ek)) print 'ztop (m) = ', Phi_il[0,-1]/g, ztop ptop = p(eta(1.)) 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(eta_il[:,0]) print 'min p, T :', pbot.min(), tmean(pbot/p0) gas_bot = thermo.set_pT(pbot, tmean(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 g, Lx, Ly = 9.81, 4e7, 6e6 nx, ny, llm = 200, 30, 22 dx,dy=Lx/nx,Ly/ny unst.setvar('g',g) thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm) T, nslice, dt = 3600., 1, 3600. context = xios.XIOS_Context_Curvilinear(mesh,1, dt) for i in range(48): context.update_calendar(i) print 'send_field', i, gas0.T.shape # context.send_field_primal('ps', lat_i) context.send_field_primal('temp', gas0.T) context.send_field_primal('p', gas0.p) # to do at the end print 'xios.context_finalize()' context.finalize() print 'xios.finalize()' xios.finalize() print 'Done' # 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) #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 s' % (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 False: # 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): # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) 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() Nslice=0 for it in range(Nslice): s=S/m ; s=.5*(s+abs(s)) for l in range(llm): w[:,:,l]=.5*params.g*(W[:,:,l+1]+W[:,:,l])/m[:,:,l] z[:,:,l]=.5*(Phi[:,:,l+1]+Phi[:,:,l])/params.g print 'ptop, model top (m) :', unst.getvar('ptop'), Phi.max()/unst.getvar('g') jj=ny/2 xx,zz,ss,ww = mesh.xx[:,jj,:]/1000., z[:,jj,:]/1000., s[:,jj,:], w[:,jj,:] f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4)) c=ax1.contourf(xx,zz,ss,20) ax1.set_xlim((-.5,.5)), ax1.set_xlabel('x (km)') ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)') plt.colorbar(c,ax=ax1) ax1.set_title(title_format % (it*T,)) # plt.show() # plt.figure(figsize=(12,5)) c=ax2.contourf(xx,zz,ww,20) ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)') 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.tight_layout() # plt.show() plt.savefig('fig_baroclinic_3D/%02d.png'%it) 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)