[8399] | 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 | |
---|