1 | #!/usr/bin/env python |
---|
2 | ################################################################################## |
---|
3 | # Utility to produce animation frames directly from the WAD_TEST_CASE output. |
---|
4 | # It can be used without arguments, i.e.: |
---|
5 | # python plotframes.py |
---|
6 | # in which case a frame is created every tenth time level from SSH data extracted from |
---|
7 | # WAD_1ts_00010101_00010101_grid_T.nc along the centre of the basin (j=17). A closed |
---|
8 | # basin is assumed and frames are named wadfr0000.png etc. Bathymetry information is |
---|
9 | # extracted from the mesh_mask.nc file. The frames are annotated with a timestamp that |
---|
10 | # assumes an 18s baroclinic timestep. |
---|
11 | # |
---|
12 | # All these settings can be overridden with command-line arguments. See: |
---|
13 | # python plotframes.py -h |
---|
14 | # for details. For example: |
---|
15 | # python plotframes.py -nt 300 -stride 30 -froot mywad |
---|
16 | # |
---|
17 | # Two major variations are also supported for specific test cases: |
---|
18 | # python plotframes.py -obc |
---|
19 | # plots the right-hand side of the basin as an open boundary (test case 7) and: |
---|
20 | # python plotframes.py -use_sal |
---|
21 | # colours each gridcell according to its salinity value (test case 6) |
---|
22 | ################################################################################## |
---|
23 | import os, sys |
---|
24 | from argparse import ArgumentParser |
---|
25 | import numpy as np |
---|
26 | import netCDF4 |
---|
27 | |
---|
28 | import matplotlib.pyplot as plt |
---|
29 | import matplotlib |
---|
30 | from matplotlib.patches import Polygon |
---|
31 | from matplotlib.collections import PatchCollection |
---|
32 | from netCDF4 import Dataset |
---|
33 | |
---|
34 | # |
---|
35 | # Turn off the unhelpful warning about open figures |
---|
36 | # |
---|
37 | matplotlib.rcParams['figure.max_open_warning'] = 0 |
---|
38 | |
---|
39 | if __name__ == '__main__': |
---|
40 | parser = ArgumentParser(description= |
---|
41 | """ |
---|
42 | produce frames for the animation of results from the WAD_TEST_CASES. |
---|
43 | Mostly this can be run without arguments but command line arguments may be |
---|
44 | used to override defaults. These are necessary in cases with open boundaries |
---|
45 | (e.g. nn_wad_test=7) and cases where it is desired to show variations in salinity |
---|
46 | (e.g. nn_wad_test=6). |
---|
47 | e.g. plotframes.py -tfile <T-grid file> -bfile <bathymetry file> |
---|
48 | -froot <root name for frames> |
---|
49 | -nt <maximum number of time frames to process> |
---|
50 | -stride <stride through time frames> |
---|
51 | -rdt <length of baroclinic timestep (s)> |
---|
52 | -obc |
---|
53 | -use_sal |
---|
54 | """) |
---|
55 | parser.add_argument('-tfile',dest='tfile',help='T-grid file if not WAD_1ts_00010101_00010101_grid_T.nc', default='WAD_1ts_00010101_00010101_grid_T.nc') |
---|
56 | parser.add_argument('-bfile',dest='bfile',help='Bathymetry file if not mesh_mask.nc', default='mesh_mask.nc') |
---|
57 | parser.add_argument('-j',dest='jrow',help='jrow; j-row to extract and plot (default: 17 (fortran index))', type=int, default=16) |
---|
58 | parser.add_argument('-froot',dest='froot',help='froot; root name for frames (default: wadfr)', default='wadfr') |
---|
59 | parser.add_argument('-nt',dest='nfmax',help='nfmax; maximum number of frames to produce', type=int, default=None) |
---|
60 | parser.add_argument('-stride',dest='tinc',help='tinc; stride through time frames (default: 10)', type=int, default=10) |
---|
61 | parser.add_argument('-rdt',dest='rdt',help='rdt; length of baroclinic timestep (s) (default: 18.0)', type=float, default=18.0) |
---|
62 | parser.add_argument('-obc',help='Right-hand side boundary is open', action="store_true") |
---|
63 | parser.add_argument('-use_sal',help='colour polygons according to salinity variations', action="store_true") |
---|
64 | args = parser.parse_args() |
---|
65 | tfile = args.tfile |
---|
66 | bfile = args.bfile |
---|
67 | jrow = args.jrow |
---|
68 | froot = args.froot |
---|
69 | nfmax = args.nfmax |
---|
70 | stride = args.tinc |
---|
71 | rdt = args.rdt |
---|
72 | obc = args.obc |
---|
73 | use_sal= args.use_sal |
---|
74 | |
---|
75 | |
---|
76 | fw = Dataset(tfile) |
---|
77 | ssh = fw.variables['sossheig'][:,jrow,:] |
---|
78 | vot = fw.variables['sosaline'][:,jrow,:] |
---|
79 | if use_sal: |
---|
80 | sal = fw.variables['vosaline'][:,:,jrow,:] |
---|
81 | nz = sal.shape[1] |
---|
82 | fw.close() |
---|
83 | |
---|
84 | fw = Dataset(bfile) |
---|
85 | bat = fw.variables['ht_wd'][0,jrow,:] |
---|
86 | if use_sal: |
---|
87 | mbat = fw.variables['mbathy'][0,jrow,:] |
---|
88 | fw.close() |
---|
89 | |
---|
90 | #print "ssh" |
---|
91 | #print ssh.shape |
---|
92 | #print "bat" |
---|
93 | #print bat.shape |
---|
94 | #print "vot" |
---|
95 | #print vot.shape |
---|
96 | nt = ssh.shape[0] |
---|
97 | nx = ssh.shape[1] |
---|
98 | #print nx,nt |
---|
99 | |
---|
100 | bat = -1.*bat |
---|
101 | batmin = np.amin(bat) |
---|
102 | batmax = np.amax(bat) |
---|
103 | brange = batmax - batmin |
---|
104 | tol = 0.1*brange |
---|
105 | #print batmin,batmax,' ho' |
---|
106 | if obc: |
---|
107 | batrhs = batmin |
---|
108 | else: |
---|
109 | batrhs = batmax |
---|
110 | |
---|
111 | if nfmax is None: |
---|
112 | nfmax = nt |
---|
113 | |
---|
114 | nf = 0 |
---|
115 | ntmax = np.minimum(nt,nfmax*stride) |
---|
116 | |
---|
117 | if not use_sal: |
---|
118 | # |
---|
119 | # plot solid single colour polygons just showing ssh variation |
---|
120 | # |
---|
121 | for t in range(0,ntmax,stride): |
---|
122 | wadfr = froot+"{:0>4d}.png".format(nf) |
---|
123 | nf = nf + 1 |
---|
124 | |
---|
125 | tfac = rdt/3600.0 |
---|
126 | t24 = np.int(np.mod(t*tfac,24)) |
---|
127 | dy = np.int(t*tfac/24.0) |
---|
128 | mn = np.int(np.rint((np.mod(t*tfac,24) - t24 )*60)) |
---|
129 | hour = "t={:0>2d}:{:0>2d}:{:0>2d} ".format(dy,t24,mn) |
---|
130 | hour2 = " (days:hrs:mins)" |
---|
131 | batpts = np.zeros((nx+4,2)) |
---|
132 | sshpts = np.zeros((2*nx,2)) |
---|
133 | votpts = np.zeros((nx,2)) |
---|
134 | |
---|
135 | for pt in range(nx): |
---|
136 | batpts[pt+2,0] = pt |
---|
137 | batpts[pt+2,1] = bat[pt] |
---|
138 | sshpts[pt,0] = pt |
---|
139 | sshpts[pt,1] = ssh[t,pt] |
---|
140 | votpts[pt,0] = pt |
---|
141 | votpts[pt,1] = np.minimum(36.,vot[t,pt]) |
---|
142 | votpts[pt,1] = np.maximum(30.0,votpts[pt,1]) |
---|
143 | votpts[pt,1] = batmin +0.2*brange + (votpts[pt,1]-30.)*brange/6.0 |
---|
144 | |
---|
145 | batpts[nx+1,1] = batrhs |
---|
146 | batpts[nx+2,0] = nx-1 |
---|
147 | batpts[nx+2,1] = batrhs |
---|
148 | batpts[nx+3,0] = nx-1 |
---|
149 | batpts[nx+3,1] = batmin |
---|
150 | batpts[0,0] = 0.0 |
---|
151 | batpts[0,1] = batmin |
---|
152 | batpts[1,0] = 0.0 |
---|
153 | batpts[1,1] = batmax |
---|
154 | batpts[2,1] = batmax |
---|
155 | sshpts[nx-1,0] = nx-1 |
---|
156 | sshpts[nx-1,1] = sshpts[nx-2,1] |
---|
157 | sshpts[0,0] = 0.0 |
---|
158 | votpts[nx-2,0] = nx-1 |
---|
159 | votpts[nx-1,0] = nx-1 |
---|
160 | votpts[0,0] = 0.0 |
---|
161 | votpts[nx-1,1] = batmax + tol |
---|
162 | votpts[0,1] = batmax + tol |
---|
163 | sshpts[0,1] = sshpts[1,1] |
---|
164 | for pt in range(nx): |
---|
165 | sshpts[pt+nx,0]=batpts[nx+1-pt,0] |
---|
166 | sshpts[pt+nx,1]=batpts[nx+1-pt,1] |
---|
167 | |
---|
168 | |
---|
169 | xs, ys = zip(*votpts) |
---|
170 | fig, ax = plt.subplots() |
---|
171 | |
---|
172 | patches = [] |
---|
173 | polygon = Polygon(batpts, True) |
---|
174 | patches.append(polygon) |
---|
175 | p = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0) |
---|
176 | p.set_facecolors(['#f1a9a9']) |
---|
177 | |
---|
178 | patches = [] |
---|
179 | polygon2 = Polygon(sshpts, True) |
---|
180 | patches.append(polygon2) |
---|
181 | p2 = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0) |
---|
182 | p2.set_facecolors(['#44a1ff']) |
---|
183 | |
---|
184 | # Maximum depth set here to -10m |
---|
185 | ax.set_ylim([-10., 6.0]) |
---|
186 | ax.set_xlim([0., 51.0]) |
---|
187 | ax.add_collection(p2) |
---|
188 | ax.add_collection(p) |
---|
189 | ax.plot(xs,ys, '--', color='black', ms=10) |
---|
190 | |
---|
191 | plt.annotate(hour,xy=(2,batmin+0.1*brange)) |
---|
192 | plt.annotate(hour2,xy=(2,batmin+0.05*brange)) |
---|
193 | plt.savefig(wadfr) |
---|
194 | |
---|
195 | else: |
---|
196 | # |
---|
197 | # plot each gridcell coloured according to its salinity value |
---|
198 | # |
---|
199 | for t in range(0,ntmax,stride): |
---|
200 | wadfr = froot+"{:0>4d}.png".format(nf) |
---|
201 | nf = nf + 1 |
---|
202 | |
---|
203 | tfac = rdt/3600.0 |
---|
204 | t24 = np.int(np.mod(t*tfac,24)) |
---|
205 | dy = np.int(t*tfac/24.0) |
---|
206 | mn = np.int(np.rint((np.mod(t*tfac,24) - t24 )*60)) |
---|
207 | hour = "t={:0>2d}:{:0>2d}:{:0>2d} ".format(dy,t24,mn) |
---|
208 | hour2 = " (days:hrs:mins)" |
---|
209 | batpts = np.zeros((nx+4,2)) |
---|
210 | votpts = np.zeros((nx,2)) |
---|
211 | salpts = np.zeros((nx*nz,6,2)) |
---|
212 | salmin = 28. |
---|
213 | salmax = 37. |
---|
214 | salrange = salmax - salmin |
---|
215 | faccol = np.zeros((nx*nz)) |
---|
216 | cl = 0 |
---|
217 | for pt in range(nx): |
---|
218 | batpts[pt+2,0] = pt |
---|
219 | batpts[pt+2,1] = bat[pt] |
---|
220 | votpts[pt,0] = pt |
---|
221 | votpts[pt,1] = np.minimum(35.,vot[t,pt]) |
---|
222 | votpts[pt,1] = np.maximum(30.0,votpts[pt,1]) |
---|
223 | votpts[pt,1] = batmin +0.2*brange + (votpts[pt,1]-30.)*brange/6.0 |
---|
224 | batpts[nx+1,1] = batmax |
---|
225 | batpts[nx+2,0] = nx-1 |
---|
226 | batpts[nx+2,1] = batmax |
---|
227 | batpts[nx+3,0] = nx-1 |
---|
228 | batpts[nx+3,1] = batmin |
---|
229 | batpts[0,0] = 0.0 |
---|
230 | batpts[0,1] = batmin |
---|
231 | batpts[1,0] = 0.0 |
---|
232 | batpts[1,1] = batmax |
---|
233 | batpts[2,1] = batmax |
---|
234 | votpts[nx-2,0] = nx-1 |
---|
235 | votpts[nx-1,0] = nx-1 |
---|
236 | votpts[0,0] = 0.0 |
---|
237 | votpts[nx-1,1] = batmax + tol |
---|
238 | votpts[0,1] = batmax + tol |
---|
239 | |
---|
240 | cl = 0 |
---|
241 | for pt in range(nx): |
---|
242 | mz = np.maximum(1,mbat[pt]) |
---|
243 | im1 = np.maximum(pt-1,1) |
---|
244 | ip1 = np.minimum(pt+1,nx-1) |
---|
245 | dz = (ssh[t,pt] - batpts[pt+2,1] )/mz |
---|
246 | dz1 = 0.5*(ssh[t,pt] + ssh[t,im1] - batpts[pt+2,1] - batpts[im1+2,1] )/mz |
---|
247 | dz2 = 0.5*(ssh[t,pt] + ssh[t,ip1] - batpts[pt+2,1] - batpts[ip1+2,1] )/mz |
---|
248 | dz = np.maximum(dz ,0.0) |
---|
249 | dz1 = np.maximum(dz1,0.0) |
---|
250 | dz2 = np.maximum(dz2,0.0) |
---|
251 | bat1 = 0.5*( batpts[pt+2,1] + batpts[im1+2,1] ) |
---|
252 | bat2 = 0.5*( batpts[pt+2,1] + batpts[ip1+2,1] ) |
---|
253 | ptm = np.maximum(pt-0.5,0.0) |
---|
254 | ptx = np.minimum(pt+0.5,nx-1) |
---|
255 | for z in range(mz): |
---|
256 | if ( sal[t,mz-1-z,pt] > 0.0 ): |
---|
257 | salpts[cl,0,0] = pt |
---|
258 | salpts[cl,1,0] = ptm |
---|
259 | salpts[cl,2,0] = ptm |
---|
260 | salpts[cl,3,0] = pt |
---|
261 | salpts[cl,4,0] = ptx |
---|
262 | salpts[cl,5,0] = ptx |
---|
263 | salpts[cl,0,1] = batpts[pt+2,1] +dz*z |
---|
264 | salpts[cl,1,1] = bat1 +dz1*z |
---|
265 | salpts[cl,2,1] = bat1 +dz1*(z+1) |
---|
266 | salpts[cl,3,1] = batpts[pt+2,1] +dz*(z+1) |
---|
267 | salpts[cl,4,1] = bat2 +dz2*(z+1) |
---|
268 | salpts[cl,5,1] = bat2 +dz2*z |
---|
269 | faccol[cl] = 100*(sal[t,mz-1-z,pt] - salmin) / salrange |
---|
270 | faccol[cl] = np.maximum(faccol[cl],0.0) |
---|
271 | faccol[cl] = np.minimum(faccol[cl],100.0) |
---|
272 | cl = cl + 1 |
---|
273 | |
---|
274 | |
---|
275 | votpts2 = votpts[2:nx-4,:] |
---|
276 | xs, ys = zip(*votpts2) |
---|
277 | fig, ax = plt.subplots() |
---|
278 | patches = [] |
---|
279 | |
---|
280 | for pt in range(cl-1): |
---|
281 | polygon = Polygon(salpts[pt,:,:], True) |
---|
282 | patches.append(polygon) |
---|
283 | |
---|
284 | p = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0) |
---|
285 | |
---|
286 | p.set_array(faccol) |
---|
287 | p.set_edgecolor('face') |
---|
288 | |
---|
289 | patches = [] |
---|
290 | polygon = Polygon(batpts, True) |
---|
291 | patches.append(polygon) |
---|
292 | p2 = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0) |
---|
293 | p2.set_facecolors(['#f1a9a9']) |
---|
294 | |
---|
295 | # Maximum depth set here to -8m (suitable for test case 6 only) |
---|
296 | ax.set_ylim([-8., 6.0]) |
---|
297 | ax.set_xlim([0., 51.0]) |
---|
298 | ax.add_collection(p) |
---|
299 | ax.add_collection(p2) |
---|
300 | ax.plot(xs,ys, '--', color='black', ms=10) |
---|
301 | |
---|
302 | plt.annotate(hour,xy=(2,batmin+0.1*brange)) |
---|
303 | plt.annotate(hour2,xy=(2,batmin+0.05*brange)) |
---|
304 | plt.savefig(wadfr) |
---|