source: trunk/anuga_core/source/anuga/file/sww.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: 44.5 KB
Line 
1""" Classes to read an SWW file.
2"""
3
4from anuga.coordinate_transforms.geo_reference import Geo_reference
5from anuga.config import netcdf_mode_r, netcdf_mode_w, netcdf_mode_a
6from anuga.config import netcdf_float, netcdf_float32, netcdf_int
7from anuga.config import max_float
8from anuga.utilities.numerical_tools import ensure_numeric
9import anuga.utilities.log as log
10from Scientific.IO.NetCDF import NetCDFFile
11
12from anuga.coordinate_transforms.geo_reference import \
13        ensure_geo_reference
14
15from file_utils import create_filename
16import numpy as num
17
18##
19# @brief Generic class for storing output to e.g. visualisation or checkpointing
20class Data_format:
21    """Generic interface to data formats
22    """
23
24    ##
25    # @brief Instantiate this instance.
26    # @param domain
27    # @param extension
28    # @param mode The mode of the underlying file.
29    def __init__(self, domain, extension, mode=netcdf_mode_w):
30        assert mode[0] in ['r', 'w', 'a'], \
31               "Mode %s must be either:\n" % mode + \
32               "   'w' (write)\n" + \
33               "   'r' (read)\n" + \
34               "   'a' (append)"
35
36        #Create filename
37        self.filename = create_filename(domain.get_datadir(),
38                                        domain.get_name(), extension)
39
40        self.timestep = 0
41        self.domain = domain
42
43        # Exclude ghosts in case this is a parallel domain
44        self.number_of_nodes = domain.number_of_full_nodes
45        self.number_of_volumes = domain.number_of_full_triangles
46        #self.number_of_volumes = len(domain)
47
48        #FIXME: Should we have a general set_precision function?
49
50
51class SWW_file(Data_format):
52    """Interface to native NetCDF format (.sww) for storing model output
53
54    There are two kinds of data
55
56    1: Constant data: Vertex coordinates and field values. Stored once
57    2: Variable data: Conserved quantities. Stored once per timestep.
58
59    All data is assumed to reside at vertex locations.
60    """
61
62    ##
63    # @brief Instantiate this instance.
64    # @param domain ??
65    # @param mode Mode of the underlying data file.
66    # @param max_size ??
67    # @param recursion ??
68    # @note Prepare the underlying data file if mode starts with 'w'.
69    def __init__(self, domain, 
70                 mode=netcdf_mode_w, max_size=2000000000, recursion=False):
71        from Scientific.IO.NetCDF import NetCDFFile
72
73        self.precision = netcdf_float32 # Use single precision for quantities
74        self.recursion = recursion
75        self.mode = mode
76        if hasattr(domain, 'max_size'):
77            self.max_size = domain.max_size # File size max is 2Gig
78        else:
79            self.max_size = max_size
80        if hasattr(domain, 'minimum_storable_height'):
81            self.minimum_storable_height = domain.minimum_storable_height
82        else:
83            self.minimum_storable_height = default_minimum_storable_height
84
85        # Call parent constructor
86        Data_format.__init__(self, domain, 'sww', mode)
87
88        # Get static and dynamic quantities from domain
89        static_quantities = []
90        dynamic_quantities = []
91       
92        for q in domain.quantities_to_be_stored:
93            flag = domain.quantities_to_be_stored[q]
94       
95            msg = 'Quantity %s is requested to be stored ' % q
96            msg += 'but it does not exist in domain.quantities'
97            assert q in domain.quantities, msg
98       
99            assert flag in [1,2]
100            if flag == 1: static_quantities.append(q)
101            if flag == 2: dynamic_quantities.append(q)               
102                       
103       
104        # NetCDF file definition
105        fid = NetCDFFile(self.filename, mode)
106        if mode[0] == 'w':
107            description = 'Output from anuga.abstract_2d_finite_volumes ' \
108                          'suitable for plotting'
109                         
110            self.writer = Write_sww(static_quantities, dynamic_quantities)
111            self.writer.store_header(fid,
112                                     domain.starttime,
113                                     self.number_of_volumes,
114                                     self.domain.number_of_full_nodes,
115                                     description=description,
116                                     smoothing=domain.smooth,
117                                     order=domain.default_order,
118                                     sww_precision=self.precision)
119
120            # Extra optional information
121            if hasattr(domain, 'texture'):
122                fid.texture = domain.texture
123
124            if domain.quantities_to_be_monitored is not None:
125                fid.createDimension('singleton', 1)
126                fid.createDimension('two', 2)
127
128                poly = domain.monitor_polygon
129                if poly is not None:
130                    N = len(poly)
131                    fid.createDimension('polygon_length', N)
132                    fid.createVariable('extrema.polygon',
133                                       self.precision,
134                                       ('polygon_length', 'two'))
135                    fid.variables['extrema.polygon'][:] = poly
136
137                interval = domain.monitor_time_interval
138                if interval is not None:
139                    fid.createVariable('extrema.time_interval',
140                                       self.precision,
141                                       ('two',))
142                    fid.variables['extrema.time_interval'][:] = interval
143
144                for q in domain.quantities_to_be_monitored:
145                    fid.createVariable(q + '.extrema', self.precision,
146                                       ('numbers_in_range',))
147                    fid.createVariable(q + '.min_location', self.precision,
148                                       ('numbers_in_range',))
149                    fid.createVariable(q + '.max_location', self.precision,
150                                       ('numbers_in_range',))
151                    fid.createVariable(q + '.min_time', self.precision,
152                                       ('singleton',))
153                    fid.createVariable(q + '.max_time', self.precision,
154                                       ('singleton',))
155
156        fid.close()
157
158    ##
159    # @brief Store connectivity data into the underlying data file.
160    def store_connectivity(self):
161        """Store information about nodes, triangles and static quantities
162
163        Writes x,y coordinates of triangles and their connectivity.
164       
165        Store also any quantity that has been identified as static.
166        """
167
168        # FIXME: Change name to reflect the fact thta this function
169        # stores both connectivity (triangulation) and static quantities
170       
171        from Scientific.IO.NetCDF import NetCDFFile
172
173        domain = self.domain
174
175        # append to the NetCDF file
176        fid = NetCDFFile(self.filename, netcdf_mode_a)
177
178        # Get X, Y from one (any) of the quantities
179        Q = domain.quantities.values()[0]
180        X,Y,_,V = Q.get_vertex_values(xy=True, precision=self.precision)
181
182        # store the connectivity data
183        points = num.concatenate((X[:,num.newaxis],Y[:,num.newaxis]), axis=1)
184        self.writer.store_triangulation(fid,
185                                        points,
186                                        V.astype(num.float32),
187                                        points_georeference=\
188                                            domain.geo_reference)
189
190
191        # Get names of static quantities
192        static_quantities = {}
193        for name in self.writer.static_quantities:
194            Q = domain.quantities[name]
195            A, _ = Q.get_vertex_values(xy=False, 
196                                       precision=self.precision)
197            static_quantities[name] = A
198       
199        # Store static quantities       
200        self.writer.store_static_quantities(fid, **static_quantities)
201                                           
202        fid.close()
203
204    ##
205    # @brief Store time and time dependent quantities
206    # to the underlying data file.
207    def store_timestep(self):
208        """Store time and time dependent quantities
209        """
210
211        from Scientific.IO.NetCDF import NetCDFFile
212        import types
213        from time import sleep
214        from os import stat
215
216        # Get NetCDF
217        retries = 0
218        file_open = False
219        while not file_open and retries < 10:
220            try:
221                # Open existing file
222                fid = NetCDFFile(self.filename, netcdf_mode_a)
223            except IOError:
224                # This could happen if someone was reading the file.
225                # In that case, wait a while and try again
226                msg = 'Warning (store_timestep): File %s could not be opened' \
227                      % self.filename
228                msg += ' - trying step %s again' % self.domain.time
229                log.critical(msg)
230                retries += 1
231                sleep(1)
232            else:
233                file_open = True
234
235        if not file_open:
236            msg = 'File %s could not be opened for append' % self.filename
237            raise DataFileNotOpenError, msg
238
239        # Check to see if the file is already too big:
240        time = fid.variables['time']
241        i = len(time) + 1
242        file_size = stat(self.filename)[6]
243        file_size_increase = file_size / i
244        if file_size + file_size_increase > self.max_size * 2**self.recursion:
245            # In order to get the file name and start time correct,
246            # I change the domain.filename and domain.starttime.
247            # This is the only way to do this without changing
248            # other modules (I think).
249
250            # Write a filename addon that won't break the anuga viewers
251            # (10.sww is bad)
252            filename_ext = '_time_%s' % self.domain.time
253            filename_ext = filename_ext.replace('.', '_')
254
255            # Remember the old filename, then give domain a
256            # name with the extension
257            old_domain_filename = self.domain.get_name()
258            if not self.recursion:
259                self.domain.set_name(old_domain_filename + filename_ext)
260
261            # Temporarily change the domain starttime to the current time
262            old_domain_starttime = self.domain.starttime
263            self.domain.starttime = self.domain.get_time()
264
265            # Build a new data_structure.
266            next_data_structure = SWW_file(self.domain, mode=self.mode,
267                                           max_size=self.max_size,
268                                           recursion=self.recursion+1)
269            if not self.recursion:
270                log.critical('    file_size = %s' % file_size)
271                log.critical('    saving file to %s'
272                             % next_data_structure.filename) 
273
274            # Set up the new data_structure
275            self.domain.writer = next_data_structure
276
277            # Store connectivity and first timestep
278            next_data_structure.store_connectivity()
279            next_data_structure.store_timestep()
280            fid.sync()
281            fid.close()
282
283            # Restore the old starttime and filename
284            self.domain.starttime = old_domain_starttime
285            self.domain.set_name(old_domain_filename)
286        else:
287            self.recursion = False
288            domain = self.domain
289
290            # Get the variables
291            time = fid.variables['time']
292            i = len(time)
293             
294            if 'stage' in self.writer.dynamic_quantities:           
295                # Select only those values for stage,
296                # xmomentum and ymomentum (if stored) where
297                # depth exceeds minimum_storable_height
298                #
299                # In this branch it is assumed that elevation
300                # is also available as a quantity           
301           
302                Q = domain.quantities['stage']
303                w, _ = Q.get_vertex_values(xy=False)
304               
305                Q = domain.quantities['elevation']
306                z, _ = Q.get_vertex_values(xy=False)               
307               
308                storable_indices = (w-z >= self.minimum_storable_height)
309            else:
310                # Very unlikely branch
311                storable_indices = None # This means take all
312           
313           
314            # Now store dynamic quantities
315            dynamic_quantities = {}
316            for name in self.writer.dynamic_quantities:
317                netcdf_array = fid.variables[name]
318               
319                Q = domain.quantities[name]
320                A, _ = Q.get_vertex_values(xy=False,
321                                           precision=self.precision)
322
323                if storable_indices is not None:
324                    if name == 'stage':
325                        A = num.choose(storable_indices, (z, A))
326
327                    if name in ['xmomentum', 'ymomentum']:
328                        # Get xmomentum where depth exceeds
329                        # minimum_storable_height
330                       
331                        # Define a zero vector of same size and type as A
332                        # for use with momenta
333                        null = num.zeros(num.size(A), A.dtype.char)
334                        A = num.choose(storable_indices, (null, A))
335               
336                dynamic_quantities[name] = A
337               
338                                       
339            # Store dynamic quantities
340            self.writer.store_quantities(fid,
341                                         time=self.domain.time,
342                                         sww_precision=self.precision,
343                                         **dynamic_quantities)
344
345
346            # Update extrema if requested
347            domain = self.domain
348            if domain.quantities_to_be_monitored is not None:
349                for q, info in domain.quantities_to_be_monitored.items():
350                    if info['min'] is not None:
351                        fid.variables[q + '.extrema'][0] = info['min']
352                        fid.variables[q + '.min_location'][:] = \
353                                        info['min_location']
354                        fid.variables[q + '.min_time'][0] = info['min_time']
355
356                    if info['max'] is not None:
357                        fid.variables[q + '.extrema'][1] = info['max']
358                        fid.variables[q + '.max_location'][:] = \
359                                        info['max_location']
360                        fid.variables[q + '.max_time'][0] = info['max_time']
361
362            # Flush and close
363            fid.sync()
364            fid.close()
365
366
367##
368# @brief Class to open an sww file so that domain can be populated with quantity values
369class Read_sww:
370
371    def __init__(self, source):
372        """The source parameter is assumed to be a NetCDF sww file.
373        """
374
375        self.source = source
376
377        self.frame_number = 0
378
379        fin = NetCDFFile(self.source, 'r')
380
381        self.time = num.array(fin.variables['time'], num.float)
382        self.last_frame_number = self.time.shape[0] - 1
383
384        self.frames = num.arange(self.last_frame_number+1)
385
386        fin.close()
387       
388        self.read_mesh()
389
390        self.quantities = {}
391
392        self.read_quantities()
393
394
395    def read_mesh(self):
396        """ Read and store the mesh data contained within this sww file.
397        """
398        fin = NetCDFFile(self.source, 'r')
399
400        self.vertices = num.array(fin.variables['volumes'], num.int)
401       
402        self.x = x = num.array(fin.variables['x'], num.float)
403        self.y = y = num.array(fin.variables['y'], num.float)
404
405        assert len(self.x) == len(self.y)
406       
407        self.xmin = num.min(x)
408        self.xmax = num.max(x)
409        self.ymin = num.min(y)
410        self.ymax = num.max(y)
411
412
413        fin.close()
414       
415    def read_quantities(self, frame_number=0):
416        """
417        Read the quantities contained in this file.
418        frame_number is the time index to load.
419        """
420        assert frame_number >= 0 and frame_number <= self.last_frame_number
421
422        self.frame_number = frame_number
423
424        M = len(self.x)/3
425       
426        fin = NetCDFFile(self.source, 'r')
427       
428        for q in filter(lambda n:n != 'x' and n != 'y' and n != 'time' and n != 'volumes' and \
429                        '_range' not in n, \
430                        fin.variables.keys()):
431            if len(fin.variables[q].shape) == 1: # Not a time-varying quantity
432                self.quantities[q] = num.ravel(num.array(fin.variables[q], num.float)).reshape(M,3)
433            else: # Time-varying, get the current timestep data
434                self.quantities[q] = num.array(fin.variables[q][self.frame_number], num.float).reshape(M,3)
435        fin.close()
436        return self.quantities
437
438    def get_bounds(self):
439        """
440            Get the bounding rect around the mesh.
441        """
442        return [self.xmin, self.xmax, self.ymin, self.ymax]
443
444    def get_last_frame_number(self):
445        """
446            Return the last loaded frame index.
447        """
448        return self.last_frame_number
449
450    def get_time(self):
451        """
452            Get time at the current frame num, in secs.
453        """
454        return self.time[self.frame_number]
455
456
457class Write_sww:
458    """
459        A class to write an SWW file.
460       
461        It is domain agnostic, and requires all the data to be fed in
462        manually.
463    """
464    RANGE = '_range'
465    EXTREMA = ':extrema'
466
467    def __init__(self, static_quantities, dynamic_quantities):
468        """Initialise Write_sww with two list af quantity names:
469       
470        static_quantities (e.g. elevation or friction):
471            Stored once at the beginning of the simulation in a 1D array
472            of length number_of_points   
473        dynamic_quantities (e.g stage):
474            Stored every timestep in a 2D array with
475            dimensions number_of_points X number_of_timesteps       
476       
477        """
478        self.static_quantities = static_quantities   
479        self.dynamic_quantities = dynamic_quantities
480
481
482    def store_header(self,
483                     outfile,
484                     times,
485                     number_of_volumes,
486                     number_of_points,
487                     description='Generated by ANUGA',
488                     smoothing=True,
489                     order=1,
490                     sww_precision=netcdf_float32,
491                     verbose=False):
492        """Write an SWW file header.
493
494        Writes the first section of the .sww file.
495
496        outfile - the open file that will be written
497        times - A list of the time slice times OR a start time
498        Note, if a list is given the info will be made relative.
499        number_of_volumes - the number of triangles
500        number_of_points - the number of vertices in the mesh
501        """
502
503        from anuga.abstract_2d_finite_volumes.util \
504            import get_revision_number
505
506        outfile.institution = 'Geoscience Australia'
507        outfile.description = description
508
509        # For sww compatibility
510        if smoothing is True:
511            # Smoothing to be depreciated
512            outfile.smoothing = 'Yes'
513            outfile.vertices_are_stored_uniquely = 'False'
514        else:
515            # Smoothing to be depreciated
516            outfile.smoothing = 'No'
517            outfile.vertices_are_stored_uniquely = 'True'
518        outfile.order = order
519
520        try:
521            revision_number = get_revision_number()
522        except:
523            # This will be triggered if the system cannot get the SVN
524            # revision number.
525            revision_number = None
526        # Allow None to be stored as a string
527        outfile.revision_number = str(revision_number)
528
529        # This is being used to seperate one number from a list.
530        # what it is actually doing is sorting lists from numeric arrays.
531        if isinstance(times, (list, num.ndarray)):
532            number_of_times = len(times)
533            times = ensure_numeric(times)
534            if number_of_times == 0:
535                starttime = 0
536            else:
537                starttime = times[0]
538                times = times - starttime  #Store relative times
539        else:
540            number_of_times = 0
541            starttime = times
542
543
544        outfile.starttime = starttime
545
546        # dimension definitions
547        outfile.createDimension('number_of_volumes', number_of_volumes)
548        outfile.createDimension('number_of_vertices', 3)
549        outfile.createDimension('numbers_in_range', 2)
550
551        if smoothing is True:
552            outfile.createDimension('number_of_points', number_of_points)
553            # FIXME(Ole): This will cause sww files for parallel domains to
554            # have ghost nodes stored (but not used by triangles).
555            # To clean this up, we have to change get_vertex_values and
556            # friends in quantity.py (but I can't be bothered right now)
557        else:
558            outfile.createDimension('number_of_points', 3*number_of_volumes)
559
560        outfile.createDimension('number_of_timesteps', number_of_times)
561
562        # variable definitions
563        outfile.createVariable('x', sww_precision, ('number_of_points',))
564        outfile.createVariable('y', sww_precision, ('number_of_points',))
565
566        outfile.createVariable('volumes', netcdf_int, ('number_of_volumes',
567                                                       'number_of_vertices'))
568
569        # Doing sww_precision instead of Float gives cast errors.
570        outfile.createVariable('time', netcdf_float,
571                               ('number_of_timesteps',))
572
573                               
574        for q in self.static_quantities:
575           
576            outfile.createVariable(q, sww_precision,
577                                   ('number_of_points',))
578           
579            outfile.createVariable(q + Write_sww.RANGE, sww_precision,
580                                   ('numbers_in_range',))
581                                   
582            # Initialise ranges with small and large sentinels.
583            # If this was in pure Python we could have used None sensibly
584            outfile.variables[q+Write_sww.RANGE][0] = max_float  # Min
585            outfile.variables[q+Write_sww.RANGE][1] = -max_float # Max
586
587        #if 'elevation' in self.static_quantities:   
588        #    # FIXME: Backwards compat - get rid of z once old view has retired
589        #    outfile.createVariable('z', sww_precision,
590        #                           ('number_of_points',))
591                               
592        for q in self.dynamic_quantities:
593            outfile.createVariable(q, sww_precision, ('number_of_timesteps',
594                                                      'number_of_points'))
595            outfile.createVariable(q + Write_sww.RANGE, sww_precision,
596                                   ('numbers_in_range',))
597
598            # Initialise ranges with small and large sentinels.
599            # If this was in pure Python we could have used None sensibly
600            outfile.variables[q+Write_sww.RANGE][0] = max_float  # Min
601            outfile.variables[q+Write_sww.RANGE][1] = -max_float # Max
602
603        if isinstance(times, (list, num.ndarray)):
604            outfile.variables['time'][:] = times    # Store time relative
605
606        if verbose:
607            log.critical('------------------------------------------------')
608            log.critical('Statistics:')
609            log.critical('    t in [%f, %f], len(t) == %d'
610                         % (num.min(times), num.max(times), len(times.flat)))
611
612
613    def store_triangulation(self,
614                            outfile,
615                            points_utm,
616                            volumes,
617                            zone=None, 
618                            new_origin=None,
619                            points_georeference=None, 
620                            verbose=False):
621        """
622        Store triangulation data in the underlying file.
623       
624        Stores the points and triangle indices in the sww file
625       
626        outfile Open handle to underlying file.
627
628        new_origin georeference that the points can be set to.
629
630        points_georeference The georeference of the points_utm.
631
632        verbose True if this function is to be verbose.
633
634        new_origin - qa georeference that the points can be set to. (Maybe
635        do this before calling this function.)
636
637        points_utm - currently a list or array of the points in UTM.
638        points_georeference - the georeference of the points_utm
639
640        How about passing new_origin and current_origin.
641        If you get both, do a convertion from the old to the new.
642
643        If you only get new_origin, the points are absolute,
644        convert to relative
645
646        if you only get the current_origin the points are relative, store
647        as relative.
648
649        if you get no georefs create a new georef based on the minimums of
650        points_utm.  (Another option would be to default to absolute)
651
652        Yes, and this is done in another part of the code.
653        Probably geospatial.
654
655        If you don't supply either geo_refs, then supply a zone. If not
656        the default zone will be used.
657
658        precon:
659            header has been called.
660        """
661
662        number_of_points = len(points_utm)
663        volumes = num.array(volumes)
664        points_utm = num.array(points_utm)
665
666        # Given the two geo_refs and the points, do the stuff
667        # described in the method header
668        # if this is needed else where, pull out as a function
669        points_georeference = ensure_geo_reference(points_georeference)
670        new_origin = ensure_geo_reference(new_origin)
671        if new_origin is None and points_georeference is not None:
672            points = points_utm
673            geo_ref = points_georeference
674        else:
675            if new_origin is None:
676                new_origin = Geo_reference(zone, min(points_utm[:,0]),
677                                                 min(points_utm[:,1]))
678            points = new_origin.change_points_geo_ref(points_utm,
679                                                      points_georeference)
680            geo_ref = new_origin
681
682        # At this stage I need a georef and points
683        # the points are relative to the georef
684        geo_ref.write_NetCDF(outfile)
685
686        # This will put the geo ref in the middle
687        #geo_ref = Geo_reference(refzone,(max(x)+min(x))/2.0,(max(x)+min(y))/2.)
688
689        x =  points[:,0]
690        y =  points[:,1]
691
692        if verbose:
693            log.critical('------------------------------------------------')
694            log.critical('More Statistics:')
695            log.critical('  Extent (/lon):')
696            log.critical('    x in [%f, %f], len(lat) == %d'
697                         % (min(x), max(x), len(x)))
698            log.critical('    y in [%f, %f], len(lon) == %d'
699                         % (min(y), max(y), len(y)))
700            #log.critical('    z in [%f, %f], len(z) == %d'
701            #             % (min(elevation), max(elevation), len(elevation)))
702            log.critical('geo_ref: %s' % str(geo_ref))
703            log.critical('------------------------------------------------')
704
705        outfile.variables['x'][:] = points[:,0] #- geo_ref.get_xllcorner()
706        outfile.variables['y'][:] = points[:,1] #- geo_ref.get_yllcorner()
707        outfile.variables['volumes'][:] = volumes.astype(num.int32) #On Opteron 64
708
709
710    def store_static_quantities(self, 
711                                outfile, 
712                                sww_precision=num.float32,
713                                verbose=False, 
714                                **quant):
715        """
716        Write the static quantity info.
717
718        **quant is extra keyword arguments passed in. These must be
719          the numpy arrays to be stored in the sww file at each timestep.
720
721        The argument sww_precision allows for storing as either
722        * single precision (default): num.float32
723        * double precision: num.float64 or num.float
724
725        Precondition:
726            store_triangulation and
727            store_header have been called.
728        """
729
730        # The dictionary quant must contain numpy arrays for each name.
731        # These will typically be quantities from Domain such as friction
732        #
733        # Arrays not listed in static_quantitiues will be ignored, silently.
734        #
735        # This method will also write the ranges for each quantity,
736        # e.g. stage_range, xmomentum_range and ymomentum_range
737        for q in self.static_quantities:
738            if not quant.has_key(q):
739                msg = 'Values for quantity %s was not specified in ' % q
740                msg += 'store_quantities so they cannot be stored.'
741                raise NewQuantity, msg
742            else:
743                q_values = ensure_numeric(quant[q])
744               
745                x = q_values.astype(sww_precision)
746                outfile.variables[q][:] = x
747       
748                # This populates the _range values
749                outfile.variables[q + Write_sww.RANGE][0] = num.min(x)
750                outfile.variables[q + Write_sww.RANGE][1] = num.max(x)
751                   
752        # FIXME: Hack for backwards compatibility with old viewer
753        #if 'elevation' in self.static_quantities:
754        #    outfile.variables['z'][:] = outfile.variables['elevation'][:]
755
756                   
757                   
758       
759       
760    def store_quantities(self, 
761                         outfile, 
762                         sww_precision=num.float32,
763                         slice_index=None,
764                         time=None,
765                         verbose=False, 
766                         **quant):
767        """
768        Write the quantity info at each timestep.
769
770        **quant is extra keyword arguments passed in. These must be
771          the numpy arrays to be stored in the sww file at each timestep.
772
773        if the time array is already been built, use the slice_index
774        to specify the index.
775
776        Otherwise, use time to increase the time dimension
777
778        Maybe make this general, but the viewer assumes these quantities,
779        so maybe we don't want it general - unless the viewer is general
780       
781        The argument sww_precision allows for storing as either
782        * single precision (default): num.float32
783        * double precision: num.float64 or num.float
784
785        Precondition:
786            store_triangulation and
787            store_header have been called.
788        """
789
790        if time is not None:
791            file_time = outfile.variables['time']
792            slice_index = len(file_time)
793            file_time[slice_index] = time
794        else:
795            slice_index = int(slice_index) # Has to be cast in case it was numpy.int   
796
797        # Write the named dynamic quantities
798        # The dictionary quant must contain numpy arrays for each name.
799        # These will typically be the conserved quantities from Domain
800        # (Typically stage,  xmomentum, ymomentum).
801        #
802        # Arrays not listed in dynamic_quantitiues will be ignored, silently.
803        #
804        # This method will also write the ranges for each quantity,
805        # e.g. stage_range, xmomentum_range and ymomentum_range
806        for q in self.dynamic_quantities:
807            if not quant.has_key(q):
808                msg = 'Values for quantity %s was not specified in ' % q
809                msg += 'store_quantities so they cannot be stored.'
810                raise NewQuantity, msg
811            else:
812                q_values = ensure_numeric(quant[q])
813               
814                x = q_values.astype(sww_precision)
815                outfile.variables[q][slice_index] = x
816                   
817       
818                # This updates the _range values
819                q_range = outfile.variables[q + Write_sww.RANGE][:]
820                q_values_min = num.min(q_values)
821                if q_values_min < q_range[0]:
822                    outfile.variables[q + Write_sww.RANGE][0] = q_values_min
823                q_values_max = num.max(q_values)
824                if q_values_max > q_range[1]:
825                    outfile.variables[q + Write_sww.RANGE][1] = q_values_max
826
827    ##
828    # @brief Print the quantities in the underlying file.
829    # @param outfile UNUSED.
830    def verbose_quantities(self, outfile):
831        log.critical('------------------------------------------------')
832        log.critical('More Statistics:')
833        for q in self.dynamic_quantities:
834            log.critical('  %s in [%f, %f]'
835                         % (q, outfile.variables[q+Write_sww.RANGE][0],
836                            outfile.variables[q+Write_sww.RANGE][1]))
837        log.critical('------------------------------------------------')
838
839
840
841
842def extent_sww(file_name):
843    """Read in an sww file, then get its extents
844
845    Input:
846    file_name - the sww file
847
848    Output:
849    A list: [min(x),max(x),min(y),max(y),min(stage.flat),max(stage.flat)]
850    """
851
852    from Scientific.IO.NetCDF import NetCDFFile
853
854    #Get NetCDF
855    fid = NetCDFFile(file_name, netcdf_mode_r)
856
857    # Get the variables
858    x = fid.variables['x'][:]
859    y = fid.variables['y'][:]
860    stage = fid.variables['stage'][:]
861
862    fid.close()
863
864    return [min(x), max(x), min(y), max(y), num.min(stage), num.max(stage)]
865
866
867def load_sww_as_domain(filename, boundary=None, t=None,
868               fail_if_NaN=True, NaN_filler=0,
869               verbose=False, very_verbose=False):
870    """
871    Load an sww file into a domain.
872
873    Usage: domain = load_sww_as_domain('file.sww',
874                        t=time (default = last time in file))
875
876    Boundary is not recommended if domain.smooth is not selected, as it
877    uses unique coordinates, but not unique boundaries. This means that
878    the boundary file will not be compatable with the coordinates, and will
879    give a different final boundary, or crash.
880    """
881   
882    from Scientific.IO.NetCDF import NetCDFFile
883    from anuga.shallow_water.shallow_water_domain import Domain
884
885    # initialise NaN.
886    NaN = 9.969209968386869e+036
887
888    if verbose: log.critical('Reading from %s' % filename)
889
890    fid = NetCDFFile(filename, netcdf_mode_r)    # Open existing file for read
891    time = fid.variables['time']       # Timesteps
892    if t is None:
893        t = time[-1]
894    time_interp = get_time_interp(time,t)
895
896    # Get the variables as numeric arrays
897    x = fid.variables['x'][:]                   # x-coordinates of vertices
898    y = fid.variables['y'][:]                   # y-coordinates of vertices
899    elevation = fid.variables['elevation']      # Elevation
900    stage = fid.variables['stage']              # Water level
901    xmomentum = fid.variables['xmomentum']      # Momentum in the x-direction
902    ymomentum = fid.variables['ymomentum']      # Momentum in the y-direction
903
904    starttime = fid.starttime[0]
905    volumes = fid.variables['volumes'][:]       # Connectivity
906    coordinates = num.transpose(num.asarray([x.tolist(), y.tolist()]))
907    # FIXME (Ole): Something like this might be better:
908    #                 concatenate((x, y), axis=1)
909    # or              concatenate((x[:,num.newaxis], x[:,num.newaxis]), axis=1)
910
911    conserved_quantities = []
912    interpolated_quantities = {}
913    other_quantities = []
914
915    # get geo_reference
916    try:                             # sww files don't have to have a geo_ref
917        geo_reference = Geo_reference(NetCDFObject=fid)
918    except: # AttributeError, e:
919        geo_reference = None
920
921    if verbose: log.critical('    getting quantities')
922
923    for quantity in fid.variables.keys():
924        dimensions = fid.variables[quantity].dimensions
925        if 'number_of_timesteps' in dimensions:
926            conserved_quantities.append(quantity)
927            interpolated_quantities[quantity] = \
928                  interpolated_quantity(fid.variables[quantity][:], time_interp)
929        else:
930            other_quantities.append(quantity)
931
932    other_quantities.remove('x')
933    other_quantities.remove('y')
934    #other_quantities.remove('z')
935    other_quantities.remove('volumes')
936    try:
937        other_quantities.remove('stage_range')
938        other_quantities.remove('xmomentum_range')
939        other_quantities.remove('ymomentum_range')
940        other_quantities.remove('elevation_range')
941    except:
942        pass
943
944    conserved_quantities.remove('time')
945
946    if verbose: log.critical('    building domain')
947
948    #    From domain.Domain:
949    #    domain = Domain(coordinates, volumes,\
950    #                    conserved_quantities = conserved_quantities,\
951    #                    other_quantities = other_quantities,zone=zone,\
952    #                    xllcorner=xllcorner, yllcorner=yllcorner)
953
954    # From shallow_water.Domain:
955    coordinates = coordinates.tolist()
956    volumes = volumes.tolist()
957    # FIXME:should this be in mesh? (peter row)
958    if fid.smoothing == 'Yes':
959        unique = False
960    else:
961        unique = True
962    if unique:
963        coordinates, volumes, boundary = weed(coordinates, volumes,boundary)
964
965     
966   
967    try:
968        domain = Domain(coordinates, volumes, boundary)
969    except AssertionError, e:
970        fid.close()
971        msg = 'Domain could not be created: %s. ' \
972              'Perhaps use "fail_if_NaN=False and NaN_filler = ..."' % e
973        raise DataDomainError, msg
974
975    if not boundary is None:
976        domain.boundary = boundary
977
978    domain.geo_reference = geo_reference
979
980    domain.starttime = float(starttime) + float(t)
981    domain.time = 0.0
982
983    for quantity in other_quantities:
984        try:
985            NaN = fid.variables[quantity].missing_value
986        except:
987            pass                       # quantity has no missing_value number
988        X = fid.variables[quantity][:]
989        if very_verbose:
990            log.critical('       %s' % str(quantity))
991            log.critical('        NaN = %s' % str(NaN))
992            log.critical('        max(X)')
993            log.critical('       %s' % str(max(X)))
994            log.critical('        max(X)==NaN')
995            log.critical('       %s' % str(max(X)==NaN))
996            log.critical('')
997        if max(X) == NaN or min(X) == NaN:
998            if fail_if_NaN:
999                msg = 'quantity "%s" contains no_data entry' % quantity
1000                raise DataMissingValuesError, msg
1001            else:
1002                data = (X != NaN)
1003                X = (X*data) + (data==0)*NaN_filler
1004        if unique:
1005            X = num.resize(X, (len(X)/3, 3))
1006        domain.set_quantity(quantity, X)
1007    #
1008    for quantity in conserved_quantities:
1009        try:
1010            NaN = fid.variables[quantity].missing_value
1011        except:
1012            pass                       # quantity has no missing_value number
1013        X = interpolated_quantities[quantity]
1014        if very_verbose:
1015            log.critical('       %s' % str(quantity))
1016            log.critical('        NaN = %s' % str(NaN))
1017            log.critical('        max(X)')
1018            log.critical('       %s' % str(max(X)))
1019            log.critical('        max(X)==NaN')
1020            log.critical('       %s' % str(max(X)==NaN))
1021            log.critical('')
1022        if max(X) == NaN or min(X) == NaN:
1023            if fail_if_NaN:
1024                msg = 'quantity "%s" contains no_data entry' % quantity
1025                raise DataMissingValuesError, msg
1026            else:
1027                data = (X != NaN)
1028                X = (X*data) + (data==0)*NaN_filler
1029        if unique:
1030            X = num.resize(X, (X.shape[0]/3, 3))
1031        domain.set_quantity(quantity, X)
1032
1033    fid.close()
1034
1035    return domain
1036
1037
1038def get_mesh_and_quantities_from_file(filename,
1039                                      quantities=None,
1040                                      verbose=False):
1041    """Get and rebuild mesh structure and associated quantities from sww file
1042
1043    Input:
1044        filename - Name os sww file
1045        quantities - Names of quantities to load
1046
1047    Output:
1048        mesh - instance of class Interpolate
1049               (including mesh and interpolation functionality)
1050        quantities - arrays with quantity values at each mesh node
1051        time - vector of stored timesteps
1052
1053    This function is used by e.g.:
1054        get_interpolated_quantities_at_polyline_midpoints
1055    """
1056
1057    # FIXME (Ole): Maybe refactor filefunction using this more fundamental code.
1058
1059    import types
1060    from Scientific.IO.NetCDF import NetCDFFile
1061    from anuga.abstract_2d_finite_volumes.neighbour_mesh import Mesh
1062
1063    if verbose: log.critical('Reading from %s' % filename)
1064
1065    fid = NetCDFFile(filename, netcdf_mode_r)    # Open existing file for read
1066    time = fid.variables['time'][:]    # Time vector
1067    time += fid.starttime[0]
1068
1069    # Get the variables as numeric arrays
1070    x = fid.variables['x'][:]                   # x-coordinates of nodes
1071    y = fid.variables['y'][:]                   # y-coordinates of nodes
1072    elevation = fid.variables['elevation'][:]   # Elevation
1073    stage = fid.variables['stage'][:]           # Water level
1074    xmomentum = fid.variables['xmomentum'][:]   # Momentum in the x-direction
1075    ymomentum = fid.variables['ymomentum'][:]   # Momentum in the y-direction
1076
1077    # Mesh (nodes (Mx2), triangles (Nx3))
1078    nodes = num.concatenate((x[:,num.newaxis], y[:,num.newaxis]), axis=1)
1079    triangles = fid.variables['volumes'][:]
1080
1081    # Get geo_reference
1082    try:
1083        geo_reference = Geo_reference(NetCDFObject=fid)
1084    except: #AttributeError, e:
1085        # Sww files don't have to have a geo_ref
1086        geo_reference = None
1087
1088    if verbose: log.critical('    building mesh from sww file %s' % filename)
1089
1090    boundary = None
1091
1092    #FIXME (Peter Row): Should this be in mesh?
1093    if fid.smoothing != 'Yes':
1094        nodes = nodes.tolist()
1095        triangles = triangles.tolist()
1096        nodes, triangles, boundary = weed(nodes, triangles, boundary)
1097
1098    try:
1099        mesh = Mesh(nodes, triangles, boundary, geo_reference=geo_reference)
1100    except AssertionError, e:
1101        fid.close()
1102        msg = 'Domain could not be created: %s. "' % e
1103        raise DataDomainError, msg
1104
1105    quantities = {}
1106    quantities['elevation'] = elevation
1107    quantities['stage'] = stage
1108    quantities['xmomentum'] = xmomentum
1109    quantities['ymomentum'] = ymomentum
1110
1111    fid.close()
1112
1113    return mesh, quantities, time
1114
1115
1116def get_time_interp(time, t=None):
1117    """Finds the ratio and index for time interpolation.
1118        time is an array of time steps
1119        t is the sample time.
1120        returns a tuple containing index into time, and ratio
1121    """
1122    if t is None:
1123        t=time[-1]
1124        index = -1
1125        ratio = 0.
1126    else:
1127        T = time
1128        tau = t
1129        index=0
1130        msg = 'Time interval derived from file %s [%s:%s]' \
1131              % ('FIXMEfilename', T[0], T[-1])
1132        msg += ' does not match model time: %s' % tau
1133        if tau < time[0]: raise DataTimeError, msg
1134        if tau > time[-1]: raise DataTimeError, msg
1135        while tau > time[index]: index += 1
1136        while tau < time[index]: index -= 1
1137        if tau == time[index]:
1138            #Protect against case where tau == time[-1] (last time)
1139            # - also works in general when tau == time[i]
1140            ratio = 0
1141        else:
1142            #t is now between index and index+1
1143            ratio = (tau - time[index])/(time[index+1] - time[index])
1144
1145    return (index, ratio)
1146
1147
1148
1149##
1150# @brief Interpolate a quantity wrt time.
1151# @param saved_quantity The quantity to interpolate.
1152# @param time_interp (index, ratio)
1153# @return A vector of interpolated values.
1154def interpolated_quantity(saved_quantity, time_interp):
1155    '''Given an index and ratio, interpolate quantity with respect to time.'''
1156
1157    index, ratio = time_interp
1158
1159    Q = saved_quantity
1160
1161    if ratio > 0:
1162        q = (1-ratio)*Q[index] + ratio*Q[index+1]
1163    else:
1164        q = Q[index]
1165
1166    #Return vector of interpolated values
1167    return q
1168
1169
1170def weed(coordinates, volumes, boundary=None):
1171    """ Excise all duplicate points.
1172    """
1173    if isinstance(coordinates, num.ndarray):
1174        coordinates = coordinates.tolist()
1175    if isinstance(volumes, num.ndarray):
1176        volumes = volumes.tolist()
1177
1178    unique = False
1179    point_dict = {}
1180    same_point = {}
1181    for i in range(len(coordinates)):
1182        point = tuple(coordinates[i])
1183        if point_dict.has_key(point):
1184            unique = True
1185            same_point[i] = point
1186            #to change all point i references to point j
1187        else:
1188            point_dict[point] = i
1189            same_point[i] = point
1190
1191    coordinates = []
1192    i = 0
1193    for point in point_dict.keys():
1194        point = tuple(point)
1195        coordinates.append(list(point))
1196        point_dict[point] = i
1197        i += 1
1198
1199    for volume in volumes:
1200        for i in range(len(volume)):
1201            index = volume[i]
1202            if index > -1:
1203                volume[i] = point_dict[same_point[index]]
1204
1205    new_boundary = {}
1206    if not boundary is None:
1207        for segment in boundary.keys():
1208            point0 = point_dict[same_point[segment[0]]]
1209            point1 = point_dict[same_point[segment[1]]]
1210            label = boundary[segment]
1211            #FIXME should the bounday attributes be concaterated
1212            #('exterior, pond') or replaced ('pond')(peter row)
1213
1214            if new_boundary.has_key((point0, point1)):
1215                new_boundary[(point0,point1)] = new_boundary[(point0, point1)]
1216
1217            elif new_boundary.has_key((point1, point0)):
1218                new_boundary[(point1,point0)] = new_boundary[(point1, point0)]
1219            else: new_boundary[(point0, point1)] = label
1220
1221        boundary = new_boundary
1222
1223    return coordinates, volumes, boundary
1224
1225
1226
Note: See TracBrowser for help on using the repository browser.