1 | """ Random utilities for reading sww file data and for plotting (in ipython, or |
---|
2 | in scripts) |
---|
3 | |
---|
4 | Functionality of note: |
---|
5 | |
---|
6 | util.get_outputs -- read the data from a single sww file into a single object |
---|
7 | |
---|
8 | util.combine_outputs -- read the data from a list of sww files into a single object |
---|
9 | |
---|
10 | util.near_transect -- for finding the indices of points 'near' to a given |
---|
11 | line, and assigning these points a coordinate along that line. |
---|
12 | This is useful for plotting outputs which are 'almost' along a |
---|
13 | transect (e.g. a channel cross-section) -- see example below |
---|
14 | |
---|
15 | util.sort_sww_filenames -- match sww filenames by a wildcard, and order |
---|
16 | them according to their 'time'. This means that |
---|
17 | they can be stuck together using 'combine_outputs' correctly |
---|
18 | |
---|
19 | util.triangle_areas -- compute the areas of every triangle in a get_outputs object [ must be vertex-based] |
---|
20 | |
---|
21 | util.water_volume -- compute the water volume at every time step in an sww file (needs both vertex and centroid value input). |
---|
22 | |
---|
23 | Here is an example ipython session which uses some of these functions: |
---|
24 | |
---|
25 | > import util |
---|
26 | > from matplotlib import pyplot as pyplot |
---|
27 | > p=util.get_output('myfile.sww',minimum_allowed_height=0.01) |
---|
28 | > p2=util.get_centroids(p,velocity_extrapolation=True) |
---|
29 | > xxx=util.near_transect(p,[95., 85.], [120.,68.],tol=2.) # Could equally well use p2 |
---|
30 | > pyplot.ion() # Interactive plotting |
---|
31 | > pyplot.scatter(xxx[1],p.vel[140,xxx[0]],color='red') # Plot along the transect |
---|
32 | |
---|
33 | FIXME: TODO -- Convert to a single function 'get_output', which can either take a |
---|
34 | single filename, a list of filenames, or a wildcard defining a number of |
---|
35 | filenames, and ensure that in each case, the output will be as desired. |
---|
36 | |
---|
37 | """ |
---|
38 | from Scientific.IO.NetCDF import NetCDFFile |
---|
39 | import numpy |
---|
40 | |
---|
41 | class combine_outputs: |
---|
42 | """ |
---|
43 | Read in a list of filenames, and combine all their outputs into a single object. |
---|
44 | e.g.: |
---|
45 | |
---|
46 | p = util.combine_outputs(['file1.sww', 'file1_time_10000.sww', 'file1_time_20000.sww'], 0.01) |
---|
47 | |
---|
48 | will make an object p which has components p.x,p.y,p.time,p.stage, .... etc, |
---|
49 | where the values of stage / momentum / velocity from the sww files are concatenated as appropriate. |
---|
50 | |
---|
51 | This is nice for interactive interrogation of model outputs, or for sticking together outputs in scripts |
---|
52 | |
---|
53 | WARNING: It is easy to use lots of memory, if the sww files are large. |
---|
54 | |
---|
55 | Note: If you want the centroid values, then you could subsequently use: |
---|
56 | |
---|
57 | p2 = util.get_centroids(p,velocity_extrapolation=False) |
---|
58 | |
---|
59 | which would make an object p2 that is like p, but holds information at centroids |
---|
60 | """ |
---|
61 | def __init__(self, filename_list, minimum_allowed_height=1.0e-03): |
---|
62 | # |
---|
63 | # Go through the sww files in 'filename_list', and combine them into one object. |
---|
64 | # |
---|
65 | |
---|
66 | for i, filename in enumerate(filename_list): |
---|
67 | print i, filename |
---|
68 | # Store output from filename |
---|
69 | p_tmp = get_output(filename, minimum_allowed_height) |
---|
70 | if(i==0): |
---|
71 | # Create self |
---|
72 | p1=p_tmp |
---|
73 | else: |
---|
74 | # Append extra data to self |
---|
75 | # Note that p1.x, p1.y, p1.vols, p1.elev should not change |
---|
76 | assert (p1.x == p_tmp.x).all() |
---|
77 | assert (p1.y == p_tmp.y).all() |
---|
78 | assert (p1.vols ==p_tmp.vols).all() |
---|
79 | p1.time = numpy.append(p1.time, p_tmp.time) |
---|
80 | p1.stage = numpy.append(p1.stage, p_tmp.stage, axis=0) |
---|
81 | p1.xmom = numpy.append(p1.xmom, p_tmp.xmom, axis=0) |
---|
82 | p1.ymom = numpy.append(p1.ymom, p_tmp.ymom, axis=0) |
---|
83 | p1.xvel = numpy.append(p1.xvel, p_tmp.xvel, axis=0) |
---|
84 | p1.yvel = numpy.append(p1.yvel, p_tmp.yvel, axis=0) |
---|
85 | p1.vel = numpy.append(p1.vel, p_tmp.vel, axis=0) |
---|
86 | |
---|
87 | self.x, self.y, self.time, self.vols, self.elev, self.stage, self.xmom, \ |
---|
88 | self.ymom, self.xvel, self.yvel, self.vel, self.minimum_allowed_height = \ |
---|
89 | p1.x, p1.y, p1.time, p1.vols, p1.elev, p1.stage, p1.xmom, p1.ymom, p1.xvel,\ |
---|
90 | p1.yvel, p1.vel, p1.minimum_allowed_height |
---|
91 | |
---|
92 | #################### |
---|
93 | |
---|
94 | def sort_sww_filenames(sww_wildcard): |
---|
95 | # Function to take a 'wildcard' sww filename, |
---|
96 | # and return a list of all filenames of this type, |
---|
97 | # sorted by their time. |
---|
98 | # This can then be used efficiently in 'combine_outputs' |
---|
99 | # if you have many filenames starting with the same pattern |
---|
100 | import glob |
---|
101 | filenames=glob.glob(sww_wildcard) |
---|
102 | |
---|
103 | # Extract time from filenames |
---|
104 | file_time=range(len(filenames)) # Predefine |
---|
105 | |
---|
106 | for i,filename in enumerate(filenames): |
---|
107 | filesplit=filename.rsplit('_time_') |
---|
108 | if(len(filesplit)>1): |
---|
109 | file_time[i]=int(filesplit[1].split('_0.sww')[0]) |
---|
110 | else: |
---|
111 | file_time[i]=0 |
---|
112 | |
---|
113 | name_and_time=zip(file_time,filenames) |
---|
114 | name_and_time.sort() # Sort by file_time |
---|
115 | |
---|
116 | output_times, output_names = zip(*name_and_time) |
---|
117 | |
---|
118 | return list(output_names) |
---|
119 | |
---|
120 | ############## |
---|
121 | |
---|
122 | class get_output: |
---|
123 | """Read in data from an .sww file in a convenient form |
---|
124 | e.g. |
---|
125 | p = util.get_output('channel3.sww', minimum_allowed_height=0.01) |
---|
126 | |
---|
127 | p then contains most relevant information as e.g., p.stage, p.elev, p.xmom, etc |
---|
128 | """ |
---|
129 | def __init__(self, filename, minimum_allowed_height=1.0e-03): |
---|
130 | self.x, self.y, self.time, self.vols, self.stage, \ |
---|
131 | self.elev, self.xmom, self.ymom, \ |
---|
132 | self.xvel, self.yvel, self.vel, self.minimum_allowed_height = \ |
---|
133 | read_output(filename, minimum_allowed_height) |
---|
134 | |
---|
135 | |
---|
136 | def read_output(filename, minimum_allowed_height): |
---|
137 | # Input: The name of an .sww file to read data from, |
---|
138 | # e.g. read_sww('channel3.sww') |
---|
139 | # |
---|
140 | # Purpose: To read the sww file, and output a number of variables as arrays that |
---|
141 | # we can then manipulate (e.g. plot, interrogate) |
---|
142 | # |
---|
143 | # Output: x, y, time, stage, elev, xmom, ymom, xvel, yvel, vel |
---|
144 | |
---|
145 | # Import modules |
---|
146 | |
---|
147 | |
---|
148 | |
---|
149 | # Open ncdf connection |
---|
150 | fid=NetCDFFile(filename) |
---|
151 | |
---|
152 | |
---|
153 | # Read variables |
---|
154 | x=fid.variables['x'] |
---|
155 | x=x.getValue() |
---|
156 | y=fid.variables['y'] |
---|
157 | y=y.getValue() |
---|
158 | |
---|
159 | stage=fid.variables['stage'] |
---|
160 | stage=stage.getValue() |
---|
161 | |
---|
162 | elev=fid.variables['elevation'] |
---|
163 | elev=elev.getValue() |
---|
164 | |
---|
165 | xmom=fid.variables['xmomentum'] |
---|
166 | xmom=xmom.getValue() |
---|
167 | |
---|
168 | ymom=fid.variables['ymomentum'] |
---|
169 | ymom=ymom.getValue() |
---|
170 | |
---|
171 | time=fid.variables['time'] |
---|
172 | time=time.getValue() |
---|
173 | |
---|
174 | vols=fid.variables['volumes'] |
---|
175 | vols=vols.getValue() |
---|
176 | |
---|
177 | |
---|
178 | # Calculate velocity = momentum/depth |
---|
179 | xvel=xmom*0.0 |
---|
180 | yvel=ymom*0.0 |
---|
181 | |
---|
182 | for i in range(xmom.shape[0]): |
---|
183 | xvel[i,:]=xmom[i,:]/(stage[i,:]-elev+1.0e-06)*(stage[i,:]> elev+minimum_allowed_height) |
---|
184 | yvel[i,:]=ymom[i,:]/(stage[i,:]-elev+1.0e-06)*(stage[i,:]> elev+minimum_allowed_height) |
---|
185 | |
---|
186 | vel = (xvel**2+yvel**2)**0.5 |
---|
187 | |
---|
188 | return x, y, time, vols, stage, elev, xmom, ymom, xvel, yvel, vel, minimum_allowed_height |
---|
189 | |
---|
190 | ############## |
---|
191 | |
---|
192 | |
---|
193 | |
---|
194 | class get_centroids: |
---|
195 | """ |
---|
196 | Extract centroid values from the output of get_output. |
---|
197 | e.g. |
---|
198 | p = util.get_output('my_sww.sww', minimum_allowed_height=0.01) # vertex values |
---|
199 | pc=util.get_centroids(p, velocity_extrapolation=True) # centroid values |
---|
200 | """ |
---|
201 | def __init__(self,p, velocity_extrapolation=False): |
---|
202 | self.time, self.x, self.y, self.stage, self.xmom,\ |
---|
203 | self.ymom, self.elev, self.xvel, \ |
---|
204 | self.yvel, self.vel= \ |
---|
205 | get_centroid_values(p, velocity_extrapolation) |
---|
206 | |
---|
207 | |
---|
208 | def get_centroid_values(p, velocity_extrapolation): |
---|
209 | # Input: p is the result of e.g. p=util.get_output('mysww.sww'). See the get_output class defined above |
---|
210 | # Output: Values of x, y, Stage, xmom, ymom, elev, xvel, yvel, vel at centroids |
---|
211 | #import numpy |
---|
212 | |
---|
213 | # Make 3 arrays, each containing one index of a vertex of every triangle. |
---|
214 | l=len(p.vols) |
---|
215 | vols0=numpy.zeros(l, dtype='int') |
---|
216 | vols1=numpy.zeros(l, dtype='int') |
---|
217 | vols2=numpy.zeros(l, dtype='int') |
---|
218 | |
---|
219 | # FIXME: 22/2/12/ - I think this loop is slow, should be able to do this |
---|
220 | # another way |
---|
221 | for i in range(l): |
---|
222 | vols0[i]=p.vols[i][0] |
---|
223 | vols1[i]=p.vols[i][1] |
---|
224 | vols2[i]=p.vols[i][2] |
---|
225 | |
---|
226 | # Then use these to compute centroid averages |
---|
227 | x_cent=(p.x[vols0]+p.x[vols1]+p.x[vols2])/3.0 |
---|
228 | y_cent=(p.y[vols0]+p.y[vols1]+p.y[vols2])/3.0 |
---|
229 | |
---|
230 | stage_cent=(p.stage[:,vols0]+p.stage[:,vols1]+p.stage[:,vols2])/3.0 |
---|
231 | elev_cent=(p.elev[vols0]+p.elev[vols1]+p.elev[vols2])/3.0 |
---|
232 | |
---|
233 | # Here, we need to treat differently the case of momentum extrapolation or |
---|
234 | # velocity extrapolation |
---|
235 | if velocity_extrapolation: |
---|
236 | xvel_cent=(p.xvel[:,vols0]+p.xvel[:,vols1]+p.xvel[:,vols2])/3.0 |
---|
237 | yvel_cent=(p.yvel[:,vols0]+p.yvel[:,vols1]+p.yvel[:,vols2])/3.0 |
---|
238 | |
---|
239 | # Now compute momenta |
---|
240 | xmom_cent=stage_cent*0.0 |
---|
241 | ymom_cent=stage_cent*0.0 |
---|
242 | |
---|
243 | t=len(p.time) |
---|
244 | |
---|
245 | for i in range(t): |
---|
246 | xmom_cent[i,:]=xvel_cent[i,:]*(stage_cent[i,:]-elev_cent+1.0e-06)*\ |
---|
247 | (stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
---|
248 | ymom_cent[i,:]=yvel_cent[i,:]*(stage_cent[i,:]-elev_cent+1.0e-06)*\ |
---|
249 | (stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
---|
250 | |
---|
251 | else: |
---|
252 | xmom_cent=(p.xmom[:,vols0]+p.xmom[:,vols1]+p.xmom[:,vols2])/3.0 |
---|
253 | ymom_cent=(p.ymom[:,vols0]+p.ymom[:,vols1]+p.ymom[:,vols2])/3.0 |
---|
254 | |
---|
255 | # Now compute velocities |
---|
256 | xvel_cent=stage_cent*0.0 |
---|
257 | yvel_cent=stage_cent*0.0 |
---|
258 | |
---|
259 | t=len(p.time) |
---|
260 | |
---|
261 | for i in range(t): |
---|
262 | xvel_cent[i,:]=xmom_cent[i,:]/(stage_cent[i,:]-elev_cent+1.0e-06)*(stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
---|
263 | yvel_cent[i,:]=ymom_cent[i,:]/(stage_cent[i,:]-elev_cent+1.0e-06)*(stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
---|
264 | |
---|
265 | |
---|
266 | # Compute velocity |
---|
267 | vel_cent=(xvel_cent**2 + yvel_cent**2)**0.5 |
---|
268 | |
---|
269 | return p.time, x_cent, y_cent, stage_cent, xmom_cent,\ |
---|
270 | ymom_cent, elev_cent, xvel_cent, yvel_cent, vel_cent |
---|
271 | |
---|
272 | # Make plot of stage over time. |
---|
273 | #pylab.close() |
---|
274 | #pylab.ion() |
---|
275 | #pylab.plot(time, stage[:,1]) |
---|
276 | #for i in range(201): |
---|
277 | # pylab.plot(time,stage[:,i]) |
---|
278 | |
---|
279 | # Momentum should be 0. |
---|
280 | #print 'Momentum max/min is', xmom.max() , xmom.min() |
---|
281 | |
---|
282 | #pylab.gca().set_aspect('equal') |
---|
283 | #pylab.scatter(x,y,c=elev,edgecolors='none') |
---|
284 | #pylab.colorbar() |
---|
285 | # |
---|
286 | #n=xvel.shape[0]-1 |
---|
287 | #pylab.quiver(x,y,xvel[n,:],yvel[n,:]) |
---|
288 | #pylab.savefig('Plot1.png') |
---|
289 | |
---|
290 | def animate_1D(time, var, x, ylab=' '): #, x=range(var.shape[1]), vmin=var.min(), vmax=var.max()): |
---|
291 | # Input: time = one-dimensional time vector; |
---|
292 | # var = array with first dimension = len(time) ; |
---|
293 | # x = (optional) vector width dimension equal to var.shape[1]; |
---|
294 | |
---|
295 | import pylab |
---|
296 | import numpy |
---|
297 | |
---|
298 | |
---|
299 | |
---|
300 | pylab.close() |
---|
301 | pylab.ion() |
---|
302 | |
---|
303 | # Initial plot |
---|
304 | vmin=var.min() |
---|
305 | vmax=var.max() |
---|
306 | line, = pylab.plot( (x.min(), x.max()), (vmin, vmax), 'o') |
---|
307 | |
---|
308 | # Lots of plots |
---|
309 | for i in range(len(time)): |
---|
310 | line.set_xdata(x) |
---|
311 | line.set_ydata(var[i,:]) |
---|
312 | pylab.draw() |
---|
313 | pylab.xlabel('x') |
---|
314 | pylab.ylabel(ylab) |
---|
315 | pylab.title('time = ' + str(time[i])) |
---|
316 | |
---|
317 | return |
---|
318 | |
---|
319 | def near_transect(p, point1, point2, tol=1.): |
---|
320 | # Function to get the indices of points in p less than 'tol' from the line |
---|
321 | # joining (x1,y1), and (x2,y2) |
---|
322 | # p comes from util.get_output('mysww.sww') |
---|
323 | # |
---|
324 | # e.g. |
---|
325 | # import util |
---|
326 | # from matplotlib import pyplot |
---|
327 | # p=util.get_output('merewether_1m.sww',0.01) |
---|
328 | # p2=util.get_centroids(p,velocity_extrapolation=True) |
---|
329 | # #xxx=transect_interpolate.near_transect(p,[95., 85.], [120.,68.],tol=2.) |
---|
330 | # xxx=util.near_transect(p,[95., 85.], [120.,68.],tol=2.) |
---|
331 | # pyplot.scatter(xxx[1],p.vel[140,xxx[0]],color='red') |
---|
332 | |
---|
333 | x1=point1[0] |
---|
334 | y1=point1[1] |
---|
335 | |
---|
336 | x2=point2[0] |
---|
337 | y2=point2[1] |
---|
338 | |
---|
339 | # Find line equation a*x + b*y + c = 0 |
---|
340 | # based on y=gradient*x +intercept |
---|
341 | if x1!=x2: |
---|
342 | gradient= (y2-y1)/(x2-x1) |
---|
343 | intercept = y1 - gradient*x1 |
---|
344 | |
---|
345 | a = -gradient |
---|
346 | b = 1. |
---|
347 | c = -intercept |
---|
348 | else: |
---|
349 | #print 'FIXME: Still need to treat 0 and infinite gradients' |
---|
350 | #assert 0==1 |
---|
351 | a=1. |
---|
352 | b=0. |
---|
353 | c=-x2 |
---|
354 | |
---|
355 | # Distance formula |
---|
356 | inv_denom = 1./(a**2 + b**2)**0.5 |
---|
357 | distp = abs(p.x*a + p.y*b + c)*inv_denom |
---|
358 | |
---|
359 | near_points = (distp<tol).nonzero()[0] |
---|
360 | |
---|
361 | # Now find a 'local' coordinate for the point, projected onto the line |
---|
362 | # g1 = unit vector parallel to the line |
---|
363 | # g2 = vector joining (x1,y1) and (p.x,p.y) |
---|
364 | g1x = x2-x1 |
---|
365 | g1y = y2-y1 |
---|
366 | g1_norm = (g1x**2 + g1y**2)**0.5 |
---|
367 | g1x=g1x/g1_norm |
---|
368 | g1y=g1x/g1_norm |
---|
369 | |
---|
370 | g2x = p.x[near_points] - x1 |
---|
371 | g2y = p.y[near_points] - y1 |
---|
372 | |
---|
373 | # Dot product = projected distance == a local coordinate |
---|
374 | local_coord = g1x*g2x + g1y*g2y |
---|
375 | |
---|
376 | return near_points, local_coord |
---|
377 | |
---|
378 | ######################## |
---|
379 | # TRIANGLE AREAS, WATER VOLUME |
---|
380 | def triangle_areas(p, subset='null'): |
---|
381 | # Compute areas of triangles in p -- assumes p contains vertex information |
---|
382 | # subset = vector of centroid indices to include in the computation. |
---|
383 | |
---|
384 | if(subset=='null'): |
---|
385 | subset=range(len(p.vols[:,0])) |
---|
386 | |
---|
387 | x0=p.x[p.vols[subset,0]] |
---|
388 | x1=p.x[p.vols[subset,1]] |
---|
389 | x2=p.x[p.vols[subset,2]] |
---|
390 | |
---|
391 | y0=p.y[p.vols[subset,0]] |
---|
392 | y1=p.y[p.vols[subset,1]] |
---|
393 | y2=p.y[p.vols[subset,2]] |
---|
394 | |
---|
395 | # Vectors for cross-product |
---|
396 | v1_x=x0-x1 |
---|
397 | v1_y=y0-y1 |
---|
398 | # |
---|
399 | v2_x=x2-x1 |
---|
400 | v2_y=y2-y1 |
---|
401 | # Area |
---|
402 | area=(v1_x*v2_y-v1_y*v2_x)*0.5 |
---|
403 | area=abs(area) |
---|
404 | return area |
---|
405 | |
---|
406 | ### |
---|
407 | |
---|
408 | def water_volume(p,p2, per_unit_area=False, subset='null'): |
---|
409 | # Compute the water volume from p(vertex values) and p2(centroid values) |
---|
410 | |
---|
411 | if(subset=='null'): |
---|
412 | subset=range(len(p2.x)) |
---|
413 | |
---|
414 | l=len(p2.time) |
---|
415 | area=triangle_areas(p, subset=subset) |
---|
416 | |
---|
417 | total_area=area.sum() |
---|
418 | volume=p2.time*0. |
---|
419 | |
---|
420 | # This accounts for how volume is measured in ANUGA |
---|
421 | # Compute in 2 steps to reduce precision error (important sometimes) |
---|
422 | for i in range(l): |
---|
423 | #volume[i]=((p2.stage[i,subset]-p2.elev[subset])*(p2.stage[i,subset]>p2.elev[subset])*area).sum() |
---|
424 | volume[i]=((p2.stage[i,subset])*(p2.stage[i,subset]>p2.elev[subset])*area).sum() |
---|
425 | volume[i]=volume[i]+((-p2.elev[subset])*(p2.stage[i,subset]>p2.elev[subset])*area).sum() |
---|
426 | |
---|
427 | if(per_unit_area): |
---|
428 | volume=volume/total_area |
---|
429 | |
---|
430 | return volume |
---|
431 | |
---|
432 | |
---|
433 | |
---|