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

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

Cleaned up unit tests to use API.

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