[790] | 1 | from dynamico import meshes |
---|
| 2 | from dynamico import unstructured as unst |
---|
| 3 | from dynamico import time_step |
---|
| 4 | from dynamico import precision as prec |
---|
| 5 | from dynamico import xios |
---|
| 6 | |
---|
| 7 | import math as math |
---|
| 8 | import numpy as np |
---|
| 9 | import argparse |
---|
| 10 | |
---|
| 11 | #------------------ Command-line arguments ------------------ |
---|
| 12 | |
---|
| 13 | parser = argparse.ArgumentParser() |
---|
| 14 | |
---|
| 15 | parser.add_argument("--mpi_ni", type=int, default=1, |
---|
| 16 | help="number of x processors") |
---|
| 17 | parser.add_argument("--mpi_nj", type=int, default=1, |
---|
| 18 | help="number of y processors") |
---|
| 19 | parser.add_argument("-T", type=float, default=10., |
---|
| 20 | help="Length of time slice") |
---|
| 21 | args = parser.parse_args() |
---|
| 22 | |
---|
| 23 | #------------------------- Functions ---------------------- |
---|
| 24 | |
---|
| 25 | def globmin(x): |
---|
| 26 | xloc = x[mesh.primal_own_loc] |
---|
| 27 | locmin= xloc.min() |
---|
| 28 | return client.min(locmin) |
---|
| 29 | |
---|
| 30 | def globmax(x): |
---|
| 31 | xloc = x[mesh.primal_own_loc] |
---|
| 32 | locmin= xloc.max() |
---|
| 33 | return client.max(locmin) |
---|
| 34 | |
---|
| 35 | def minmax(name, x): |
---|
| 36 | gmin,gmax = globmin(x),globmax(x) |
---|
| 37 | if(mpi_rank==0) : print('Min/max %s :'%name, gmin, gmax ) |
---|
| 38 | |
---|
| 39 | def laplace(mesh, p): |
---|
| 40 | left, right, ne = mesh.left, mesh.right, mesh.primal_ne |
---|
| 41 | gradp = mesh.field_u().flatten() |
---|
| 42 | for ij in range(mesh.left.size): |
---|
| 43 | p_left = p[left[ij]-1] # left,right start at 1 (CHECK) |
---|
| 44 | p_right = p[right[ij]-1] # left,right start at 1 (CHECK) |
---|
| 45 | gradp[ij]= mesh.le_de[ij]*(p_right-p_left) # convert to contravariant |
---|
| 46 | lap = mesh.field_mass().flatten() |
---|
| 47 | for ij in range(lap.size): |
---|
| 48 | lap_ij=0. |
---|
| 49 | for k in range(mesh.primal_deg[ij]): |
---|
| 50 | edge=mesh.primal_edge[ij,k]-1 # primal_edge starts at 1 (CHECK) |
---|
| 51 | lap_ij = lap_ij + mesh.primal_ne[ij,k]*gradp[edge] |
---|
| 52 | lap[ij]=lap_ij |
---|
| 53 | return lap/mesh.Ai |
---|
| 54 | |
---|
| 55 | def curlgrad(mesh,p): |
---|
| 56 | left, right, ne = mesh.left, mesh.right, mesh.primal_ne |
---|
| 57 | gradp = mesh.field_u().flatten() |
---|
| 58 | for ij in range(mesh.left.size): |
---|
| 59 | p_left = p[left[ij]-1] # left,right start at 1 (CHECK) |
---|
| 60 | p_right = p[right[ij]-1] # left,right start at 1 (CHECK) |
---|
| 61 | gradp[ij]= p_right-p_left # covariant |
---|
| 62 | curl = mesh.field_z().flatten() |
---|
| 63 | for ij in range(curl.size): |
---|
| 64 | curl_ij=0. |
---|
| 65 | for k in range(mesh.dual_deg[ij]): |
---|
| 66 | edge=mesh.dual_edge[ij,k]-1 # dual_edge starts at 1 (CHECK) |
---|
| 67 | curl_ij = curl_ij + mesh.dual_ne[ij,k]*gradp[edge] |
---|
| 68 | curl[ij]=curl_ij |
---|
| 69 | return curl/mesh.Av |
---|
| 70 | |
---|
| 71 | def curl(mesh, ucov): |
---|
| 72 | curl = mesh.field_z().flatten() |
---|
| 73 | for ij in range(curl.size): |
---|
| 74 | curl_ij=0. |
---|
| 75 | for k in range(mesh.dual_deg[ij]): |
---|
| 76 | edge=mesh.dual_edge[ij,k]-1 # dual_edge starts at 1 (CHECK) |
---|
| 77 | curl_ij = curl_ij + mesh.dual_ne[ij,k]*ucov[edge] |
---|
| 78 | curl[ij]=curl_ij |
---|
| 79 | return curl/mesh.Av |
---|
| 80 | |
---|
| 81 | def test_laplace_seq(): |
---|
| 82 | p=p.flatten() |
---|
| 83 | print p.shape |
---|
| 84 | curl = curlgrad(mesh,p) |
---|
| 85 | print 'curl(grad(p)) min max :', curl.min(), curl.max() |
---|
| 86 | lap = laplace(mesh,p) |
---|
| 87 | p = p+lap/vp |
---|
| 88 | print 'p+lap(p)/vp min max :', p.min(), p.max() |
---|
| 89 | |
---|
| 90 | def test_laplace_parallel(): |
---|
| 91 | transfer_primal, transfer_edge = mesh.com_primal.index, mesh.com_edges.index |
---|
| 92 | with xios.Context_Curvilinear(mesh,1, 24*3600) as context: |
---|
| 93 | # now XIOS knows about the mesh and we can write to disk |
---|
| 94 | context.update_calendar(0) |
---|
| 95 | for it in range(4): |
---|
| 96 | context.update_calendar(it+1) |
---|
| 97 | unst.ker.dynamico_update_halo(mesh.com_primal.index, 1, mesh.primal_num, p) |
---|
| 98 | lap = laplace(mesh,p) |
---|
| 99 | # compute eigenvalue and save residual (should be 0.) |
---|
| 100 | context.send_field_primal('p', p) |
---|
| 101 | p = p+lap/vp |
---|
| 102 | context.send_field_primal('ps', p) |
---|
| 103 | # next iteration will start from lap(p) |
---|
| 104 | p=-lap/vp |
---|
| 105 | |
---|
| 106 | def periodize(z,nz): return (z+2*nz)%nz |
---|
| 107 | def index(x,y): return periodize(x,nx)+nx*periodize(y,ny) |
---|
| 108 | def indexu(x,y): return 2*index(x,y) |
---|
| 109 | def indexv(x,y): return 2*index(x,y)+1 |
---|
| 110 | |
---|
| 111 | def test_RSW(): |
---|
| 112 | x1,x2= xx-1., xx+1. |
---|
| 113 | if False: |
---|
| 114 | u0 = mesh.field_u() # flow initally at rest |
---|
| 115 | h0 = 1+0.1*(np.exp(-2.*(x1*x1+yy*yy))+ |
---|
| 116 | np.exp(-2.*(x2*x2+yy*yy))) |
---|
| 117 | else: # uniform flow |
---|
| 118 | ulon = 1.+0.*mesh.lat_e |
---|
| 119 | ulat = 0.*ulon |
---|
| 120 | u0 = mesh.ucov2D(ulon,ulat) |
---|
| 121 | h0 = mesh.field_mass() |
---|
| 122 | h0[:]=1. |
---|
| 123 | |
---|
| 124 | dt=0.1 |
---|
| 125 | scheme = time_step.RK1(None, dt) |
---|
| 126 | step = unst.caldyn_step_TRSW(mesh,scheme,1) |
---|
| 127 | |
---|
| 128 | # check halo exchange |
---|
| 129 | step.u[:] = mesh.edges_E2 # global index of edges |
---|
| 130 | unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, step.u) |
---|
| 131 | for i in range(step.u.size): |
---|
| 132 | if step.u[i] != mesh.edges_E2[i]: |
---|
| 133 | print '[%d] minmax halo exchanges mismatch : ', i, step.u[i], mesh.edges_E2[i] |
---|
| 134 | |
---|
| 135 | step.mass[:]=h0 |
---|
| 136 | step.theta_rhodz[:]=h0 |
---|
| 137 | step.u[:]=u0 |
---|
| 138 | |
---|
| 139 | def minmax(name, data): print '[%d] minmax %s :'%(mpi_rank,name), data.min(), data.max() |
---|
| 140 | |
---|
| 141 | print 'minmax dual_deg : ', mesh.dual_deg.min(), mesh.dual_deg.max() |
---|
| 142 | with xios.Context_Curvilinear(mesh,1, 24*3600) as context: |
---|
| 143 | for it in range(3): |
---|
| 144 | step.next() |
---|
| 145 | context.update_calendar(it+1) |
---|
| 146 | u, dh, fast, slow = step.u[:], step.drhodz[:], step.du_fast[:], step.du_slow[:] |
---|
| 147 | qv = step.qv[:] |
---|
| 148 | # better exchange halos before evaluating min/max |
---|
| 149 | unst.ker.dynamico_update_halo(mesh.com_primal.index, 1, mesh.primal_num, dh) |
---|
| 150 | unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, fast) |
---|
| 151 | unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, slow) |
---|
| 152 | # unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, qu) |
---|
| 153 | # unst.ker.dynamico_update_halo(mesh.com_dual.index, 1, mesh.dual_num, qv) |
---|
| 154 | minmax('dh',dh) |
---|
| 155 | minmax('du_fast',fast) |
---|
| 156 | minmax('du_slow',slow) |
---|
| 157 | # minmax('qu',qu) |
---|
| 158 | minmax('qv',qv) |
---|
| 159 | context.send_field_primal('p', dh ) |
---|
| 160 | |
---|
| 161 | unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, u) |
---|
| 162 | zeta = curl(mesh, u) |
---|
| 163 | minmax('zeta', zeta) |
---|
| 164 | for i in range(zeta.size): |
---|
| 165 | if zeta[i]<0. : |
---|
| 166 | if mpi_rank==0: |
---|
| 167 | ij = mesh.vertices_V1[i] # global index of dual cell (starting at 0) |
---|
| 168 | x,y = ij%nx, ij/nx |
---|
| 169 | edges1 = [ indexv(x+1,y), indexu(x,y+1), indexv(x,y), indexu(x,y) ] |
---|
| 170 | edges2 = [ mesh.edges_E2[k-1] for k in mesh.dual_edge[i,:] ] |
---|
| 171 | if edges1!=edges2: print '[%d] minmax zeta '%mpi_rank, ij, x, y |
---|
| 172 | print '[%d] minmax zeta '%mpi_rank, [ u[edge-1] for edge in mesh.dual_edge[i,:] ] |
---|
| 173 | print '[%d] minmax zeta '%mpi_rank, mesh.dual_ne[i,:] |
---|
| 174 | |
---|
| 175 | #--------------------------------- Main program ----------------------------- |
---|
| 176 | |
---|
| 177 | with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator |
---|
| 178 | comm = client.comm |
---|
| 179 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
| 180 | print '%d/%d starting'%(mpi_rank,mpi_size) |
---|
| 181 | |
---|
| 182 | # nx,ny,Lx,Ly=200,30,4e7,6e6 |
---|
| 183 | nx,ny,Lx,Ly = 32,32,8.,8. |
---|
| 184 | llm, nqdyn, g, f = 1,1,1.,0. |
---|
| 185 | dx,dy=Lx/nx,Ly/ny |
---|
| 186 | |
---|
| 187 | cfl = .8 |
---|
| 188 | dt = cfl/math.sqrt((nx/Lx)**2+(ny/Ly)**2) |
---|
| 189 | T=args.T |
---|
| 190 | N=int(T/dt)+1 |
---|
| 191 | dt=T/N |
---|
| 192 | print N,dt,Lx/nx |
---|
| 193 | |
---|
| 194 | filename, radius = 'cart_%03d_%03d.nc'%(nx,ny), None |
---|
| 195 | unst.setvar('g',g) |
---|
| 196 | |
---|
| 197 | def coriolis(lon,lat): |
---|
| 198 | return f+0.*lon |
---|
| 199 | |
---|
| 200 | print 'Reading Cartesian mesh ...' |
---|
| 201 | meshfile = meshes.DYNAMICO_Format(filename) |
---|
| 202 | |
---|
| 203 | parallel=True |
---|
| 204 | laplace=False |
---|
| 205 | |
---|
| 206 | if parallel: |
---|
| 207 | pmesh = meshes.Unstructured_PMesh(comm,meshfile) |
---|
| 208 | pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj) |
---|
| 209 | # pmesh.partition_metis() |
---|
| 210 | mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis) |
---|
| 211 | else: |
---|
| 212 | mesh = meshes.Cartesian_mesh(nx,ny,llm,nqdyn,Lx,Ly,f) |
---|
| 213 | |
---|
| 214 | xx,yy = mesh.lon_i, mesh.lat_i |
---|
| 215 | |
---|
| 216 | if laplace: |
---|
| 217 | p = np.cos(2*np.pi*xx/Lx)*np.cos(2.*np.pi*yy/Ly) |
---|
| 218 | |
---|
| 219 | # analytic eigenvalue : vp = -((2*np/pi/Lx)**2+(2*np.pi/Ly)**2) |
---|
| 220 | kx = 2.*np.sin(dx*np.pi/Lx)/dx |
---|
| 221 | ky = 2.*np.sin(dy*np.pi/Ly)/dy |
---|
| 222 | vp = kx**2+ky**2 |
---|
| 223 | |
---|
| 224 | if parallel: |
---|
| 225 | test_laplace_parallel() |
---|
| 226 | else: |
---|
| 227 | test_laplace_seq() |
---|
| 228 | |
---|
| 229 | else: |
---|
| 230 | test_RSW() |
---|