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() |
---|