| 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 | |
|---|
| 16 | |
|---|
| 17 | Here is an example ipython session which uses some of these functions: |
|---|
| 18 | |
|---|
| 19 | > import util |
|---|
| 20 | > from matplotlib import pyplot as pyplot |
|---|
| 21 | > p=util.get_output('myfile.sww',minimum_allowed_height=0.01) |
|---|
| 22 | > p2=util.get_centroids(p,velocity_extrapolation=True) |
|---|
| 23 | > xxx=util.near_transect(p,[95., 85.], [120.,68.],tol=2.) # Could equally well use p2 |
|---|
| 24 | > pyplot.ion() # Interactive plotting |
|---|
| 25 | > pyplot.scatter(xxx[1],p.vel[140,xxx[0]],color='red') # Plot along the transect |
|---|
| 26 | |
|---|
| 27 | """ |
|---|
| 28 | from Scientific.IO.NetCDF import NetCDFFile |
|---|
| 29 | import numpy |
|---|
| 30 | |
|---|
| 31 | class combine_outputs: |
|---|
| 32 | """ |
|---|
| 33 | Read in a list of filenames, and combine all their outputs into a single object. |
|---|
| 34 | e.g.: |
|---|
| 35 | |
|---|
| 36 | p = util.combine_outputs(['file1.sww', 'file1_time_10000.sww', 'file1_time_20000.sww'], 0.01) |
|---|
| 37 | |
|---|
| 38 | will make an object p which has components p.x,p.y,p.time,p.stage, .... etc, |
|---|
| 39 | where the values of stage / momentum / velocity from the sww files are concatenated as appropriate. |
|---|
| 40 | |
|---|
| 41 | This is nice for interactive interrogation of model outputs, or for sticking together outputs in scripts |
|---|
| 42 | |
|---|
| 43 | WARNING: It is easy to use lots of memory, if the sww files are large. |
|---|
| 44 | |
|---|
| 45 | Note: If you want the centroid values, then you could subsequently use: |
|---|
| 46 | |
|---|
| 47 | p2 = util.get_centroids(p,velocity_extrapolation=False) |
|---|
| 48 | |
|---|
| 49 | which would make an object p2 that is like p, but holds information at centroids |
|---|
| 50 | """ |
|---|
| 51 | def __init__(self, filename_list, minimum_allowed_height=1.0e-03): |
|---|
| 52 | # |
|---|
| 53 | # Go through the sww files in 'filename_list', and combine them into one object. |
|---|
| 54 | # |
|---|
| 55 | |
|---|
| 56 | for i, filename in enumerate(filename_list): |
|---|
| 57 | print i, filename |
|---|
| 58 | # Store output from filename |
|---|
| 59 | p_tmp = get_output(filename, minimum_allowed_height) |
|---|
| 60 | if(i==0): |
|---|
| 61 | # Create self |
|---|
| 62 | p1=p_tmp |
|---|
| 63 | else: |
|---|
| 64 | # Append extra data to self |
|---|
| 65 | # Note that p1.x, p1.y, p1.vols, p1.elev should not change |
|---|
| 66 | assert (p1.x == p_tmp.x).all() |
|---|
| 67 | assert (p1.y == p_tmp.y).all() |
|---|
| 68 | assert (p1.vols ==p_tmp.vols).all() |
|---|
| 69 | p1.time = numpy.append(p1.time, p_tmp.time) |
|---|
| 70 | p1.stage = numpy.append(p1.stage, p_tmp.stage, axis=0) |
|---|
| 71 | p1.xmom = numpy.append(p1.xmom, p_tmp.xmom, axis=0) |
|---|
| 72 | p1.ymom = numpy.append(p1.ymom, p_tmp.ymom, axis=0) |
|---|
| 73 | p1.xvel = numpy.append(p1.xvel, p_tmp.xvel, axis=0) |
|---|
| 74 | p1.yvel = numpy.append(p1.yvel, p_tmp.yvel, axis=0) |
|---|
| 75 | p1.vel = numpy.append(p1.vel, p_tmp.vel, axis=0) |
|---|
| 76 | |
|---|
| 77 | self.x, self.y, self.time, self.vols, self.elev, self.stage, self.xmom, \ |
|---|
| 78 | self.ymom, self.xvel, self.yvel, self.vel, self.minimum_allowed_height = \ |
|---|
| 79 | p1.x, p1.y, p1.time, p1.vols, p1.elev, p1.stage, p1.xmom, p1.ymom, p1.xvel,\ |
|---|
| 80 | p1.yvel, p1.vel, p1.minimum_allowed_height |
|---|
| 81 | |
|---|
| 82 | |
|---|
| 83 | |
|---|
| 84 | class get_output: |
|---|
| 85 | """Read in data from an .sww file in a convenient form |
|---|
| 86 | e.g. |
|---|
| 87 | p = util.get_output('channel3.sww', minimum_allowed_height=0.01) |
|---|
| 88 | |
|---|
| 89 | p then contains most relevant information as e.g., p.stage, p.elev, p.xmom, etc |
|---|
| 90 | """ |
|---|
| 91 | def __init__(self, filename, minimum_allowed_height=1.0e-03): |
|---|
| 92 | self.x, self.y, self.time, self.vols, self.stage, \ |
|---|
| 93 | self.elev, self.xmom, self.ymom, \ |
|---|
| 94 | self.xvel, self.yvel, self.vel, self.minimum_allowed_height = \ |
|---|
| 95 | read_output(filename, minimum_allowed_height) |
|---|
| 96 | |
|---|
| 97 | |
|---|
| 98 | def read_output(filename, minimum_allowed_height): |
|---|
| 99 | # Input: The name of an .sww file to read data from, |
|---|
| 100 | # e.g. read_sww('channel3.sww') |
|---|
| 101 | # |
|---|
| 102 | # Purpose: To read the sww file, and output a number of variables as arrays that |
|---|
| 103 | # we can then manipulate (e.g. plot, interrogate) |
|---|
| 104 | # |
|---|
| 105 | # Output: x, y, time, stage, elev, xmom, ymom, xvel, yvel, vel |
|---|
| 106 | |
|---|
| 107 | # Import modules |
|---|
| 108 | |
|---|
| 109 | |
|---|
| 110 | |
|---|
| 111 | # Open ncdf connection |
|---|
| 112 | fid=NetCDFFile(filename) |
|---|
| 113 | |
|---|
| 114 | |
|---|
| 115 | # Read variables |
|---|
| 116 | x=fid.variables['x'] |
|---|
| 117 | x=x.getValue() |
|---|
| 118 | y=fid.variables['y'] |
|---|
| 119 | y=y.getValue() |
|---|
| 120 | |
|---|
| 121 | stage=fid.variables['stage'] |
|---|
| 122 | stage=stage.getValue() |
|---|
| 123 | |
|---|
| 124 | elev=fid.variables['elevation'] |
|---|
| 125 | elev=elev.getValue() |
|---|
| 126 | |
|---|
| 127 | xmom=fid.variables['xmomentum'] |
|---|
| 128 | xmom=xmom.getValue() |
|---|
| 129 | |
|---|
| 130 | ymom=fid.variables['ymomentum'] |
|---|
| 131 | ymom=ymom.getValue() |
|---|
| 132 | |
|---|
| 133 | time=fid.variables['time'] |
|---|
| 134 | time=time.getValue() |
|---|
| 135 | |
|---|
| 136 | vols=fid.variables['volumes'] |
|---|
| 137 | vols=vols.getValue() |
|---|
| 138 | |
|---|
| 139 | |
|---|
| 140 | # Calculate velocity = momentum/depth |
|---|
| 141 | xvel=xmom*0.0 |
|---|
| 142 | yvel=ymom*0.0 |
|---|
| 143 | |
|---|
| 144 | for i in range(xmom.shape[0]): |
|---|
| 145 | xvel[i,:]=xmom[i,:]/(stage[i,:]-elev+1.0e-06)*(stage[i,:]> elev+minimum_allowed_height) |
|---|
| 146 | yvel[i,:]=ymom[i,:]/(stage[i,:]-elev+1.0e-06)*(stage[i,:]> elev+minimum_allowed_height) |
|---|
| 147 | |
|---|
| 148 | vel = (xvel**2+yvel**2)**0.5 |
|---|
| 149 | |
|---|
| 150 | return x, y, time, vols, stage, elev, xmom, ymom, xvel, yvel, vel, minimum_allowed_height |
|---|
| 151 | |
|---|
| 152 | ############## |
|---|
| 153 | |
|---|
| 154 | |
|---|
| 155 | |
|---|
| 156 | class get_centroids: |
|---|
| 157 | """ |
|---|
| 158 | Extract centroid values from the output of get_output. |
|---|
| 159 | e.g. |
|---|
| 160 | p = util.get_output('my_sww.sww', minimum_allowed_height=0.01) # vertex values |
|---|
| 161 | pc=util.get_centroids(p, velocity_extrapolation=True) # centroid values |
|---|
| 162 | """ |
|---|
| 163 | def __init__(self,p, velocity_extrapolation=False): |
|---|
| 164 | self.time, self.x, self.y, self.stage, self.xmom,\ |
|---|
| 165 | self.ymom, self.elev, self.xvel, \ |
|---|
| 166 | self.yvel, self.vel= \ |
|---|
| 167 | get_centroid_values(p, velocity_extrapolation) |
|---|
| 168 | |
|---|
| 169 | |
|---|
| 170 | def get_centroid_values(p, velocity_extrapolation): |
|---|
| 171 | # Input: p is the result of e.g. p=util.get_output('mysww.sww'). See the get_output class defined above |
|---|
| 172 | # Output: Values of x, y, Stage, xmom, ymom, elev, xvel, yvel, vel at centroids |
|---|
| 173 | #import numpy |
|---|
| 174 | |
|---|
| 175 | # Make 3 arrays, each containing one index of a vertex of every triangle. |
|---|
| 176 | l=len(p.vols) |
|---|
| 177 | vols0=numpy.zeros(l, dtype='int') |
|---|
| 178 | vols1=numpy.zeros(l, dtype='int') |
|---|
| 179 | vols2=numpy.zeros(l, dtype='int') |
|---|
| 180 | |
|---|
| 181 | # FIXME: 22/2/12/ - I think this loop is slow, should be able to do this |
|---|
| 182 | # another way |
|---|
| 183 | for i in range(l): |
|---|
| 184 | vols0[i]=p.vols[i][0] |
|---|
| 185 | vols1[i]=p.vols[i][1] |
|---|
| 186 | vols2[i]=p.vols[i][2] |
|---|
| 187 | |
|---|
| 188 | # Then use these to compute centroid averages |
|---|
| 189 | x_cent=(p.x[vols0]+p.x[vols1]+p.x[vols2])/3.0 |
|---|
| 190 | y_cent=(p.y[vols0]+p.y[vols1]+p.y[vols2])/3.0 |
|---|
| 191 | |
|---|
| 192 | stage_cent=(p.stage[:,vols0]+p.stage[:,vols1]+p.stage[:,vols2])/3.0 |
|---|
| 193 | elev_cent=(p.elev[vols0]+p.elev[vols1]+p.elev[vols2])/3.0 |
|---|
| 194 | |
|---|
| 195 | # Here, we need to treat differently the case of momentum extrapolation or |
|---|
| 196 | # velocity extrapolation |
|---|
| 197 | if velocity_extrapolation: |
|---|
| 198 | xvel_cent=(p.xvel[:,vols0]+p.xvel[:,vols1]+p.xvel[:,vols2])/3.0 |
|---|
| 199 | yvel_cent=(p.yvel[:,vols0]+p.yvel[:,vols1]+p.yvel[:,vols2])/3.0 |
|---|
| 200 | |
|---|
| 201 | # Now compute momenta |
|---|
| 202 | xmom_cent=stage_cent*0.0 |
|---|
| 203 | ymom_cent=stage_cent*0.0 |
|---|
| 204 | |
|---|
| 205 | t=len(p.time) |
|---|
| 206 | |
|---|
| 207 | for i in range(t): |
|---|
| 208 | xmom_cent[i,:]=xvel_cent[i,:]*(stage_cent[i,:]-elev_cent+1.0e-06)*\ |
|---|
| 209 | (stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
|---|
| 210 | ymom_cent[i,:]=yvel_cent[i,:]*(stage_cent[i,:]-elev_cent+1.0e-06)*\ |
|---|
| 211 | (stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
|---|
| 212 | |
|---|
| 213 | else: |
|---|
| 214 | xmom_cent=(p.xmom[:,vols0]+p.xmom[:,vols1]+p.xmom[:,vols2])/3.0 |
|---|
| 215 | ymom_cent=(p.ymom[:,vols0]+p.ymom[:,vols1]+p.ymom[:,vols2])/3.0 |
|---|
| 216 | |
|---|
| 217 | # Now compute velocities |
|---|
| 218 | xvel_cent=stage_cent*0.0 |
|---|
| 219 | yvel_cent=stage_cent*0.0 |
|---|
| 220 | |
|---|
| 221 | t=len(p.time) |
|---|
| 222 | |
|---|
| 223 | for i in range(t): |
|---|
| 224 | xvel_cent[i,:]=xmom_cent[i,:]/(stage_cent[i,:]-elev_cent+1.0e-06)*(stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
|---|
| 225 | yvel_cent[i,:]=ymom_cent[i,:]/(stage_cent[i,:]-elev_cent+1.0e-06)*(stage_cent[i,:]>elev_cent+p.minimum_allowed_height) |
|---|
| 226 | |
|---|
| 227 | |
|---|
| 228 | # Compute velocity |
|---|
| 229 | vel_cent=(xvel_cent**2 + yvel_cent**2)**0.5 |
|---|
| 230 | |
|---|
| 231 | return p.time, x_cent, y_cent, stage_cent, xmom_cent,\ |
|---|
| 232 | ymom_cent, elev_cent, xvel_cent, yvel_cent, vel_cent |
|---|
| 233 | |
|---|
| 234 | # Make plot of stage over time. |
|---|
| 235 | #pylab.close() |
|---|
| 236 | #pylab.ion() |
|---|
| 237 | #pylab.plot(time, stage[:,1]) |
|---|
| 238 | #for i in range(201): |
|---|
| 239 | # pylab.plot(time,stage[:,i]) |
|---|
| 240 | |
|---|
| 241 | # Momentum should be 0. |
|---|
| 242 | #print 'Momentum max/min is', xmom.max() , xmom.min() |
|---|
| 243 | |
|---|
| 244 | #pylab.gca().set_aspect('equal') |
|---|
| 245 | #pylab.scatter(x,y,c=elev,edgecolors='none') |
|---|
| 246 | #pylab.colorbar() |
|---|
| 247 | # |
|---|
| 248 | #n=xvel.shape[0]-1 |
|---|
| 249 | #pylab.quiver(x,y,xvel[n,:],yvel[n,:]) |
|---|
| 250 | #pylab.savefig('Plot1.png') |
|---|
| 251 | |
|---|
| 252 | def animate_1D(time, var, x, ylab=' '): #, x=range(var.shape[1]), vmin=var.min(), vmax=var.max()): |
|---|
| 253 | # Input: time = one-dimensional time vector; |
|---|
| 254 | # var = array with first dimension = len(time) ; |
|---|
| 255 | # x = (optional) vector width dimension equal to var.shape[1]; |
|---|
| 256 | |
|---|
| 257 | import pylab |
|---|
| 258 | import numpy |
|---|
| 259 | |
|---|
| 260 | |
|---|
| 261 | |
|---|
| 262 | pylab.close() |
|---|
| 263 | pylab.ion() |
|---|
| 264 | |
|---|
| 265 | # Initial plot |
|---|
| 266 | vmin=var.min() |
|---|
| 267 | vmax=var.max() |
|---|
| 268 | line, = pylab.plot( (x.min(), x.max()), (vmin, vmax), 'o') |
|---|
| 269 | |
|---|
| 270 | # Lots of plots |
|---|
| 271 | for i in range(len(time)): |
|---|
| 272 | line.set_xdata(x) |
|---|
| 273 | line.set_ydata(var[i,:]) |
|---|
| 274 | pylab.draw() |
|---|
| 275 | pylab.xlabel('x') |
|---|
| 276 | pylab.ylabel(ylab) |
|---|
| 277 | pylab.title('time = ' + str(time[i])) |
|---|
| 278 | |
|---|
| 279 | return |
|---|
| 280 | |
|---|
| 281 | def near_transect(p, point1, point2, tol=1.): |
|---|
| 282 | # Function to get the indices of points in p less than 'tol' from the line |
|---|
| 283 | # joining (x1,y1), and (x2,y2) |
|---|
| 284 | # p comes from util.get_output('mysww.sww') |
|---|
| 285 | # |
|---|
| 286 | # e.g. |
|---|
| 287 | # import util |
|---|
| 288 | # from matplotlib import pyplot |
|---|
| 289 | # p=util.get_output('merewether_1m.sww',0.01) |
|---|
| 290 | # p2=util.get_centroids(p,velocity_extrapolation=True) |
|---|
| 291 | # #xxx=transect_interpolate.near_transect(p,[95., 85.], [120.,68.],tol=2.) |
|---|
| 292 | # xxx=util.near_transect(p,[95., 85.], [120.,68.],tol=2.) |
|---|
| 293 | # pyplot.scatter(xxx[1],p.vel[140,xxx[0]],color='red') |
|---|
| 294 | |
|---|
| 295 | x1=point1[0] |
|---|
| 296 | y1=point1[1] |
|---|
| 297 | |
|---|
| 298 | x2=point2[0] |
|---|
| 299 | y2=point2[1] |
|---|
| 300 | |
|---|
| 301 | # Find line equation a*x + b*y + c = 0 |
|---|
| 302 | # based on y=gradient*x +intercept |
|---|
| 303 | if x1!=x2: |
|---|
| 304 | gradient= (y2-y1)/(x2-x1) |
|---|
| 305 | intercept = y1 - gradient*x1 |
|---|
| 306 | |
|---|
| 307 | a = -gradient |
|---|
| 308 | b = 1. |
|---|
| 309 | c = -intercept |
|---|
| 310 | else: |
|---|
| 311 | #print 'FIXME: Still need to treat 0 and infinite gradients' |
|---|
| 312 | #assert 0==1 |
|---|
| 313 | a=1. |
|---|
| 314 | b=0. |
|---|
| 315 | c=-x2 |
|---|
| 316 | |
|---|
| 317 | # Distance formula |
|---|
| 318 | inv_denom = 1./(a**2 + b**2)**0.5 |
|---|
| 319 | distp = abs(p.x*a + p.y*b + c)*inv_denom |
|---|
| 320 | |
|---|
| 321 | near_points = (distp<tol).nonzero()[0] |
|---|
| 322 | |
|---|
| 323 | # Now find a 'local' coordinate for the point, projected onto the line |
|---|
| 324 | # g1 = unit vector parallel to the line |
|---|
| 325 | # g2 = vector joining (x1,y1) and (p.x,p.y) |
|---|
| 326 | g1x = x2-x1 |
|---|
| 327 | g1y = y2-y1 |
|---|
| 328 | g1_norm = (g1x**2 + g1y**2)**0.5 |
|---|
| 329 | g1x=g1x/g1_norm |
|---|
| 330 | g1y=g1x/g1_norm |
|---|
| 331 | |
|---|
| 332 | g2x = p.x[near_points] - x1 |
|---|
| 333 | g2y = p.y[near_points] - y1 |
|---|
| 334 | |
|---|
| 335 | # Dot product = projected distance == a local coordinate |
|---|
| 336 | local_coord = g1x*g2x + g1y*g2y |
|---|
| 337 | |
|---|
| 338 | return near_points, local_coord |
|---|
