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

Last change on this file since 8137 was 8137, checked in by wilsonr, 14 years ago

Added a 'KeyError?' type to a bare exception catch.
File size: 17.1 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,
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    grid_points = num.zeros ((ncols*nrows, 2), num.float)
305
306    for i in xrange(nrows):
307        if out_ext == '.asc':
308            yg = i * cellsize
309        else:
310            # this will flip the order of the y values for ers
311            yg = (nrows-i) * cellsize
312
313        for j in xrange(ncols):
314            xg = j * cellsize
315            k = i*ncols + j
316
317            grid_points[k, 0] = xg
318            grid_points[k, 1] = yg
319
320    # Interpolate
321    from anuga.fit_interpolate.interpolate import Interpolate
322
323    # Remove loners from vertex_points, volumes here
324    vertex_points, volumes = remove_lone_verts(vertex_points, volumes)
325    # export_mesh_file('monkey.tsh',{'vertices':vertex_points, 'triangles':volumes})
326
327    interp = Interpolate(vertex_points, volumes, verbose = verbose)
328
329    # Interpolate using quantity values
330    if verbose: log.critical('Interpolating')
331    grid_values = interp.interpolate(result, grid_points).flatten()
332
333    if verbose:
334        log.critical('Interpolated values are in [%f, %f]'
335                     % (num.min(grid_values), num.max(grid_values)))
336
337    # Assign NODATA_value to all points outside bounding polygon (from interpolation mesh)
338    P = interp.mesh.get_boundary_polygon()
339    outside_indices = outside_polygon(grid_points, P, closed=True)
340
341    for i in outside_indices:
342        grid_values[i] = NODATA_value
343
344    if out_ext == '.ers':
345        # setup ERS header information
346        grid_values = num.reshape(grid_values, (nrows, ncols))
347        header = {}
348        header['datum'] = '"' + datum + '"'
349        # FIXME The use of hardwired UTM and zone number needs to be made optional
350        # FIXME Also need an automatic test for coordinate type (i.e. EN or LL)
351        header['projection'] = '"UTM-' + str(zone) + '"'
352        header['coordinatetype'] = 'EN'
353        if header['coordinatetype'] == 'LL':
354            header['longitude'] = str(newxllcorner)
355            header['latitude'] = str(newyllcorner)
356        elif header['coordinatetype'] == 'EN':
357            header['eastings'] = str(newxllcorner)
358            header['northings'] = str(newyllcorner)
359        header['nullcellvalue'] = str(NODATA_value)
360        header['xdimension'] = str(cellsize)
361        header['ydimension'] = str(cellsize)
362        header['value'] = '"' + quantity + '"'
363        #header['celltype'] = 'IEEE8ByteReal'  #FIXME: Breaks unit test
364
365        #Write
366        if verbose:
367            log.critical('Writing %s' % name_out)
368
369        import ermapper_grids
370
371        ermapper_grids.write_ermapper_grid(name_out, grid_values, header)
372
373        fid.close()
374    else:
375        #Write to Ascii format
376        #Write prj file
377        prjfile = basename_out + '.prj'
378
379        if verbose: log.critical('Writing %s' % prjfile)
380        prjid = open(prjfile, 'w')
381        prjid.write('Projection    %s\n' %'UTM')
382        prjid.write('Zone          %d\n' %zone)
383        prjid.write('Datum         %s\n' %datum)
384        prjid.write('Zunits        NO\n')
385        prjid.write('Units         METERS\n')
386        prjid.write('Spheroid      %s\n' %datum)
387        prjid.write('Xshift        %d\n' %false_easting)
388        prjid.write('Yshift        %d\n' %false_northing)
389        prjid.write('Parameters\n')
390        prjid.close()
391
392        if verbose: log.critical('Writing %s' % name_out)
393
394        ascid = open(name_out, 'w')
395
396        ascid.write('ncols         %d\n' %ncols)
397        ascid.write('nrows         %d\n' %nrows)
398        ascid.write('xllcorner     %d\n' %newxllcorner)
399        ascid.write('yllcorner     %d\n' %newyllcorner)
400        ascid.write('cellsize      %f\n' %cellsize)
401        ascid.write('NODATA_value  %d\n' %NODATA_value)
402
403        #Get bounding polygon from mesh
404        #P = interp.mesh.get_boundary_polygon()
405        #inside_indices = inside_polygon(grid_points, P)
406
407        # change printoptions so that a long string of zeros in not
408        # summarized as [0.0, 0.0, 0.0, ... 0.0, 0.0, 0.0]
409        #printoptions = num.get_printoptions()
410        #num.set_printoptions(threshold=sys.maxint)
411
412        format = '%.'+'%g' % number_of_decimal_places +'e'
413        for i in range(nrows):
414            if verbose and i % ((nrows+10)/10) == 0:
415                log.critical('Doing row %d of %d' % (i, nrows))
416
417            base_index = (nrows-i-1)*ncols
418
419            slice = grid_values[base_index:base_index+ncols]
420
421            num.savetxt(ascid, slice.reshape(1,ncols), format, ' ' )
422           
423       
424        #Close
425        ascid.close()
426        fid.close()
427
428        return basename_out
429
430
431
432def sww2dem_batch(basename_in, extra_name_out=None,
433                quantities=None, # defaults to elevation
434                reduction=None,
435                cellsize=10,
436                number_of_decimal_places=None,
437                NODATA_value=-9999,
438                easting_min=None,
439                easting_max=None,
440                northing_min=None,
441                northing_max=None,
442                verbose=False,
443                origin=None,
444                datum='WGS84',
445                format='ers'):
446    """Wrapper for sww2dem.
447    See sww2dem to find out what most of the parameters do. Note that since this
448    is a batch command, the normal filename naming conventions do not apply.
449
450    basename_in is a path to sww file/s, without the .sww extension.
451    extra_name_out is a postfix to add to the output filename.
452
453    Quantities is a list of quantities.  Each quantity will be
454    calculated for each sww file.
455
456    This returns the basenames of the files returned, which is made up
457    of the dir and all of the file name, except the extension.
458
459    This function returns the names of the files produced.
460
461    It will also produce as many output files as there are input sww files.
462    """
463
464    if quantities is None:
465        quantities = ['elevation']
466
467    if type(quantities) is str:
468            quantities = [quantities]
469
470    # How many sww files are there?
471    dir, base = os.path.split(basename_in)
472
473    iterate_over = get_all_swwfiles(dir, base, verbose)
474
475    if dir == "":
476        dir = "." # Unix compatibility
477
478    files_out = []
479    for sww_file in iterate_over:
480        for quantity in quantities:
481            if extra_name_out is None:
482                basename_out = sww_file + '_' + quantity
483            else:
484                basename_out = sww_file + '_' + quantity + '_' + extra_name_out
485
486            swwin = dir+os.sep+sww_file+'.sww'
487            demout = dir+os.sep+basename_out+'.'+format
488
489            if verbose:
490                log.critical('sww2dem: %s => %s' % (swwin, demout))
491
492            file_out = sww2dem(swwin,
493                               demout,
494                               quantity,
495                               reduction,
496                               cellsize,
497                               number_of_decimal_places,
498                               NODATA_value,
499                               easting_min,
500                               easting_max,
501                               northing_min,
502                               northing_max,
503                               verbose,
504                               origin,
505                               datum)
506                               
507            files_out.append(file_out)
508    return files_out
509
Note: See TracBrowser for help on using the repository browser.