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

Last change on this file since 7841 was 7841, checked in by hudson, 14 years ago

Refactorings to allow tests to pass.
File size: 16.9 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
76    The parameter quantity must be the name of an existing quantity or
77    an expression involving existing quantities. The default is
78    'elevation'. Quantity is not a list of quantities.
79
80    if timestep (an index) is given, output quantity at that timestep
81
82    if reduction is given and its an index, output quantity at that timestep. If reduction is given
83    and is a built in function, use that to reduce quantity over all timesteps.
84
85    datum
86
87    format can be either 'asc' or 'ers'
88    block_size - sets the number of slices along the non-time axis to
89                 process in one block.
90    """
91
92    import sys
93    import types
94
95    from anuga.geometry.polygon import inside_polygon, outside_polygon
96    from anuga.abstract_2d_finite_volumes.util import \
97         apply_expression_to_dictionary
98
99    basename_in, in_ext = os.path.splitext(name_in)
100    basename_out, out_ext = os.path.splitext(name_out)
101    out_ext = out_ext.lower()
102
103    if in_ext != '.sww':
104        raise IOError('Input format for %s must be .sww' % name_in)
105
106    if out_ext not in ['.asc', '.ers']:
107        raise IOError('Format for %s must be either asc or ers.' % name_out)
108
109    false_easting = 500000
110    false_northing = 10000000
111
112    if quantity is None:
113        quantity = 'elevation'
114   
115    if reduction is None:
116        reduction = max
117
118    if quantity_formula.has_key(quantity):
119        quantity = quantity_formula[quantity]
120
121    if number_of_decimal_places is None:
122        number_of_decimal_places = 3
123
124    if block_size is None:
125        block_size = DEFAULT_BLOCK_SIZE
126
127    assert(isinstance(block_size, (int, long, float)))
128
129    # Read sww file
130    if verbose:
131        log.critical('Reading from %s' % name_in)
132        log.critical('Output directory is %s' % name_out)
133
134    from Scientific.IO.NetCDF import NetCDFFile
135    fid = NetCDFFile(name_in)
136
137    #Get extent and reference
138    x = fid.variables['x'][:]
139    y = fid.variables['y'][:]
140    volumes = fid.variables['volumes'][:]
141    if type(reduction) is not types.BuiltinFunctionType:
142        times = fid.variables['time'][reduction]
143    else:
144        times = fid.variables['time'][:]
145
146    number_of_timesteps = fid.dimensions['number_of_timesteps']
147    number_of_points = fid.dimensions['number_of_points']
148
149    if origin is None:
150        # Get geo_reference
151        # sww files don't have to have a geo_ref
152        try:
153            geo_reference = Geo_reference(NetCDFObject=fid)
154        except AttributeError, e:
155            geo_reference = Geo_reference() # Default georef object
156
157        xllcorner = geo_reference.get_xllcorner()
158        yllcorner = geo_reference.get_yllcorner()
159        zone = geo_reference.get_zone()
160    else:
161        zone = origin[0]
162        xllcorner = origin[1]
163        yllcorner = origin[2]
164
165    # FIXME: Refactor using code from Interpolation_function.statistics
166    # (in interpolate.py)
167    # Something like print swwstats(swwname)
168    if verbose:
169        log.critical('------------------------------------------------')
170        log.critical('Statistics of SWW file:')
171        log.critical('  Name: %s' % name_in)
172        log.critical('  Reference:')
173        log.critical('    Lower left corner: [%f, %f]' % (xllcorner, yllcorner))
174        if type(reduction) is not types.BuiltinFunctionType:
175            log.critical('    Time: %f' % times)
176        else:
177            log.critical('    Start time: %f' % fid.starttime[0])
178        log.critical('  Extent:')
179        log.critical('    x [m] in [%f, %f], len(x) == %d'
180                     %(num.min(x), num.max(x), len(x.flat)))
181        log.critical('    y [m] in [%f, %f], len(y) == %d'
182                     % (num.min(y), num.max(y), len(y.flat)))
183        if type(reduction) is not types.BuiltinFunctionType:
184            log.critical('    t [s] = %f, len(t) == %d' % (times, 1))
185        else:
186            log.critical('    t [s] in [%f, %f], len(t) == %d'
187                         % (min(times), max(times), len(times)))
188        log.critical('  Quantities [SI units]:')
189       
190        # Comment out for reduced memory consumption
191        for name in ['stage', 'xmomentum', 'ymomentum']:
192            q = fid.variables[name][:].flatten()
193            if type(reduction) is not types.BuiltinFunctionType:
194                q = q[reduction*len(x):(reduction+1)*len(x)]
195            if verbose: log.critical('    %s in [%f, %f]'
196                                     % (name, min(q), max(q)))
197        for name in ['elevation']:
198            q = fid.variables[name][:].flatten()
199            if verbose: log.critical('    %s in [%f, %f]'
200                                     % (name, min(q), max(q)))
201
202    # Get the variables in the supplied expression.
203    # This may throw a SyntaxError exception.
204    var_list = get_vars_in_expression(quantity)
205
206    # Check that we have the required variables in the SWW file.
207    missing_vars = []
208    for name in var_list:
209        try:
210            _ = fid.variables[name]
211        except:
212            missing_vars.append(name)
213    if missing_vars:
214        msg = ("In expression '%s', variables %s are not in the SWW file '%s'"
215               % (quantity, name_in))
216        raise Exception, msg
217
218    # Create result array and start filling, block by block.
219    result = num.zeros(number_of_points, num.float)
220
221    if verbose:
222        msg = 'Slicing sww file, num points: ' + str(number_of_points)
223        msg += ', block size: ' + str(block_size)
224        log.critical(msg)
225
226    for start_slice in xrange(0, number_of_points, block_size):
227        # Limit slice size to array end if at last block
228        end_slice = min(start_slice + block_size, number_of_points)
229       
230        # Get slices of all required variables
231        q_dict = {}
232        for name in var_list:
233            # check if variable has time axis
234            if len(fid.variables[name].shape) == 2:
235                q_dict[name] = fid.variables[name][:,start_slice:end_slice]
236            else:       # no time axis
237                q_dict[name] = fid.variables[name][start_slice:end_slice]
238
239        # Evaluate expression with quantities found in SWW file
240        res = apply_expression_to_dictionary(quantity, q_dict)
241
242        if len(res.shape) == 2:
243            new_res = num.zeros(res.shape[1], num.float)
244            for k in xrange(res.shape[1]):
245                if type(reduction) is not types.BuiltinFunctionType:
246                    new_res[k] = res[reduction,k]
247                else:
248                    new_res[k] = reduction(res[:,k])
249            res = new_res
250
251        result[start_slice:end_slice] = res
252                                   
253    # Post condition: Now q has dimension: number_of_points
254    assert len(result.shape) == 1
255    assert result.shape[0] == number_of_points
256
257    if verbose:
258        log.critical('Processed values for %s are in [%f, %f]'
259                     % (quantity, min(result), max(result)))
260
261    # Create grid and update xll/yll corner and x,y
262    # Relative extent
263    if easting_min is None:
264        xmin = min(x)
265    else:
266        xmin = easting_min - xllcorner
267
268    if easting_max is None:
269        xmax = max(x)
270    else:
271        xmax = easting_max - xllcorner
272
273    if northing_min is None:
274        ymin = min(y)
275    else:
276        ymin = northing_min - yllcorner
277
278    if northing_max is None:
279        ymax = max(y)
280    else:
281        ymax = northing_max - yllcorner
282
283    msg = 'xmax must be greater than or equal to xmin.\n'
284    msg += 'I got xmin = %f, xmax = %f' %(xmin, xmax)
285    assert xmax >= xmin, msg
286
287    msg = 'ymax must be greater than or equal to xmin.\n'
288    msg += 'I got ymin = %f, ymax = %f' %(ymin, ymax)
289    assert ymax >= ymin, msg
290
291    if verbose: log.critical('Creating grid')
292    ncols = int((xmax-xmin)/cellsize) + 1
293    nrows = int((ymax-ymin)/cellsize) + 1
294
295    # New absolute reference and coordinates
296    newxllcorner = xmin + xllcorner
297    newyllcorner = ymin + yllcorner
298
299    x = x + xllcorner - newxllcorner
300    y = y + yllcorner - newyllcorner
301
302    vertex_points = num.concatenate ((x[:,num.newaxis], y[:,num.newaxis]), axis=1)
303    assert len(vertex_points.shape) == 2
304
305    grid_points = num.zeros ((ncols*nrows, 2), num.float)
306
307    for i in xrange(nrows):
308        if out_ext == '.asc':
309            yg = i * cellsize
310        else:
311            # this will flip the order of the y values for ers
312            yg = (nrows-i) * cellsize
313
314        for j in xrange(ncols):
315            xg = j * cellsize
316            k = i*ncols + j
317
318            grid_points[k, 0] = xg
319            grid_points[k, 1] = yg
320
321    # Interpolate
322    from anuga.fit_interpolate.interpolate import Interpolate
323
324    # Remove loners from vertex_points, volumes here
325    vertex_points, volumes = remove_lone_verts(vertex_points, volumes)
326    # export_mesh_file('monkey.tsh',{'vertices':vertex_points, 'triangles':volumes})
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.