source: trunk/anuga_work/development/gareth/experimental/bal_and/plot_utils.py @ 9056

Last change on this file since 9056 was 9042, checked in by davies, 11 years ago

Adding code to make georeferenced rasters from sww to plot_utils
File size: 24.2 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    #FIXME: 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.height = numpy.append(p1.height, p_tmp.height, axis=0)
82                p1.xmom = numpy.append(p1.xmom, p_tmp.xmom, axis=0)
83                p1.ymom = numpy.append(p1.ymom, p_tmp.ymom, axis=0)
84                p1.xvel = numpy.append(p1.xvel, p_tmp.xvel, axis=0)
85                p1.yvel = numpy.append(p1.yvel, p_tmp.yvel, axis=0)
86                p1.vel = numpy.append(p1.vel, p_tmp.vel, axis=0)
87
88        self.x, self.y, self.time, self.vols, self.elev, self.stage, self.xmom, \
89                self.ymom, self.xvel, self.yvel, self.vel, self.minimum_allowed_height = \
90                p1.x, p1.y, p1.time, p1.vols, p1.elev, p1.stage, p1.xmom, p1.ymom, p1.xvel,\
91                p1.yvel, p1.vel, p1.minimum_allowed_height
92
93####################
94
95def sort_sww_filenames(sww_wildcard):
96    # Function to take a 'wildcard' sww filename,
97    # and return a list of all filenames of this type,
98    # sorted by their time.
99    # This can then be used efficiently in 'combine_outputs'
100    # if you have many filenames starting with the same pattern
101    import glob
102    filenames=glob.glob(sww_wildcard)
103   
104    # Extract time from filenames
105    file_time=range(len(filenames)) # Predefine
106     
107    for i,filename in enumerate(filenames):
108        filesplit=filename.rsplit('_time_')
109        if(len(filesplit)>1):
110            file_time[i]=int(filesplit[1].split('_0.sww')[0])
111        else:
112            file_time[i]=0         
113   
114    name_and_time=zip(file_time,filenames)
115    name_and_time.sort() # Sort by file_time
116   
117    output_times, output_names = zip(*name_and_time)
118   
119    return list(output_names)
120
121##############
122
123class get_output:
124    """Read in data from an .sww file in a convenient form
125       e.g.
126        p = util.get_output('channel3.sww', minimum_allowed_height=0.01)
127       
128       p then contains most relevant information as e.g., p.stage, p.elev, p.xmom, etc
129    """
130    def __init__(self, filename, minimum_allowed_height=1.0e-03):
131        self.x, self.y, self.time, self.vols, self.stage, \
132                self.height, self.elev, self.xmom, self.ymom, \
133                self.xvel, self.yvel, self.vel, self.minimum_allowed_height = \
134                read_output(filename, minimum_allowed_height)
135        self.filename=filename
136
137
138def read_output(filename, minimum_allowed_height):
139    # Input: The name of an .sww file to read data from,
140    #                    e.g. read_sww('channel3.sww')
141    #
142    # Purpose: To read the sww file, and output a number of variables as arrays that
143    #          we can then manipulate (e.g. plot, interrogate)
144    #
145    # Output: x, y, time, stage, height, elev, xmom, ymom, xvel, yvel, vel
146    #         x,y are only stored at one time
147    #         elevation may be stored at one or multiple times
148    #         everything else is stored every time step for vertices
149
150    # Import modules
151
152    # Open ncdf connection
153    fid=NetCDFFile(filename)
154
155
156    # Read variables
157    x=fid.variables['x'][:]
158   
159    y=fid.variables['y'][:]
160
161    stage=fid.variables['stage'][:]
162   
163    elev=fid.variables['elevation'][:]
164
165    if(fid.variables.has_key('height')):
166        height=fid.variables['height'][:]
167    else:
168        # Back calculate height if it is not stored
169        height=fid.variables['stage'][:]
170        if(len(stage.shape)==len(elev.shape)):
171            for i in range(stage.shape[0]):
172                height[i,:]=stage[i,:]-elev[i,:]
173        else:
174            for i in range(stage.shape[0]):
175                height[i,:]=stage[i,:]-elev
176
177
178    xmom=fid.variables['xmomentum'][:]
179
180    ymom=fid.variables['ymomentum'][:]
181
182    time=fid.variables['time'][:]
183
184    vols=fid.variables['volumes'][:]
185
186
187    # Calculate velocity = momentum/depth
188    xvel=xmom*0.0
189    yvel=ymom*0.0
190
191    for i in range(xmom.shape[0]):
192        xvel[i,:]=xmom[i,:]/(height[i,:]+1.0e-06)*(height[i,:]>minimum_allowed_height)
193        yvel[i,:]=ymom[i,:]/(height[i,:]+1.0e-06)*(height[i,:]>minimum_allowed_height)
194
195    vel = (xvel**2+yvel**2)**0.5
196
197    return x, y, time, vols, stage, height, elev, xmom, ymom, xvel, yvel, vel, minimum_allowed_height
198
199##############
200
201
202class get_centroids:
203    """
204    Extract centroid values from the output of get_output. 
205    e.g.
206        p = util.get_output('my_sww.sww', minimum_allowed_height=0.01) # vertex values
207        pc=util.get_centroids(p, velocity_extrapolation=True) # centroid values
208
209    NOTE: elevation is only stored once in the output, even if it was stored every timestep
210         This is done because presently centroid elevations do not change over time.
211    """
212    def __init__(self,p, velocity_extrapolation=False):
213       
214        self.time, self.x, self.y, self.stage, self.xmom,\
215             self.ymom, self.height, self.elev, self.xvel, \
216             self.yvel, self.vel= \
217             get_centroid_values(p, velocity_extrapolation)
218                                 
219
220def get_centroid_values(p, velocity_extrapolation):
221    # Input: p is the result of e.g. p=util.get_output('mysww.sww'). See the get_output class defined above
222    # Output: Values of x, y, Stage, xmom, ymom, elev, xvel, yvel, vel at centroids
223    #import numpy
224   
225    # Make 3 arrays, each containing one index of a vertex of every triangle.
226    l=len(p.vols)
227    #vols0=numpy.zeros(l, dtype='int')
228    #vols1=numpy.zeros(l, dtype='int')
229    #vols2=numpy.zeros(l, dtype='int')
230
231    # FIXME: 22/2/12/ - I think this loop is slow, should be able to do this
232    # another way
233#    for i in range(l):
234#        vols0[i]=p.vols[i][0]
235#        vols1[i]=p.vols[i][1]
236#        vols2[i]=p.vols[i][2]
237
238
239
240    vols0=p.vols[:,0]
241    vols1=p.vols[:,1]
242    vols2=p.vols[:,2]
243
244    #print vols0.shape
245    #print p.vols.shape
246
247    # Then use these to compute centroid averages
248    x_cent=(p.x[vols0]+p.x[vols1]+p.x[vols2])/3.0
249    y_cent=(p.y[vols0]+p.y[vols1]+p.y[vols2])/3.0
250       
251    fid=NetCDFFile(p.filename)
252    if(fid.variables.has_key('stage_c')==False):
253
254        stage_cent=(p.stage[:,vols0]+p.stage[:,vols1]+p.stage[:,vols2])/3.0
255        height_cent=(p.height[:,vols0]+p.height[:,vols1]+p.height[:,vols2])/3.0
256        #elev_cent=(p.elev[:,vols0]+p.elev[:,vols1]+p.elev[:,vols2])/3.0
257        # Only store elevation centroid once (since it doesn't change)
258        if(len(p.elev.shape)==2):
259            elev_cent=(p.elev[0,vols0]+p.elev[0,vols1]+p.elev[0,vols2])/3.0
260        else:
261            elev_cent=(p.elev[vols0]+p.elev[vols1]+p.elev[vols2])/3.0
262
263        # Here, we need to treat differently the case of momentum extrapolation or
264        # velocity extrapolation
265        if velocity_extrapolation:
266            xvel_cent=(p.xvel[:,vols0]+p.xvel[:,vols1]+p.xvel[:,vols2])/3.0
267            yvel_cent=(p.yvel[:,vols0]+p.yvel[:,vols1]+p.yvel[:,vols2])/3.0
268           
269            # Now compute momenta
270            xmom_cent=stage_cent*0.0
271            ymom_cent=stage_cent*0.0
272
273            t=len(p.time)
274
275            for i in range(t):
276                xmom_cent[i,:]=xvel_cent[i,:]*(height_cent[i,:]+1e-06)*\
277                                                (height_cent[i,:]>p.minimum_allowed_height)
278                ymom_cent[i,:]=yvel_cent[i,:]*(height_cent[i,:]+1e-06)*\
279                                                (height_cent[i,:]>p.minimum_allowed_height)
280        else:
281            xmom_cent=(p.xmom[:,vols0]+p.xmom[:,vols1]+p.xmom[:,vols2])/3.0
282            ymom_cent=(p.ymom[:,vols0]+p.ymom[:,vols1]+p.ymom[:,vols2])/3.0
283
284            # Now compute velocities
285            xvel_cent=stage_cent*0.0
286            yvel_cent=stage_cent*0.0
287
288            t=len(p.time)
289
290            for i in range(t):
291                xvel_cent[i,:]=xmom_cent[i,:]/(height_cent[i,:]+1.0e-06)*(height_cent[i,:]>p.minimum_allowed_height)
292                yvel_cent[i,:]=ymom_cent[i,:]/(height_cent[i,:]+1.0e-06)*(height_cent[i,:]>p.minimum_allowed_height)
293   
294    else:
295        # Get centroid values from file
296        print 'Reading centroids from file'
297        stage_cent=fid.variables['stage_c'][:]
298        height_cent=fid.variables['height_c'][:]
299        elev_cent=fid.variables['elevation_c'][:]
300        if(len(elev_cent.shape)==2):
301            elev_cent=elev_cent[0,:] # Only store elevation centroid once -- since it is constant
302        xmom_cent=fid.variables['xmomentum_c'][:]*(height_cent>p.minimum_allowed_height)
303        ymom_cent=fid.variables['ymomentum_c'][:]*(height_cent>p.minimum_allowed_height)
304        xvel_cent=xmom_cent/(height_cent+1.0e-06)
305        yvel_cent=ymom_cent/(height_cent+1.0e-06)
306       
307   
308    # Compute velocity
309    vel_cent=(xvel_cent**2 + yvel_cent**2)**0.5
310
311    return p.time, x_cent, y_cent, stage_cent, xmom_cent,\
312             ymom_cent, height_cent, elev_cent, xvel_cent, yvel_cent, vel_cent
313
314def animate_1D(time, var, x, ylab=' '): #, x=range(var.shape[1]), vmin=var.min(), vmax=var.max()):
315    # Input: time = one-dimensional time vector;
316    #        var =  array with first dimension = len(time) ;
317    #        x = (optional) vector width dimension equal to var.shape[1];
318   
319    import pylab
320    import numpy
321   
322   
323
324    pylab.close()
325    pylab.ion()
326
327    # Initial plot
328    vmin=var.min()
329    vmax=var.max()
330    line, = pylab.plot( (x.min(), x.max()), (vmin, vmax), 'o')
331
332    # Lots of plots
333    for i in range(len(time)):
334        line.set_xdata(x)
335        line.set_ydata(var[i,:])
336        pylab.draw()
337        pylab.xlabel('x')
338        pylab.ylabel(ylab)
339        pylab.title('time = ' + str(time[i]))
340   
341    return
342
343def near_transect(p, point1, point2, tol=1.):
344    # Function to get the indices of points in p less than 'tol' from the line
345    # joining (x1,y1), and (x2,y2)
346    # p comes from util.get_output('mysww.sww')
347    #
348    # e.g.
349    # import util
350    # from matplotlib import pyplot
351    # p=util.get_output('merewether_1m.sww',0.01)
352    # p2=util.get_centroids(p,velocity_extrapolation=True)
353    # #xxx=transect_interpolate.near_transect(p,[95., 85.], [120.,68.],tol=2.)
354    # xxx=util.near_transect(p,[95., 85.], [120.,68.],tol=2.)
355    # pyplot.scatter(xxx[1],p.vel[140,xxx[0]],color='red')
356
357    x1=point1[0]
358    y1=point1[1]
359
360    x2=point2[0]
361    y2=point2[1]
362
363    # Find line equation a*x + b*y + c = 0
364    # based on y=gradient*x +intercept
365    if x1!=x2:
366        gradient= (y2-y1)/(x2-x1)
367        intercept = y1 - gradient*x1
368
369        a = -gradient
370        b = 1.
371        c = -intercept
372    else:
373        #print 'FIXME: Still need to treat 0 and infinite gradients'
374        #assert 0==1
375        a=1.
376        b=0.
377        c=-x2
378
379    # Distance formula
380    inv_denom = 1./(a**2 + b**2)**0.5
381    distp = abs(p.x*a + p.y*b + c)*inv_denom
382
383    near_points = (distp<tol).nonzero()[0]
384   
385    # Now find a 'local' coordinate for the point, projected onto the line
386    # g1 = unit vector parallel to the line
387    # g2 = vector joining (x1,y1) and (p.x,p.y)
388    g1x = x2-x1
389    g1y = y2-y1
390    g1_norm = (g1x**2 + g1y**2)**0.5
391    g1x=g1x/g1_norm
392    g1y=g1x/g1_norm
393
394    g2x = p.x[near_points] - x1
395    g2y = p.y[near_points] - y1
396   
397    # Dot product = projected distance == a local coordinate
398    local_coord = g1x*g2x + g1y*g2y
399
400    return near_points, local_coord
401
402########################
403# TRIANGLE AREAS, WATER VOLUME
404def triangle_areas(p, subset=None):
405    # Compute areas of triangles in p -- assumes p contains vertex information
406    # subset = vector of centroid indices to include in the computation.
407
408    if(subset is None):
409        subset=range(len(p.vols[:,0]))
410   
411    x0=p.x[p.vols[subset,0]]
412    x1=p.x[p.vols[subset,1]]
413    x2=p.x[p.vols[subset,2]]
414   
415    y0=p.y[p.vols[subset,0]]
416    y1=p.y[p.vols[subset,1]]
417    y2=p.y[p.vols[subset,2]]
418   
419    # Vectors for cross-product
420    v1_x=x0-x1
421    v1_y=y0-y1
422    #
423    v2_x=x2-x1
424    v2_y=y2-y1
425    # Area
426    area=(v1_x*v2_y-v1_y*v2_x)*0.5
427    area=abs(area)
428    return area
429
430###
431
432#def water_volume(p,p2, per_unit_area=False, subset=None):
433#    # Compute the water volume from p(vertex values) and p2(centroid values)
434#
435#    if(subset is None):
436#        subset=range(len(p2.x))
437#
438#    l=len(p2.time)
439#    area=triangle_areas(p, subset=subset)
440#   
441#    total_area=area.sum()
442#    volume=p2.time*0.
443#   
444#    # This accounts for how volume is measured in ANUGA
445#    # Compute in 2 steps to reduce precision error (important sometimes)
446#    for i in range(l):
447#        #volume[i]=((p2.stage[i,subset]-p2.elev[subset])*(p2.stage[i,subset]>p2.elev[subset])*area).sum()
448#        volume[i]=((p2.stage[i,subset])*(p2.stage[i,subset]>p2.elev[subset])*area).sum()
449#        volume[i]=volume[i]+((-p2.elev[subset])*(p2.stage[i,subset]>p2.elev[subset])*area).sum()
450#   
451#    if(per_unit_area):
452#        volume=volume/total_area
453#   
454#    return volume
455
456
457def get_triangle_containing_point(p,point):
458
459    V = p.vols
460
461    x = p.x
462    y = p.y
463
464    l = len(x)
465
466    from anuga.geometry.polygon import is_outside_polygon,is_inside_polygon
467
468    # FIXME: Horrible brute force
469    for i in xrange(l):
470        i0 = V[i,0]
471        i1 = V[i,1]
472        i2 = V[i,2]
473        poly = [ [x[i0], y[i0]], [x[i1], y[i1]], [x[i2], y[i2]] ]
474
475        if is_inside_polygon(point, poly, closed=True):
476            return i
477
478    msg = 'Point %s not found within a triangle' %str(point)
479    raise Exception(msg)
480
481
482def get_extent(p):
483
484    import numpy
485
486    x_min = numpy.min(p.x)
487    x_max = numpy.max(p.x)
488
489    y_min = numpy.min(p.y)
490    y_max = numpy.max(p.y)
491
492    return x_min, x_max, y_min, y_max
493
494
495
496## FOLLOWING FUNCTION MODIFIED FROM STACK OVERFLOW
497def make_grid(data, lats, lons, fileName, EPSG_CODE=None):
498    """
499        Convert data,lats,lons to a georeferenced raster tif
500        INPUT: data -- array with desired raster cell values
501               lats -- 1d array with 'latitude' or 'y' range
502               lons -- 1D array with 'longitude' or 'x' range
503               fileName -- name of file to write to
504               EPSG_CODE -- Integer code with projection information in EPSG format
505    """
506    try:
507        import gdal
508        import osr
509    except:
510        raise Exception, 'Cannot find gdal and/or osr python modules'
511
512    #data = scipy.random.rand(5,6)
513    #lats = scipy.array([-5.5, -5.0, -4.5, -4.0, -3.5])
514    #lons = scipy.array([135.0, 135.5, 136.0, 136.5, 137.0, 137.5])
515    #import pdb
516    #pdb.set_trace()
517    xres = lons[1] - lons[0]
518    yres = lats[1] - lats[0]
519
520    ysize = len(lats)
521    xsize = len(lons)
522
523    # Assume data/lats/longs refer to cell centres, as per gdal
524    #ulx = lons[0] - (xres / 2.)
525    #uly = lats[-1] - (yres / 2.)
526    llx = lons[0] - (xres / 2.)
527    lly = lats[0] - (yres / 2.)
528
529    driver = gdal.GetDriverByName('GTiff')
530    ds = driver.Create(fileName, xsize, ysize, 1, gdal.GDT_Float32)
531
532    # this assumes the projection is Geographic lat/lon WGS 84
533    srs = osr.SpatialReference()
534    srs.ImportFromEPSG(EPSG_CODE)
535    ds.SetProjection(srs.ExportToWkt())
536
537    #gt = [ulx, xres, 0, uly, 0, yres ]
538    gt = [llx, xres, 0, lly, 0, yres ]
539    ds.SetGeoTransform(gt)
540
541    outband = ds.GetRasterBand(1)
542    outband.WriteArray(data)
543
544    ds = None
545    return
546
547##################################################################################
548
549def Make_Geotif(swwFile, output_quantity='depth',
550             myTimeStep=1, CellSize=5.0, 
551             lower_left=None, upper_right=None,
552             EPSG_CODE=None, 
553             velocity_extrapolation=True,
554             min_allowed_height=1.0e-05,
555             output_dir='TIFS',
556             verbose=False):
557    """
558        Make a georeferenced tif by nearest-neighbour interpolation of sww file outputs.
559        INPUTS: swwFile -- name of sww file
560                output_quantity -- quantity to plot ('depth' or 'velocity' or 'stage' or 'elevation')
561                myTimeStep -- time-index of swwFile to plot, or 'last', or 'max'
562                CellSize -- approximate pixel size for output raster [adapted to fit lower_left / upper_right]
563                lower_left -- [x0,y0] of lower left corner. If None, use extent of swwFile.
564                upper_right -- [x1,y1] of upper right corner. If None, use extent of swwFile.
565                EPSG_CODE -- Projection information as an integer EPSG code (e.g. 3123 for PRS92 Zone 3).
566                             Google for info on EPSG Codes
567                velocity_extrapolation -- Compute velocity assuming the code extrapolates with velocity (instead of momentum)?
568                min_allowed_height -- Minimum allowed height from ANUGA
569                output_dir -- Write outputs to this directory
570               
571    """
572    try:
573        import gdal
574        import osr
575        import scipy.io
576        import scipy.interpolate
577        import anuga
578        import os
579    except:
580        raise Exception, 'Required modules not installed for mkgeotif'
581
582
583    if(EPSG_CODE is None):
584        raise Exception, 'Must specify an integer EPSG_CODE describing the file projection'
585    ###
586    ## INPUTS
587    ###
588    #swwFile='taguig_rainfall_ondoy.sww'
589    #myTimeStep=40 # Index of time-slice in sWW to plot
590    ## Plot extents
591    #lower_left=None #[x0,y0]
592    #upper_right=None #[x1,y1]
593    ## Projection as an epsg code
594    #EPSG_CODE=3123
595    ## CellSize for output raster
596    #CellSize=5.0 # APPROXIMATE CellSize
597    ## ANUGA Parameters
598    #min_allowed_height=1.0e-05
599    #velocity_extrapolation=True
600
601    #output_quantity='depth' # 'velocity' 'depth' 'stage'
602    #output_dir='TEST'
603
604    ##
605    # END INPUTS
606    ##
607
608    # Make output_dir
609    try:
610        os.mkdir(output_dir)
611    except:
612        pass
613
614    # Read in ANUGA outputs
615    # FIXME: It would be good to support reading of data subsets
616    if(verbose):
617        print 'Reading sww File ...'
618    p=util.get_output(swwFile,min_allowed_height)
619    p2=util.get_centroids(p,velocity_extrapolation)
620    swwIn=scipy.io.netcdf_file(swwFile)
621    xllcorner=swwIn.xllcorner
622    yllcorner=swwIn.yllcorner
623
624    if(myTimeStep=='last'):
625        myTimeStep=len(p.time)-1
626   
627
628    if(verbose):
629        print 'Extracting required data ...'
630    # Get ANUGA points
631    swwX=p2.x+xllcorner
632    swwY=p2.y+yllcorner
633
634    if(myTimeStep!='max'):
635        swwStage=p2.stage[myTimeStep,:]
636        swwHeight=p2.height[myTimeStep,:]
637        swwVel=p2.vel[myTimeStep,:]
638        timestepString=str(round(p2.time[myTimeStep]))
639    else:
640        swwStage=p2.stage.max(axis=0)
641        swwHeight=p2.height.max(axis=0)
642        swwVel=p2.vel.max(axis=0)
643        timestepString='max'
644       
645    #import pdb
646    #pdb.set_trace()
647
648    # Make name for output file
649    output_name=output_dir+'/'+os.path.splitext(os.path.basename(swwFile))[0] + '_'+\
650                output_quantity+'_'+timestepString+\
651                '_'+str(myTimeStep)+'.tif'
652
653    if(verbose):
654        print 'Computing grid of output locations...'
655    # Get points where we want raster cells
656    if(lower_left is None):
657        lower_left=[swwX.min(),swwY.min()]
658    if(upper_right is None):
659        upper_right=[swwX.max(),swwY.max()]
660
661    nx=round((upper_right[0]-lower_left[0])*1.0/(1.0*CellSize)) + 1
662    xres=(upper_right[0]-lower_left[0])*1.0/(1.0*(nx-1))
663    desiredX=scipy.arange(lower_left[0], upper_right[0],xres )
664    ny=round((upper_right[1]-lower_left[1])*1.0/(1.0*CellSize)) + 1
665    yres=(upper_right[1]-lower_left[1])*1.0/(1.0*(ny-1))
666    desiredY=scipy.arange(lower_left[1], upper_right[1], yres)
667
668    gridX, gridY=scipy.meshgrid(desiredX,desiredY)
669
670    # Make functions to interpolate quantities at points
671    if(verbose):
672        print 'Making interpolation functions...'
673    swwXY=scipy.array([swwX[:],swwY[:]]).transpose()
674    #swwXY=scipy.array([scipy.concatenate(gridX), scipy.concatenate(gridY)]).transpose()
675    if(output_quantity=='stage'):
676        qFun=scipy.interpolate.NearestNDInterpolator(swwXY,swwStage.transpose())
677    elif(output_quantity=='velocity'):
678        qFun=scipy.interpolate.NearestNDInterpolator(swwXY,swwVel.transpose())
679    elif(output_quantity=='elevation'):
680        qFun=scipy.interpolate.NearestNDInterpolator(swwXY,p2.elev.transpose())
681    elif(output_quantity=='depth'):
682        qFun=scipy.interpolate.NearestNDInterpolator(swwXY,swwHeight.transpose())
683    else:
684        raise Exception, 'output_quantity not available -- check if it is spelled correctly'
685
686    if(verbose):
687        print 'Interpolating'
688    gridq=qFun(scipy.array([scipy.concatenate(gridX),scipy.concatenate(gridY)]).transpose()).transpose()
689    gridq.shape=(len(desiredY),len(desiredX))
690    #gridq=gridq.transpose()
691
692    if(verbose):
693        print 'Making raster ...'
694    make_grid(gridq,desiredY,desiredX, output_name,EPSG_CODE)
Note: See TracBrowser for help on using the repository browser.