source: trunk/anuga_core/source/anuga/utilities/sww_merge.py @ 8809

Last change on this file since 8809 was 8809, checked in by steve, 12 years ago

Problem with netcdf4 and scientific python way of dealing with dimensions
File size: 22.7 KB
Line 
1"""
2    Merge a list of .sww files together into a single file.
3"""
4
5import numpy as num
6from anuga.utilities.numerical_tools import ensure_numeric
7
8from anuga.file.netcdf import NetCDFFile
9from anuga.config import netcdf_mode_r, netcdf_mode_w, netcdf_mode_a
10from anuga.config import netcdf_float, netcdf_float32, netcdf_int
11from anuga.file.sww import SWW_file, Write_sww
12
13def sww_merge(domain_global_name, np, verbose=False):
14
15    output = domain_global_name+".sww"
16    swwfiles = [ domain_global_name+"_P"+str(np)+"_"+str(v)+".sww" for v in range(np)]
17
18    _sww_merge(swwfiles, output, verbose)
19
20
21def sww_merge_parallel(domain_global_name, np, verbose=False, delete_old=False):
22
23    output = domain_global_name+".sww"
24    swwfiles = [ domain_global_name+"_P"+str(np)+"_"+str(v)+".sww" for v in range(np)]
25
26    fid = NetCDFFile(swwfiles[0], netcdf_mode_r)
27
28    try: # works with netcdf4
29        number_of_volumes = len(fid.dimensions['number_of_volumes'])
30        number_of_points = len(fid.dimensions['number_of_points'])
31    except: # works with scientific.io.netcdf
32        number_of_volumes = int(fid.dimensions['number_of_volumes'])
33        number_of_points = int(fid.dimensions['number_of_points'])
34
35    fid.close()
36
37    if 3*number_of_volumes == number_of_points:
38        _sww_merge_parallel_non_smooth(swwfiles, output, verbose, delete_old)
39    else:
40        _sww_merge_parallel_smooth(swwfiles, output, verbose, delete_old)
41       
42
43def _sww_merge(swwfiles, output, verbose=False):
44    """
45        Merge a list of sww files into a single file.
46       
47        May be useful for parallel runs. Note that colinear points and
48        edges are not merged: there will essentially be multiple meshes within
49        the one sww file.
50       
51        The sww files to be merged must have exactly the same timesteps. Note
52        that some advanced information and custom quantities may not be
53        exported.
54       
55        swwfiles is a list of .sww files to merge.
56        output is the output filename, including .sww extension.
57        verbose True to log output information
58    """
59
60    if verbose:
61        print "MERGING SWW Files"
62       
63    static_quantities = ['elevation']
64    dynamic_quantities = ['stage', 'xmomentum', 'ymomentum']
65   
66    first_file = True
67    tri_offset = 0
68    for filename in swwfiles:
69        if verbose:
70            print 'Reading file ', filename, ':'   
71   
72        fid = NetCDFFile(filename, netcdf_mode_r)
73
74
75
76       
77       
78        tris = fid.variables['volumes'][:]       
79         
80        if first_file:
81            times = fid.variables['time'][:]
82            x = []
83            y = []
84            out_tris = list(tris) 
85            out_s_quantities = {}
86            out_d_quantities = {}
87
88
89            xllcorner = fid.xllcorner
90            yllcorner = fid.yllcorner
91
92            order      = fid.order
93            xllcorner  = fid.xllcorner;
94            yllcorner  = fid.yllcorner ;
95            zone       = fid.zone;
96            false_easting  = fid.false_easting;
97            false_northing = fid.false_northing;
98            datum      = fid.datum;
99            projection = fid.projection;
100
101           
102            for quantity in static_quantities:
103                out_s_quantities[quantity] = []
104
105            # Quantities are stored as a 2D array of timesteps x data.
106            for quantity in dynamic_quantities:
107                out_d_quantities[quantity] = [ [] for _ in range(len(times))]
108                 
109            description = 'merged:' + getattr(fid, 'description')         
110            first_file = False
111        else:
112            for tri in tris:
113                # Advance new tri indices to point at newly appended points.
114                verts = [vertex+tri_offset for vertex in tri]
115                out_tris.append(verts)
116
117
118
119        try: # works with netcdf4
120            num_pts = len(fid.dimensions['number_of_points'])
121        except: # works with scientific.io.netcdf
122            num_pts = int(fid.dimensions['number_of_points'])
123
124        tri_offset += num_pts
125       
126        if verbose:
127            print '  new triangle index offset is ', tri_offset
128           
129        x.extend(list(fid.variables['x'][:]))
130        y.extend(list(fid.variables['y'][:]))
131       
132        # Grow the list of static quantities associated with the x,y points
133        for quantity in static_quantities:
134            out_s_quantities[quantity].extend(fid.variables[quantity][:])
135           
136        #Collate all dynamic quantities according to their timestep
137        for quantity in dynamic_quantities:
138            time_chunks = fid.variables[quantity][:]
139            for i, time_chunk in enumerate(time_chunks):
140                out_d_quantities[quantity][i].extend(time_chunk)           
141   
142    # Mash all points into a single big list   
143    points = [[xx, yy] for xx, yy in zip(x, y)]
144
145    points = num.asarray(points).astype(netcdf_float32)
146
147    fid.close()
148
149    #---------------------------
150    # Write out the SWW file
151    #---------------------------
152
153    if verbose:
154        print 'Writing file ', output, ':'
155    fido = NetCDFFile(output, netcdf_mode_w)
156    sww = Write_sww(static_quantities, dynamic_quantities)
157    sww.store_header(fido, times,
158                             len(out_tris),
159                             len(points),
160                             description=description,
161                             sww_precision=netcdf_float32)
162
163
164
165    from anuga.coordinate_transforms.geo_reference import Geo_reference
166    geo_reference = Geo_reference()
167   
168    sww.store_triangulation(fido, points, out_tris, points_georeference=geo_reference)
169
170    fido.order      = order
171    fido.xllcorner  = xllcorner;
172    fido.yllcorner  = yllcorner ;
173    fido.zone       = zone;
174    fido.false_easting  = false_easting;
175    fido.false_northing = false_northing;
176    fido.datum      = datum;
177    fido.projection = projection;
178       
179    sww.store_static_quantities(fido, verbose=verbose, **out_s_quantities)
180
181    # Write out all the dynamic quantities for each timestep
182    for q in dynamic_quantities:
183        q_values = out_d_quantities[q]
184        for i, time_slice in enumerate(q_values):
185            fido.variables[q][i] = num.array(time_slice, netcdf_float32)
186       
187        # This updates the _range values
188        q_range = fido.variables[q + Write_sww.RANGE][:]
189        q_values_min = num.min(q_values)
190        if q_values_min < q_range[0]:
191            fido.variables[q + Write_sww.RANGE][0] = q_values_min
192        q_values_max = num.max(q_values)
193        if q_values_max > q_range[1]:
194            fido.variables[q + Write_sww.RANGE][1] = q_values_max       
195
196                                       
197    fido.close()
198
199
200def _sww_merge_parallel_smooth(swwfiles, output,  verbose=False, delete_old=False):
201    """
202        Merge a list of sww files into a single file.
203       
204        Use to merge files created by parallel runs.
205
206        The sww files to be merged must have exactly the same timesteps.
207
208        It is assumed that the separate sww files have been stored in non_smooth
209        format.
210
211        Note that some advanced information and custom quantities may not be
212        exported.
213       
214        swwfiles is a list of .sww files to merge.
215        output is the output filename, including .sww extension.
216        verbose True to log output information
217    """
218
219    if verbose:
220        print "MERGING SWW Files"
221       
222   
223    first_file = True
224    tri_offset = 0
225    for filename in swwfiles:
226        if verbose:
227            print 'Reading file ', filename, ':'   
228   
229        fid = NetCDFFile(filename, netcdf_mode_r)
230         
231        if first_file:
232
233            times    = fid.variables['time'][:]
234            n_steps = len(times)
235            #number_of_timesteps = fid.dimensions['number_of_timesteps']
236            #print n_steps, number_of_timesteps
237            starttime = int(fid.starttime)
238           
239            out_s_quantities = {}
240            out_d_quantities = {}
241
242
243            xllcorner = fid.xllcorner
244            yllcorner = fid.yllcorner
245
246            number_of_global_triangles = int(fid.number_of_global_triangles)
247            number_of_global_nodes     = int(fid.number_of_global_nodes)
248
249            order      = fid.order
250            xllcorner  = fid.xllcorner;
251            yllcorner  = fid.yllcorner ;
252            zone       = fid.zone;
253            false_easting  = fid.false_easting;
254            false_northing = fid.false_northing;
255            datum      = fid.datum;
256            projection = fid.projection;
257
258            g_volumes = num.zeros((number_of_global_triangles,3),num.int)
259            g_x = num.zeros((number_of_global_nodes,),num.float32)
260            g_y = num.zeros((number_of_global_nodes,),num.float32)
261
262            g_points = num.zeros((number_of_global_nodes,2),num.float32)
263
264            quantities = set(['elevation', 'friction', 'stage', 'xmomentum',
265                              'ymomentum', 'xvelocity', 'yvelocity', 'height'])
266            variables = set(fid.variables.keys())
267
268            quantities = list(quantities & variables)
269           
270            static_quantities = []
271            dynamic_quantities = []
272
273            for quantity in quantities:
274                # Test if elevation is static
275                if n_steps == fid.variables[quantity].shape[0]:
276                    dynamic_quantities.append(quantity)
277                else:
278                    static_quantities.append(quantity)
279               
280            for quantity in static_quantities:
281                out_s_quantities[quantity] = num.zeros((number_of_global_nodes,),num.float32)
282
283            # Quantities are stored as a 2D array of timesteps x data.
284            for quantity in dynamic_quantities:
285                out_d_quantities[quantity] = \
286                      num.zeros((n_steps,number_of_global_nodes),num.float32)
287                 
288            description = 'merged:' + getattr(fid, 'description')         
289            first_file = False
290
291
292        # Read in from files and add to global arrays
293
294        tri_l2g  = fid.variables['tri_l2g'][:]
295        node_l2g = fid.variables['node_l2g'][:]
296        tri_full_flag = fid.variables['tri_full_flag'][:]
297        volumes = num.array(fid.variables['volumes'][:],dtype=num.int)
298        l_volumes = num.zeros_like(volumes)
299        l_old_volumes = num.zeros_like(volumes)
300
301
302        # Change the local node ids to global id in the
303        # volume array
304
305        # FIXME SR: Surely we can knock up a numpy way of doing this
306        #for i in range(len(l_volumes)):
307        #    g_n0 = node_l2g[volumes[i,0]]
308        #    g_n1 = node_l2g[volumes[i,1]]
309        #    g_n2 = node_l2g[volumes[i,2]]
310        #
311        #    l_old_volumes[i,:] = [g_n0,g_n1,g_n2]
312
313        g_n0 = node_l2g[volumes[:,0]].reshape(-1,1)
314        g_n1 = node_l2g[volumes[:,1]].reshape(-1,1)
315        g_n2 = node_l2g[volumes[:,2]].reshape(-1,1)
316
317        #print g_n0.shape
318        l_volumes = num.hstack((g_n0,g_n1,g_n2))
319
320        #assert num.allclose(l_volumes, l_old_volumes)
321
322        # Just pick out the full triangles
323        ftri_l2g = num.compress(tri_full_flag, tri_l2g)
324
325        #print l_volumes
326        #print tri_full_flag
327        #print tri_l2g
328        #print ftri_l2g
329
330        fg_volumes = num.compress(tri_full_flag,l_volumes,axis=0)
331        g_volumes[ftri_l2g] = fg_volumes
332
333
334
335
336        #g_x[node_l2g] = fid.variables['x']
337        #g_y[node_l2g] = fid.variables['y']
338
339        g_points[node_l2g,0] = fid.variables['x']
340        g_points[node_l2g,1] = fid.variables['y']
341       
342
343        #print number_of_timesteps
344
345
346        # FIXME SR: It seems that some of the "ghost" node quantity values
347        # are being storded. We should only store those nodes which are associated with
348        # full triangles. So we need an index array of "full" nodes, ie those in
349        # full triangles
350
351        #use numpy.compress and numpy.unique to get "full nodes
352
353        f_volumes = num.compress(tri_full_flag,volumes,axis=0)
354        fl_nodes = num.unique(f_volumes)
355        f_node_l2g = node_l2g[fl_nodes]
356
357        #print len(node_l2g)
358        #print len(fl_nodes)
359
360        # Read in static quantities
361        for quantity in static_quantities:
362            #out_s_quantities[quantity][node_l2g] = \
363            #             num.array(fid.variables[quantity],dtype=num.float32)
364            q = fid.variables[quantity]
365            #print quantity, q.shape
366            out_s_quantities[quantity][f_node_l2g] = \
367                         num.array(q,dtype=num.float32)[fl_nodes]
368
369       
370        #Collate all dynamic quantities according to their timestep
371        for quantity in dynamic_quantities:
372            q = fid.variables[quantity]
373            #print q.shape
374            for i in range(n_steps):
375                #out_d_quantities[quantity][i][node_l2g] = \
376                #           num.array(q[i],dtype=num.float32)
377                out_d_quantities[quantity][i][f_node_l2g] = \
378                           num.array(q[i],dtype=num.float32)[fl_nodes]
379
380
381
382
383        fid.close()
384
385
386    #---------------------------
387    # Write out the SWW file
388    #---------------------------
389    #print g_points.shape
390
391    #print number_of_global_triangles
392    #print number_of_global_nodes
393
394
395    if verbose:
396            print 'Writing file ', output, ':'
397    fido = NetCDFFile(output, netcdf_mode_w)
398    sww = Write_sww(static_quantities, dynamic_quantities)
399    sww.store_header(fido, starttime,
400                             number_of_global_triangles,
401                             number_of_global_nodes,
402                             description=description,
403                             sww_precision=netcdf_float32)
404
405
406
407    from anuga.coordinate_transforms.geo_reference import Geo_reference
408    geo_reference = Geo_reference()
409   
410    sww.store_triangulation(fido, g_points, g_volumes, points_georeference=geo_reference)
411
412    fido.order      = order
413    fido.xllcorner  = xllcorner;
414    fido.yllcorner  = yllcorner ;
415    fido.zone       = zone;
416    fido.false_easting  = false_easting;
417    fido.false_northing = false_northing;
418    fido.datum      = datum;
419    fido.projection = projection;
420       
421    sww.store_static_quantities(fido, verbose=verbose, **out_s_quantities)
422
423    # Write out all the dynamic quantities for each timestep
424
425    for i in range(n_steps):
426        fido.variables['time'][i] = times[i]
427       
428    for q in dynamic_quantities:
429        q_values = out_d_quantities[q]
430        for i in range(n_steps):
431            fido.variables[q][i] = q_values[i]
432       
433        # This updates the _range values
434        q_range = fido.variables[q + Write_sww.RANGE][:]
435        q_values_min = num.min(q_values)
436        if q_values_min < q_range[0]:
437            fido.variables[q + Write_sww.RANGE][0] = q_values_min
438        q_values_max = num.max(q_values)
439        if q_values_max > q_range[1]:
440            fido.variables[q + Write_sww.RANGE][1] = q_values_max       
441
442                                       
443    #print out_s_quantities
444    #print out_d_quantities
445   
446    #print g_x
447    #print g_y
448
449    #print g_volumes
450   
451    fido.close()
452   
453    if delete_old:
454        import os
455        for filename in swwfiles:
456
457            if verbose:
458                print 'Deleting file ', filename, ':'
459            os.remove(filename)
460
461def _sww_merge_parallel_non_smooth(swwfiles, output,  verbose=False, delete_old=False):
462    """
463        Merge a list of sww files into a single file.
464
465        Used to merge files created by parallel runs.
466
467        The sww files to be merged must have exactly the same timesteps.
468
469        It is assumed that the separate sww files have been stored in non_smooth
470        format.
471
472        Note that some advanced information and custom quantities may not be
473        exported.
474
475        swwfiles is a list of .sww files to merge.
476        output is the output filename, including .sww extension.
477        verbose True to log output information
478    """
479
480    if verbose:
481        print "MERGING SWW Files"
482
483
484    first_file = True
485    tri_offset = 0
486    for filename in swwfiles:
487        if verbose:
488            print 'Reading file ', filename, ':'
489
490        fid = NetCDFFile(filename, netcdf_mode_r)
491
492        if first_file:
493
494            times    = fid.variables['time'][:]
495            n_steps = len(times)
496            number_of_timesteps = fid.dimensions['number_of_timesteps']
497            #print n_steps, number_of_timesteps
498            starttime = int(fid.starttime)
499
500            out_s_quantities = {}
501            out_d_quantities = {}
502
503
504            xllcorner = fid.xllcorner
505            yllcorner = fid.yllcorner
506
507            number_of_global_triangles = int(fid.number_of_global_triangles)
508            number_of_global_nodes     = int(fid.number_of_global_nodes)
509            number_of_global_triangle_vertices = 3*number_of_global_triangles
510
511
512            order      = fid.order
513            xllcorner  = fid.xllcorner;
514            yllcorner  = fid.yllcorner ;
515            zone       = fid.zone;
516            false_easting  = fid.false_easting;
517            false_northing = fid.false_northing;
518            datum      = fid.datum;
519            projection = fid.projection;
520
521            g_volumes = num.arange(number_of_global_triangles*3).reshape(-1,3)
522
523
524
525            g_x = num.zeros((number_of_global_triangle_vertices,),num.float32)
526            g_y = num.zeros((number_of_global_triangle_vertices,),num.float32)
527
528            g_points = num.zeros((number_of_global_triangle_vertices,2),num.float32)
529
530
531            quantities = set(['elevation', 'friction', 'stage', 'xmomentum',
532                              'ymomentum', 'xvelocity', 'yvelocity', 'height'])
533            variables = set(fid.variables.keys())
534
535            quantities = list(quantities & variables)
536
537            static_quantities = []
538            dynamic_quantities = []
539
540            for quantity in quantities:
541                # Test if elevation is static
542                if n_steps == fid.variables[quantity].shape[0]:
543                    dynamic_quantities.append(quantity)
544                else:
545                    static_quantities.append(quantity)
546
547            for quantity in static_quantities:
548                out_s_quantities[quantity] = num.zeros((3*number_of_global_triangles,),num.float32)
549
550            # Quantities are stored as a 2D array of timesteps x data.
551            for quantity in dynamic_quantities:
552                out_d_quantities[quantity] = \
553                      num.zeros((n_steps,3*number_of_global_triangles),num.float32)
554
555            description = 'merged:' + getattr(fid, 'description')
556            first_file = False
557
558
559        # Read in from files and add to global arrays
560
561        tri_l2g  = fid.variables['tri_l2g'][:]
562        node_l2g = fid.variables['node_l2g'][:]
563        tri_full_flag = fid.variables['tri_full_flag'][:]
564
565
566
567
568        f_ids = num.argwhere(tri_full_flag==1).reshape(-1,)
569
570        f_gids = tri_l2g[f_ids]
571
572
573        g_vids = (3*f_gids.reshape(-1,1) + num.array([0,1,2])).reshape(-1,)
574        l_vids = (3*f_ids.reshape(-1,1) + num.array([0,1,2])).reshape(-1,)
575
576
577        l_x = num.array(fid.variables['x'][:],dtype=num.float32)
578        l_y = num.array(fid.variables['y'][:],dtype=num.float32)
579
580       
581        g_x[g_vids] = l_x[l_vids]
582        g_y[g_vids] = l_y[l_vids]
583
584        g_points[g_vids,0] = g_x[g_vids]
585        g_points[g_vids,1] = g_y[g_vids]
586
587
588        # Read in static quantities
589        for quantity in static_quantities:
590            #out_s_quantities[quantity][node_l2g] = \
591            #             num.array(fid.variables[quantity],dtype=num.float32)
592            q = fid.variables[quantity]
593            #print quantity, q.shape
594            out_s_quantities[quantity][g_vids] = \
595                         num.array(q,dtype=num.float32)[l_vids]
596
597
598        #Collate all dynamic quantities according to their timestep
599        for quantity in dynamic_quantities:
600            q = fid.variables[quantity]
601            #print q.shape
602            for i in range(n_steps):
603                #out_d_quantities[quantity][i][node_l2g] = \
604                #           num.array(q[i],dtype=num.float32)
605                out_d_quantities[quantity][i][g_vids] = \
606                           num.array(q[i],dtype=num.float32)[l_vids]
607
608
609        fid.close()
610
611    #---------------------------
612    # Write out the SWW file
613    #---------------------------
614
615    if verbose:
616            print 'Writing file ', output, ':'
617
618    fido = NetCDFFile(output, netcdf_mode_w)
619    sww = Write_sww(static_quantities, dynamic_quantities)
620    sww.store_header(fido, starttime,
621                             number_of_global_triangles,
622                             number_of_global_triangles*3,
623                             description=description,
624                             sww_precision=netcdf_float32)
625
626
627    from anuga.coordinate_transforms.geo_reference import Geo_reference
628    geo_reference = Geo_reference()
629
630    sww.store_triangulation(fido, g_points, g_volumes, points_georeference=geo_reference)
631
632    fido.order      = order
633    fido.xllcorner  = xllcorner;
634    fido.yllcorner  = yllcorner ;
635    fido.zone       = zone;
636    fido.false_easting  = false_easting;
637    fido.false_northing = false_northing;
638    fido.datum      = datum;
639    fido.projection = projection;
640
641    sww.store_static_quantities(fido, verbose=verbose, **out_s_quantities)
642
643    # Write out all the dynamic quantities for each timestep
644
645    for i in range(n_steps):
646        fido.variables['time'][i] = times[i]
647
648    for q in dynamic_quantities:
649        q_values = out_d_quantities[q]
650        for i in range(n_steps):
651            fido.variables[q][i] = q_values[i]
652
653        # This updates the _range values
654        q_range = fido.variables[q + Write_sww.RANGE][:]
655        q_values_min = num.min(q_values)
656        if q_values_min < q_range[0]:
657            fido.variables[q + Write_sww.RANGE][0] = q_values_min
658        q_values_max = num.max(q_values)
659        if q_values_max > q_range[1]:
660            fido.variables[q + Write_sww.RANGE][1] = q_values_max
661
662
663    fido.close()
664
665    if delete_old:
666        import os
667        for filename in swwfiles:
668
669            if verbose:
670                print 'Deleting file ', filename, ':'
671            os.remove(filename)
672
673
674
675
676
677
678if __name__ == "__main__":
679
680    import argparse
681    from anuga.anuga_exceptions import ANUGAError
682
683
684    parser = argparse.ArgumentParser(description='Merge sww files created from parallel run')
685    parser.add_argument('-np', type=int, default = 4,
686                   help='number of processors used to produce sww files')
687    parser.add_argument('-f', type=str, default="domain",
688                   help='domain global name')
689    parser.add_argument('-v', nargs='?', type=bool, const=True, default=False,
690                   help='verbosity')
691    parser.add_argument('-delete_old', nargs='?', type=bool, const=True, default=False,
692                   help='Flag to delete the input files')
693    args = parser.parse_args()
694
695    np = args.np
696    domain_global_name = args.f
697    verbose = args.v
698    delete_old = args.delete_old
699
700
701    try:
702        sww_merge_parallel(domain_global_name, np, verbose, delete_old)
703    except:
704        msg = 'ERROR: When merging sww files %s '% domain_global_name
705        print msg
706        raise
Note: See TracBrowser for help on using the repository browser.