- Timestamp:
- 10/18/18 15:13:21 (6 years ago)
- File:
-
- 1 edited
Legend:
- Unmodified
- Added
- Removed
-
codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py
r761 r763 8 8 from dynamico.meshes import Cartesian_mesh as Mesh 9 9 10 from mpi4py import MPI11 comm = MPI.COMM_WORLD12 mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()13 print '%d/%d starting'%(mpi_rank,mpi_size)14 15 10 import math as math 16 11 import numpy as np 17 12 import time 18 19 13 from numpy import pi, log, exp, sin, cos 20 14 … … 57 51 meshfile = meshes.DYNAMICO_Format(filename) 58 52 nqdyn, radius = 1, None 59 # mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, coriolis)60 53 pmesh = meshes.Unstructured_PMesh(comm,meshfile) 54 pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj) 61 55 mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis) 62 56 … … 118 112 return thermo, mesh, params, prec.asnum([mass_ik,Sik,ujk,Phi_il,Wil]), gas 119 113 114 with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator 115 comm = client.comm 116 mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() 117 print '%d/%d starting'%(mpi_rank,mpi_size) 120 118 121 g, Lx, Ly = 9.81, 4e7, 6e6122 nx, ny, llm = 200, 30, 22123 dx,dy=Lx/nx,Ly/ny119 g, Lx, Ly = 9.81, 4e7, 6e6 120 nx, ny, llm = 200, 30, 22 121 dx,dy=Lx/nx,Ly/ny 124 122 125 unst.setvar('g',g)123 unst.setvar('g',g) 126 124 127 thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm)125 thermo, mesh, params, flow0, gas0 = baroclinic_3D(Lx,nx,Ly,ny,llm) 128 126 129 T, nslice, dt = 3600., 1, 3600.127 T, nslice, dt = 3600., 1, 3600. 130 128 131 context = xios.XIOS_Context_Curvilinear(mesh,1, dt) 132 133 for i in range(48): 134 context.update_calendar(i) 135 print 'send_field', i, gas0.T.shape 136 # context.send_field_primal('ps', lat_i) 137 context.send_field_primal('temp', gas0.T) 138 context.send_field_primal('p', gas0.p) 139 140 # to do at the end 141 print 'xios.context_finalize()' 142 context.finalize() 143 print 'xios.finalize()' 144 xios.finalize() 145 print 'Done' 146 147 148 149 # compute hybrid coefs from initial distribution of mass 150 mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0]) 151 print 'Type of mass_bl, mass_dak, mass_dbk : ', [x.dtype for x in mass_bl, mass_dak, mass_dbk] 152 unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk) 153 154 dz = flow0[3].max()/(params.g*llm) 155 #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2+dz**-2)) 156 #dt = courant*.5/np.sqrt(gas0.c2.max()*(dx**-2+dy**-2)) 157 nt = int(math.ceil(T/dt)) 158 dt = T/nt 159 print 'Time step : %d x %g s' % (nt,dt) 160 161 #caldyn_thermo, caldyn_eta = unst.thermo_theta, unst.eta_mass 162 caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass 163 #caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_lag 164 165 if False: # time stepping in Python 166 caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g) 167 scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt) 168 def next_flow(m,S,u,Phi,W): 169 # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) 170 return scheme.advance((m,S,u,Phi,W),nt) 171 172 else: # time stepping in Fortran 173 scheme = time_step.ARK2(None, dt) 174 caldyn_step = unst.caldyn_step_NH(mesh,scheme,nt, caldyn_thermo,caldyn_eta, 175 thermo,params,params.g) 176 def next_flow(m,S,u,Phi,W): 177 # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.) 178 caldyn_step.mass[:,:,:], caldyn_step.theta_rhodz[:,:,:], caldyn_step.u[:,:,:] = m,S,u 179 caldyn_step.geopot[:,:,:], caldyn_step.W[:,:,:] = Phi,W 180 caldyn_step.next() 181 return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(), 182 caldyn_step.geopot.copy(), caldyn_step.W.copy()) 183 184 m,S,u,Phi,W=flow0 185 if caldyn_thermo == unst.thermo_theta: 186 s=S/m 187 theta = thermo.T0*np.exp(s/thermo.Cpd) 188 S=m*theta 189 title_format = 'Potential temperature at t=%g s (K)' 190 else: 191 title_format = 'Specific entropy at t=%g s (J/K/kg)' 192 193 w=mesh.field_mass() 194 z=mesh.field_mass() 195 196 Nslice=0 197 for it in range(Nslice): 198 s=S/m ; s=.5*(s+abs(s)) 199 for l in range(llm): 200 w[:,:,l]=.5*params.g*(W[:,:,l+1]+W[:,:,l])/m[:,:,l] 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 207 f, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(12,4)) 208 209 c=ax1.contourf(xx,zz,ss,20) 210 ax1.set_xlim((-.5,.5)), ax1.set_xlabel('x (km)') 211 ax1.set_ylim((0.,1.)), ax1.set_ylabel('z (km)') 212 plt.colorbar(c,ax=ax1) 213 ax1.set_title(title_format % (it*T,)) 214 215 # plt.show() 216 217 # plt.figure(figsize=(12,5)) 218 c=ax2.contourf(xx,zz,ww,20) 219 ax2.set_xlim((-.5,.5)), ax2.set_xlabel('x (km)') 220 ax2.set_ylim((0.,1.)), ax2.set_ylabel('z (km)') 221 plt.colorbar(c,ax=ax2) 222 ax2.set_title('Vertical velocity at t=%g s (m/s)' % (it*T,)) 223 # plt.tight_layout() 224 # plt.show() 225 plt.savefig('fig_baroclinic_3D/%02d.png'%it) 226 227 time1, elapsed1 =time.time(), unst.getvar('elapsed') 228 m,S,u,Phi,W = next_flow(m,S,u,Phi,W) 229 time2, elapsed2 =time.time(), unst.getvar('elapsed') 230 factor = 1000./nt 231 print 'ms per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) 232 factor = 1e9/(4*nt*nx*ny*llm) 233 print 'nanosec per gridpoint per full time step : ', factor*(time2-time1), factor*(elapsed2-elapsed1) 234 129 with xios.Context_Curvilinear(mesh,1, dt) as context: 130 # now XIOS knows about the mesh and we can write to disk 131 for i in range(48): 132 context.update_calendar(i) 133 print 'send_field', i, gas0.T.shape 134 # context.send_field_primal('ps', lat_i) 135 context.send_field_primal('temp', gas0.T) 136 context.send_field_primal('p', gas0.p)
Note: See TracChangeset
for help on using the changeset viewer.