source: trunk/anuga_core/source/anuga/file_conversion/sww2dem.py @ 8562

Last change on this file since 8562 was 8562, checked in by steve, 13 years ago

Setting up sww2dem for Yuyang to change to faster code.
File size: 17.5 KB
Line 
1"""
2    Module to convert SWW to DEM files.
3"""
4
5# external modules
6import os
7import numpy as num
8
9# ANUGA modules
10from anuga.abstract_2d_finite_volumes.util import remove_lone_verts     
11from anuga.coordinate_transforms.geo_reference import Geo_reference
12from anuga.utilities.system_tools import get_vars_in_expression
13import anuga.utilities.log as log
14from anuga.utilities.file_utils import get_all_swwfiles
15
16
17######
18# formula mappings
19######
20
21quantity_formula = {'momentum':'(xmomentum**2 + ymomentum**2)**0.5',
22                    'depth':'stage-elevation',
23                    'speed': \
24 '(xmomentum**2 + ymomentum**2)**0.5/(stage-elevation+1.e-6/(stage-elevation))'}
25
26
27
28# Default block size for sww2dem()
29DEFAULT_BLOCK_SIZE = 10000
30
31def sww2dem(name_in, name_out,
32            quantity=None, # defaults to elevation
33            reduction=None,
34            cellsize=10,
35            number_of_decimal_places=None,
36            NODATA_value=-9999.0,
37            easting_min=None,
38            easting_max=None,
39            northing_min=None,
40            northing_max=None,
41            verbose=False,
42            origin=None,
43            datum='WGS84',
44            block_size=None):
45    """Read SWW file and convert to Digitial Elevation model format
46    (.asc or .ers)
47
48    Example (ASC):
49    ncols         3121
50    nrows         1800
51    xllcorner     722000
52    yllcorner     5893000
53    cellsize      25
54    NODATA_value  -9999
55    138.3698 137.4194 136.5062 135.5558 ..........
56
57    The number of decimal places can be specified by the user to save
58    on disk space requirements by specifying in the call to sww2dem.
59
60    Also write accompanying file with same basename_in but extension .prj
61    used to fix the UTM zone, datum, false northings and eastings.
62
63    The prj format is assumed to be as
64
65    Projection    UTM
66    Zone          56
67    Datum         WGS84
68    Zunits        NO
69    Units         METERS
70    Spheroid      WGS84
71    Xshift        0.0000000000
72    Yshift        10000000.0000000000
73    Parameters
74
75    The parameter quantity must be the name of an existing quantity or
76    an expression involving existing quantities. The default is
77    'elevation'. Quantity is not a list of quantities.
78
79    If reduction is given and it's an index, sww2dem will output the quantity at that time-step.
80    If reduction is given and it's a built in function (eg max, min, mean), then that
81    function is used to reduce the quantity over all time-steps. If reduction is not given,
82    reduction is set to "max" by default.
83
84    datum
85
86    format can be either 'asc' or 'ers'
87    block_size - sets the number of slices along the non-time axis to
88                 process in one block.
89    """
90
91    import sys
92    import types
93
94    from anuga.geometry.polygon import inside_polygon, outside_polygon
95    from anuga.abstract_2d_finite_volumes.util import \
96         apply_expression_to_dictionary
97
98    basename_in, in_ext = os.path.splitext(name_in)
99    basename_out, out_ext = os.path.splitext(name_out)
100    out_ext = out_ext.lower()
101
102    if in_ext != '.sww':
103        raise IOError('Input format for %s must be .sww' % name_in)
104
105    if out_ext not in ['.asc', '.ers']:
106        raise IOError('Format for %s must be either asc or ers.' % name_out)
107
108    false_easting = 500000
109    false_northing = 10000000
110
111    if quantity is None:
112        quantity = 'elevation'
113   
114    if reduction is None:
115        reduction = max
116
117    if quantity_formula.has_key(quantity):
118        quantity = quantity_formula[quantity]
119
120    if number_of_decimal_places is None:
121        number_of_decimal_places = 3
122
123    if block_size is None:
124        block_size = DEFAULT_BLOCK_SIZE
125
126    assert(isinstance(block_size, (int, long, float)))
127
128    # Read sww file
129    if verbose:
130        log.critical('Reading from %s' % name_in)
131        log.critical('Output directory is %s' % name_out)
132
133    from Scientific.IO.NetCDF import NetCDFFile
134    fid = NetCDFFile(name_in)
135
136    #Get extent and reference
137    x = num.array(fid.variables['x'], num.float)
138    y = num.array(fid.variables['y'], num.float)
139    volumes = num.array(fid.variables['volumes'], num.int)
140    if type(reduction) is not types.BuiltinFunctionType:
141        times = fid.variables['time'][reduction]
142    else:
143        times = fid.variables['time'][:]
144
145    number_of_timesteps = fid.dimensions['number_of_timesteps']
146    number_of_points = fid.dimensions['number_of_points']
147
148    if origin is None:
149        # Get geo_reference
150        # sww files don't have to have a geo_ref
151        try:
152            geo_reference = Geo_reference(NetCDFObject=fid)
153        except AttributeError, e:
154            geo_reference = Geo_reference() # Default georef object
155
156        xllcorner = geo_reference.get_xllcorner()
157        yllcorner = geo_reference.get_yllcorner()
158        zone = geo_reference.get_zone()
159    else:
160        zone = origin[0]
161        xllcorner = origin[1]
162        yllcorner = origin[2]
163
164    # FIXME: Refactor using code from Interpolation_function.statistics
165    # (in interpolate.py)
166    # Something like print swwstats(swwname)
167    if verbose:
168        log.critical('------------------------------------------------')
169        log.critical('Statistics of SWW file:')
170        log.critical('  Name: %s' % name_in)
171        log.critical('  Reference:')
172        log.critical('    Lower left corner: [%f, %f]' % (xllcorner, yllcorner))
173        if type(reduction) is not types.BuiltinFunctionType:
174            log.critical('    Time: %f' % times)
175        else:
176            log.critical('    Start time: %f' % fid.starttime[0])
177        log.critical('  Extent:')
178        log.critical('    x [m] in [%f, %f], len(x) == %d'
179                     %(num.min(x), num.max(x), len(x.flat)))
180        log.critical('    y [m] in [%f, %f], len(y) == %d'
181                     % (num.min(y), num.max(y), len(y.flat)))
182        if type(reduction) is not types.BuiltinFunctionType:
183            log.critical('    t [s] = %f, len(t) == %d' % (times, 1))
184        else:
185            log.critical('    t [s] in [%f, %f], len(t) == %d'
186                         % (min(times), max(times), len(times)))
187        log.critical('  Quantities [SI units]:')
188       
189        # Comment out for reduced memory consumption
190        for name in ['stage', 'xmomentum', 'ymomentum']:
191            q = fid.variables[name][:].flatten()
192            if type(reduction) is not types.BuiltinFunctionType:
193                q = q[reduction*len(x):(reduction+1)*len(x)]
194            if verbose: log.critical('    %s in [%f, %f]'
195                                     % (name, min(q), max(q)))
196        for name in ['elevation']:
197            q = fid.variables[name][:].flatten()
198            if verbose: log.critical('    %s in [%f, %f]'
199                                     % (name, min(q), max(q)))
200
201    # Get the variables in the supplied expression.
202    # This may throw a SyntaxError exception.
203    var_list = get_vars_in_expression(quantity)
204
205    # Check that we have the required variables in the SWW file.
206    missing_vars = []
207    for name in var_list:
208        try:
209            _ = fid.variables[name]
210        except KeyError:
211            missing_vars.append(name)
212    if missing_vars:
213        msg = ("In expression '%s', variables %s are not in the SWW file '%s'"
214               % (quantity, str(missing_vars), name_in))
215        raise Exception, msg
216
217    # Create result array and start filling, block by block.
218    result = num.zeros(number_of_points, num.float)
219
220    if verbose:
221        msg = 'Slicing sww file, num points: ' + str(number_of_points)
222        msg += ', block size: ' + str(block_size)
223        log.critical(msg)
224
225    for start_slice in xrange(0, number_of_points, block_size):
226        # Limit slice size to array end if at last block
227        end_slice = min(start_slice + block_size, number_of_points)
228       
229        # Get slices of all required variables
230        q_dict = {}
231        for name in var_list:
232            # check if variable has time axis
233            if len(fid.variables[name].shape) == 2:
234                q_dict[name] = fid.variables[name][:,start_slice:end_slice]
235            else:       # no time axis
236                q_dict[name] = fid.variables[name][start_slice:end_slice]
237
238        # Evaluate expression with quantities found in SWW file
239        res = apply_expression_to_dictionary(quantity, q_dict)
240
241        if len(res.shape) == 2:
242            new_res = num.zeros(res.shape[1], num.float)
243            for k in xrange(res.shape[1]):
244                if type(reduction) is not types.BuiltinFunctionType:
245                    new_res[k] = res[reduction,k]
246                else:
247                    new_res[k] = reduction(res[:,k])
248            res = new_res
249
250        result[start_slice:end_slice] = res
251                                   
252    # Post condition: Now q has dimension: number_of_points
253    assert len(result.shape) == 1
254    assert result.shape[0] == number_of_points
255
256    if verbose:
257        log.critical('Processed values for %s are in [%f, %f]'
258                     % (quantity, min(result), max(result)))
259
260    # Create grid and update xll/yll corner and x,y
261    # Relative extent
262    if easting_min is None:
263        xmin = min(x)
264    else:
265        xmin = easting_min - xllcorner
266
267    if easting_max is None:
268        xmax = max(x)
269    else:
270        xmax = easting_max - xllcorner
271
272    if northing_min is None:
273        ymin = min(y)
274    else:
275        ymin = northing_min - yllcorner
276
277    if northing_max is None:
278        ymax = max(y)
279    else:
280        ymax = northing_max - yllcorner
281
282    msg = 'xmax must be greater than or equal to xmin.\n'
283    msg += 'I got xmin = %f, xmax = %f' %(xmin, xmax)
284    assert xmax >= xmin, msg
285
286    msg = 'ymax must be greater than or equal to xmin.\n'
287    msg += 'I got ymin = %f, ymax = %f' %(ymin, ymax)
288    assert ymax >= ymin, msg
289
290    if verbose: log.critical('Creating grid')
291    ncols = int((xmax-xmin)/cellsize) + 1
292    nrows = int((ymax-ymin)/cellsize) + 1
293
294    # New absolute reference and coordinates
295    newxllcorner = xmin + xllcorner
296    newyllcorner = ymin + yllcorner
297
298    x = x + xllcorner - newxllcorner
299    y = y + yllcorner - newyllcorner
300
301    vertex_points = num.concatenate ((x[:,num.newaxis], y[:,num.newaxis]), axis=1)
302    assert len(vertex_points.shape) == 2
303
304
305    def calc_grid_values(vertex_points, volumes, result):
306
307        grid_points = num.zeros ((ncols*nrows, 2), num.float)
308
309        for i in xrange(nrows):
310            if out_ext == '.asc':
311                yg = i * cellsize
312            else:
313                # this will flip the order of the y values for ers
314                yg = (nrows-i) * cellsize
315
316            for j in xrange(ncols):
317                xg = j * cellsize
318                k = i*ncols + j
319
320                grid_points[k, 0] = xg
321                grid_points[k, 1] = yg
322
323        # Interpolate
324        from anuga.fit_interpolate.interpolate import Interpolate
325
326        # Remove loners from vertex_points, volumes here
327        vertex_points, volumes = remove_lone_verts(vertex_points, volumes)
328        # export_mesh_file('monkey.tsh',{'vertices':vertex_points, 'triangles':volumes})
329
330
331
332        #
333
334        interp = Interpolate(vertex_points, volumes, verbose = verbose)
335
336        # Interpolate using quantity values
337        if verbose: log.critical('Interpolating')
338        grid_values = interp.interpolate(result, grid_points).flatten()
339        outside_indices = interp.get_outside_poly_indices()
340
341        for i in outside_indices:
342            #print 'change grid_value',NODATA_value
343            grid_values[i] = NODATA_value
344
345        return grid_values
346
347
348
349    grid_values = calc_grid_values(vertex_points, volumes, result)
350
351
352
353    #print outside_indices
354
355    if verbose:
356        log.critical('Interpolated values are in [%f, %f]'
357                     % (num.min(grid_values), num.max(grid_values)))
358
359    # Assign NODATA_value to all points outside bounding polygon (from interpolation mesh)
360   
361#    P = interp.mesh.get_boundary_polygon()
362#    outside_indices = outside_polygon(grid_points, P, closed=True)
363
364
365
366    if out_ext == '.ers':
367        # setup ERS header information
368        grid_values = num.reshape(grid_values, (nrows, ncols))
369        header = {}
370        header['datum'] = '"' + datum + '"'
371        # FIXME The use of hardwired UTM and zone number needs to be made optional
372        # FIXME Also need an automatic test for coordinate type (i.e. EN or LL)
373        header['projection'] = '"UTM-' + str(zone) + '"'
374        header['coordinatetype'] = 'EN'
375        if header['coordinatetype'] == 'LL':
376            header['longitude'] = str(newxllcorner)
377            header['latitude'] = str(newyllcorner)
378        elif header['coordinatetype'] == 'EN':
379            header['eastings'] = str(newxllcorner)
380            header['northings'] = str(newyllcorner)
381        header['nullcellvalue'] = str(NODATA_value)
382        header['xdimension'] = str(cellsize)
383        header['ydimension'] = str(cellsize)
384        header['value'] = '"' + quantity + '"'
385        #header['celltype'] = 'IEEE8ByteReal'  #FIXME: Breaks unit test
386
387        #Write
388        if verbose:
389            log.critical('Writing %s' % name_out)
390
391        import ermapper_grids
392
393        ermapper_grids.write_ermapper_grid(name_out, grid_values, header)
394
395        fid.close()
396    else:
397        #Write to Ascii format
398        #Write prj file
399        prjfile = basename_out + '.prj'
400
401        if verbose: log.critical('Writing %s' % prjfile)
402        prjid = open(prjfile, 'w')
403        prjid.write('Projection    %s\n' %'UTM')
404        prjid.write('Zone          %d\n' %zone)
405        prjid.write('Datum         %s\n' %datum)
406        prjid.write('Zunits        NO\n')
407        prjid.write('Units         METERS\n')
408        prjid.write('Spheroid      %s\n' %datum)
409        prjid.write('Xshift        %d\n' %false_easting)
410        prjid.write('Yshift        %d\n' %false_northing)
411        prjid.write('Parameters\n')
412        prjid.close()
413
414        if verbose: log.critical('Writing %s' % name_out)
415
416        ascid = open(name_out, 'w')
417
418        ascid.write('ncols         %d\n' %ncols)
419        ascid.write('nrows         %d\n' %nrows)
420        ascid.write('xllcorner     %d\n' %newxllcorner)
421        ascid.write('yllcorner     %d\n' %newyllcorner)
422        ascid.write('cellsize      %f\n' %cellsize)
423        ascid.write('NODATA_value  %d\n' %NODATA_value)
424
425        #Get bounding polygon from mesh
426        #P = interp.mesh.get_boundary_polygon()
427        #inside_indices = inside_polygon(grid_points, P)
428
429        # change printoptions so that a long string of zeros in not
430        # summarized as [0.0, 0.0, 0.0, ... 0.0, 0.0, 0.0]
431        #printoptions = num.get_printoptions()
432        #num.set_printoptions(threshold=sys.maxint)
433
434        format = '%.'+'%g' % number_of_decimal_places +'e'
435        for i in range(nrows):
436            if verbose and i % ((nrows+10)/10) == 0:
437                log.critical('Doing row %d of %d' % (i, nrows))
438
439            base_index = (nrows-i-1)*ncols
440
441            slice = grid_values[base_index:base_index+ncols]
442
443            num.savetxt(ascid, slice.reshape(1,ncols), format, ' ' )
444           
445       
446        #Close
447        ascid.close()
448        fid.close()
449
450        return basename_out
451
452
453
454def sww2dem_batch(basename_in, extra_name_out=None,
455                quantities=None, # defaults to elevation
456                reduction=None,
457                cellsize=10,
458                number_of_decimal_places=None,
459                NODATA_value=-9999,
460                easting_min=None,
461                easting_max=None,
462                northing_min=None,
463                northing_max=None,
464                verbose=False,
465                origin=None,
466                datum='WGS84',
467                format='ers'):
468    """Wrapper for sww2dem.
469    See sww2dem to find out what most of the parameters do. Note that since this
470    is a batch command, the normal filename naming conventions do not apply.
471
472    basename_in is a path to sww file/s, without the .sww extension.
473    extra_name_out is a postfix to add to the output filename.
474
475    Quantities is a list of quantities.  Each quantity will be
476    calculated for each sww file.
477
478    This returns the basenames of the files returned, which is made up
479    of the dir and all of the file name, except the extension.
480
481    This function returns the names of the files produced.
482
483    It will also produce as many output files as there are input sww files.
484    """
485
486    if quantities is None:
487        quantities = ['elevation']
488
489    if type(quantities) is str:
490            quantities = [quantities]
491
492    # How many sww files are there?
493    dir, base = os.path.split(basename_in)
494
495    iterate_over = get_all_swwfiles(dir, base, verbose)
496
497    if dir == "":
498        dir = "." # Unix compatibility
499
500    files_out = []
501    for sww_file in iterate_over:
502        for quantity in quantities:
503            if extra_name_out is None:
504                basename_out = sww_file + '_' + quantity
505            else:
506                basename_out = sww_file + '_' + quantity + '_' + extra_name_out
507
508            swwin = dir+os.sep+sww_file+'.sww'
509            demout = dir+os.sep+basename_out+'.'+format
510
511            if verbose:
512                log.critical('sww2dem: %s => %s' % (swwin, demout))
513
514            file_out = sww2dem(swwin,
515                               demout,
516                               quantity,
517                               reduction,
518                               cellsize,
519                               number_of_decimal_places,
520                               NODATA_value,
521                               easting_min,
522                               easting_max,
523                               northing_min,
524                               northing_max,
525                               verbose,
526                               origin,
527                               datum)
528                               
529            files_out.append(file_out)
530    return files_out
531
Note: See TracBrowser for help on using the repository browser.