source: trunk/anuga_core/source/anuga/file_conversion/sww2dem_new.py @ 8626

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

Fast version of sww2dem (now in the correct directory)
File size: 18.8 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
15from anuga.fit_interpolate.interpolate_ext import *
16
17
18######
19# formula mappings
20######
21
22quantity_formula = {'momentum':'(xmomentum**2 + ymomentum**2)**0.5',
23                    'depth':'stage-elevation',
24                    'speed': \
25 '(xmomentum**2 + ymomentum**2)**0.5/(stage-elevation+1.e-6/(stage-elevation))'}
26
27
28
29# Default block size for sww2dem()
30DEFAULT_BLOCK_SIZE = 10000
31
32def sww2dem(name_in, name_out,
33            quantity=None, # defaults to elevation
34            reduction=None,
35            cellsize=10,
36            number_of_decimal_places=None,
37            NODATA_value=-9999.0,
38            easting_min=None,
39            easting_max=None,
40            northing_min=None,
41            northing_max=None,
42            verbose=False,
43            origin=None,
44            datum='WGS84',
45            block_size=None):
46    """Read SWW file and convert to Digitial Elevation model format
47    (.asc or .ers)
48
49    Example (ASC):
50    ncols         3121
51    nrows         1800
52    xllcorner     722000
53    yllcorner     5893000
54    cellsize      25
55    NODATA_value  -9999
56    138.3698 137.4194 136.5062 135.5558 ..........
57
58    The number of decimal places can be specified by the user to save
59    on disk space requirements by specifying in the call to sww2dem.
60
61    Also write accompanying file with same basename_in but extension .prj
62    used to fix the UTM zone, datum, false northings and eastings.
63
64    The prj format is assumed to be as
65
66    Projection    UTM
67    Zone          56
68    Datum         WGS84
69    Zunits        NO
70    Units         METERS
71    Spheroid      WGS84
72    Xshift        0.0000000000
73    Yshift        10000000.0000000000
74    Parameters
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 reduction is given and it's an index, sww2dem will output the quantity at that time-step.
81    If reduction is given and it's a built in function (eg max, min, mean), then that
82    function is used to reduce the quantity over all time-steps. If reduction is not given,
83    reduction is set to "max" by default.
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 = num.array(fid.variables['x'], num.float)
139    y = num.array(fid.variables['y'], num.float)
140    volumes = num.array(fid.variables['volumes'], num.int)
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 KeyError:
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, str(missing_vars), 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
306    def calc_grid_values_old(vertex_points, volumes, result):
307
308        grid_points = num.zeros ((ncols*nrows, 2), num.float)
309
310        for i in xrange(nrows):
311            if out_ext == '.asc':
312                yg = i * cellsize
313            else:
314                # this will flip the order of the y values for ers
315                yg = (nrows-i) * cellsize
316
317            for j in xrange(ncols):
318                xg = j * cellsize
319                k = i*ncols + j
320
321                grid_points[k, 0] = xg
322                grid_points[k, 1] = yg
323
324        # Interpolate
325        from anuga.fit_interpolate.interpolate import Interpolate
326
327        # Remove loners from vertex_points, volumes here
328        vertex_points, volumes = remove_lone_verts(vertex_points, volumes)
329        # export_mesh_file('monkey.tsh',{'vertices':vertex_points, 'triangles':volumes})
330
331
332        interp = Interpolate(vertex_points, volumes, verbose = verbose)
333
334        bprint = 0
335
336        # Interpolate using quantity values
337        if verbose: log.critical('Interpolating')
338        grid_values = interp.interpolate(bprint, 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    def calc_grid_values(vertex_points, volumes, result):
348
349        grid_points = num.zeros ((ncols*nrows, 2), num.float)
350
351        for i in xrange(nrows):
352            if out_ext == '.asc':
353                yg = i * cellsize
354            else:
355                #this will flip the order of the y values for ers
356                yg = (nrows-i) * cellsize
357   
358            for j in xrange(ncols):
359                xg = j * cellsize
360                k = i*ncols + j
361
362                grid_points[k, 0] = xg
363                grid_points[k, 1] = yg
364   
365        grid_values = num.zeros(ncols*nrows, num.float)
366
367        eval_grid(nrows, ncols, NODATA_value, grid_points, vertex_points.flatten(), volumes, result, grid_values);
368        return grid_values.flatten()
369
370    grid_values = calc_grid_values(vertex_points, volumes, result)
371
372    if verbose:
373        log.critical('Interpolated values are in [%f, %f]'
374                     % (num.min(grid_values), num.max(grid_values)))
375
376    # Assign NODATA_value to all points outside bounding polygon (from interpolation mesh)
377   
378#    P = interp.mesh.get_boundary_polygon()
379#    outside_indices = outside_polygon(grid_points, P, closed=True)
380
381    if out_ext == '.ers':
382        # setup ERS header information
383        grid_values = num.reshape(grid_values, (nrows, ncols))
384        header = {}
385        header['datum'] = '"' + datum + '"'
386        # FIXME The use of hardwired UTM and zone number needs to be made optional
387        # FIXME Also need an automatic test for coordinate type (i.e. EN or LL)
388        header['projection'] = '"UTM-' + str(zone) + '"'
389        header['coordinatetype'] = 'EN'
390        if header['coordinatetype'] == 'LL':
391            header['longitude'] = str(newxllcorner)
392            header['latitude'] = str(newyllcorner)
393        elif header['coordinatetype'] == 'EN':
394            header['eastings'] = str(newxllcorner)
395            header['northings'] = str(newyllcorner)
396        header['nullcellvalue'] = str(NODATA_value)
397        header['xdimension'] = str(cellsize)
398        header['ydimension'] = str(cellsize)
399        header['value'] = '"' + quantity + '"'
400        #header['celltype'] = 'IEEE8ByteReal'  #FIXME: Breaks unit test
401
402        #Write
403        if verbose:
404            log.critical('Writing %s' % name_out)
405
406        import ermapper_grids
407
408        ermapper_grids.write_ermapper_grid(name_out, grid_values, header)
409
410        fid.close()
411   
412    else:
413        #Write to Ascii format
414        #Write prj file
415        prjfile = basename_out + '.prj'
416
417        if verbose: log.critical('Writing %s' % prjfile)
418        prjid = open(prjfile, 'w')
419        prjid.write('Projection    %s\n' %'UTM')
420        prjid.write('Zone          %d\n' %zone)
421        prjid.write('Datum         %s\n' %datum)
422        prjid.write('Zunits        NO\n')
423        prjid.write('Units         METERS\n')
424        prjid.write('Spheroid      %s\n' %datum)
425        prjid.write('Xshift        %d\n' %false_easting)
426        prjid.write('Yshift        %d\n' %false_northing)
427        prjid.write('Parameters\n')
428        prjid.close()
429
430        if verbose: log.critical('Writing %s' % name_out)
431
432        ascid = open(name_out, 'w')
433
434        ascid.write('ncols         %d\n' %ncols)
435        ascid.write('nrows         %d\n' %nrows)
436        ascid.write('xllcorner     %d\n' %newxllcorner)
437        ascid.write('yllcorner     %d\n' %newyllcorner)
438        ascid.write('cellsize      %f\n' %cellsize)
439        ascid.write('NODATA_value  %d\n' %NODATA_value)
440
441        #Get bounding polygon from mesh
442        #P = interp.mesh.get_boundary_polygon()
443        #inside_indices = inside_polygon(grid_points, P)
444
445        # change printoptions so that a long string of zeros in not
446        # summarized as [0.0, 0.0, 0.0, ... 0.0, 0.0, 0.0]
447        #printoptions = num.get_printoptions()
448        #num.set_printoptions(threshold=sys.maxint)
449
450        format = '%.'+'%g' % number_of_decimal_places +'e'
451        for i in range(nrows):
452            if verbose and i % ((nrows+10)/10) == 0:
453                log.critical('Doing row %d of %d' % (i, nrows))
454
455            base_index = (nrows-i-1)*ncols
456
457            slice = grid_values[base_index:base_index+ncols]
458
459            num.savetxt(ascid, slice.reshape(1,ncols), format, ' ' )
460           
461       
462        #Close
463        ascid.close()
464        fid.close()
465     
466        return basename_out
467
468
469
470def sww2dem_batch(basename_in, extra_name_out=None,
471                quantities=None, # defaults to elevation
472                reduction=None,
473                cellsize=10,
474                number_of_decimal_places=None,
475                NODATA_value=-9999,
476                easting_min=None,
477                easting_max=None,
478                northing_min=None,
479                northing_max=None,
480                verbose=False,
481                origin=None,
482                datum='WGS84',
483                format='ers'):
484    """Wrapper for sww2dem.
485    See sww2dem to find out what most of the parameters do. Note that since this
486    is a batch command, the normal filename naming conventions do not apply.
487
488    basename_in is a path to sww file/s, without the .sww extension.
489    extra_name_out is a postfix to add to the output filename.
490
491    Quantities is a list of quantities.  Each quantity will be
492    calculated for each sww file.
493
494    This returns the basenames of the files returned, which is made up
495    of the dir and all of the file name, except the extension.
496
497    This function returns the names of the files produced.
498
499    It will also produce as many output files as there are input sww files.
500    """
501
502    if quantities is None:
503        quantities = ['elevation']
504
505    if type(quantities) is str:
506            quantities = [quantities]
507
508    # How many sww files are there?
509    dir, base = os.path.split(basename_in)
510
511    iterate_over = get_all_swwfiles(dir, base, verbose)
512
513    if dir == "":
514        dir = "." # Unix compatibility
515
516    files_out = []
517    for sww_file in iterate_over:
518        for quantity in quantities:
519            if extra_name_out is None:
520                basename_out = sww_file + '_' + quantity
521            else:
522                basename_out = sww_file + '_' + quantity + '_' + extra_name_out
523
524            swwin = dir+os.sep+sww_file+'.sww'
525            demout = dir+os.sep+basename_out+'.'+format
526
527            if verbose:
528                log.critical('sww2dem: %s => %s' % (swwin, demout))
529
530            file_out = sww2dem(swwin,
531                               demout,
532                               quantity,
533                               reduction,
534                               cellsize,
535                               number_of_decimal_places,
536                               NODATA_value,
537                               easting_min,
538                               easting_max,
539                               northing_min,
540                               northing_max,
541                               verbose,
542                               origin,
543                               datum)
544                               
545            files_out.append(file_out)
546    return files_out
547
Note: See TracBrowser for help on using the repository browser.