source: trunk/anuga_core/source/anuga/utilities/plot_utils.py @ 8780

Last change on this file since 8780 was 8780, checked in by steve, 11 years ago

Some changes to allow netcdf4 use

File size: 14.7 KB
Line 
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"""
38from anuga.file.netcdf import NetCDFFile
39import numpy
40
41class 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
94def 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
122class 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
136def 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
194class 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
208def 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
290def 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
319def 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
380def 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
408def 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
Note: See TracBrowser for help on using the repository browser.