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 import math as math import matplotlib.pyplot as plt import numpy as np import time import argparse parser = argparse.ArgumentParser() parser.add_argument("-mpi_ni", type=int, default=64, choices=None, help="number of x processors") parser.add_argument("-mpi_nj", type=int, default=64, choices=None, help="number of y processors") args = parser.parse_args() 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) phi0 = 45. # Reference latitude North pi/4 (deg) omega = 7.292e-5 # Angular velocity of the Earth (s^-1) f0 = 2*omega*np.sin(phi0) lap = 0.005 # Lapse rate (K m^-1) Rd = 287.0 # Gas constant for dryy air (j kg^-1 K^-1) def eta(alpha) : return (1-(lap*ztop*alpha/(T0)))**(g/(Rd*lap)) # roughly equispaced levels #def eta(alpha) : return alpha/float(llm) filename = 'cart_%03d_%03d.nc'%(nx,ny) print 'Reading Cartesian mesh ...' def coriolis(lon,lat): return f0+0.*lon meshfile = meshes.DYNAMICO_Format(filename) radius = None #mesh = Mesh(nx,ny,llm,nqdyn,Lx,Ly,0.) print('----read--------') pmesh = meshes.Unstructured_PMesh(comm,meshfile) pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj) mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis) print(mesh.__dict__.keys()) 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('alpha_l=',alpha_l) thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0) eta_il = eta(alpha_il) eta_ik = eta(alpha_ik) eta_ek = eta(alpha_ek) print('eta_il=',eta_il) #Phi_il = Phi(mesh.llp1/float(llm)) #llp is not defined in local mesh #Phi_ik = Phi((mesh.ll+.5)/llm) Phi_il = Phi(eta_il) Phi_ik = Phi(eta_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 #pbot = p(Phi_il[:,:,0]) #gas_bot = thermo.set_pT(pbot, temp(pbot)) # define parameters for lower BC pbot = p(eta_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 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) #Lx, nx, llm, thetac, T, Nslice, courant = 2000., 100, 50, 30., 5., 10, 2.8 Lx, nx, llm, thetac = 2000., 200, 79, 30 #Lx, nx, llm, thetac, T, Nslice, courant = 3000., 75, 25, -30, 5., 10, 2.8 nqdyn, dx = 1, Lx/nx Ly,ny,dy = Lx,nx,dx g=9.81 unst.setvar('g',g) print('bubble call-----------') thermo, mesh, params, flow0, gas0 = thermal_bubble_3D(Lx,nx,Ly,ny,llm, thetac=thetac) print('bubble 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) T, nslice, dt = 3600., 1, 3600. #T, nslice = 3600., 4 with xios.Context_Curvilinear(mesh,1, dt) as context: # now XIOS knows about the mesh and we can write to disk for i in range(48): # 2 days 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) print(gas0.__dict__.keys()) print('size of T, p',np.shape(gas0.T),np.shape(gas0.p)) print('************DONE************') # #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() # # 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] # w[:,l]=.5*params.g*(W[:,l+1]+W[:,l])/m[:,l] # #z[:,:,l]=.5*(Phi[:,:,l+1]+Phi[:,:,l])/params.g # 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,:,:] # #xx,zz,ss,ww = x_ik[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_NH_3D_bubble/%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) #