source: storm_surge/pyvolution/data_manager.py @ 907

Last change on this file since 907 was 1, checked in by duncan, 20 years ago

initial import
File size: 23.5 KB
Line 
1"""Functions to store and retrieve data for the Shallow Water Wave equation.
2There are two kinds of data
3
4  1: Constant data: Vertex coordinates and field values. Stored once
5  2: Variable data: Conserved quantities. Stored once per timestep.
6
7All data is assumed to reside at vertex locations. 
8
9"""
10
11from Numeric import concatenate       
12
13
14def make_filename(s):
15    """Transform argument string into a suitable filename
16    """
17
18    s = s.strip()
19    s = s.replace(' ', '_')
20    s = s.replace('(', '')
21    s = s.replace(')', '')
22    s = s.replace('__', '_')               
23       
24    return s
25
26
27def check_dir(path, verbose=None):
28    """Check that specified path exists.
29    If path does not exist it will be created if possible   
30
31    USAGE:
32       checkdir(path, verbose):
33
34    ARGUMENTS:
35        path -- Directory
36        verbose -- Flag verbose output (default: None)
37
38    RETURN VALUE:
39        Verified path including trailing separator
40       
41    """
42
43    import os, sys
44    import os.path
45
46    if sys.platform in ['nt', 'dos', 'win32', 'what else?']:
47        unix = 0
48    else:
49        unix = 1
50
51
52    if path[-1] != os.sep: 
53        path = path + os.sep  # Add separator for directories
54
55    path = os.path.expanduser(path) # Expand ~ or ~user in pathname
56    if not (os.access(path,os.R_OK and os.W_OK) or path == ''):
57        try:
58            exitcode=os.mkdir(path)
59
60            # Change access rights if possible
61            #
62            if unix:
63                exitcode=os.system('chmod 775 '+path)
64            else:
65                pass  # FIXME: What about acces rights under Windows?
66           
67            if verbose: print 'MESSAGE: Directory', path, 'created.'
68         
69        except:
70            print 'WARNING: Directory', path, 'could not be created.'
71            if unix:
72                path = '/tmp/'
73            else:
74                path = 'C:'
75               
76            print 'Using directory %s instead' %path
77
78    return(path)
79
80   
81
82def del_dir(path):
83    """Recursively delete directory path and all its contents
84    """
85
86    import os
87
88    if os.path.isdir(path):
89        for file in os.listdir(path):
90            X = os.path.join(path, file)
91
92
93            if os.path.isdir(X) and not os.path.islink(X):
94                del_dir(X) 
95            else:
96                try:
97                    os.remove(X)
98                except:
99                    print "Could not remove file %s" %X
100
101        os.rmdir(path)
102
103
104       
105def create_filename(filename, format, size, time=None):
106
107    import os
108    from config import data_dir
109
110           
111    FN = check_dir(data_dir) + filename + '_size%d' %size
112       
113    if time is not None:
114        FN += '_time%.2f' %time
115       
116    FN += '.' + format
117       
118    return FN
119   
120
121def get_files(filename, format, size):
122    """Get all file (names) with gven name, size and format
123    """
124
125    import glob
126
127    import os
128    from config import data_dir
129
130           
131    dir = check_dir(data_dir)
132
133    pattern = dir + os.sep + filename + '_size=%d*.%s' %(size, format)
134    return glob.glob(pattern)
135
136
137
138#Generic class for storing output to e.g. visualisation or checkpointing
139class Data_format:
140    """Generic interface to data formats
141    """
142
143       
144    def __init__(self, domain, extension, mode = 'w'):
145        assert mode in ['r', 'w', 'a'], '''Mode %s must be either:''' %mode +\
146                                        '''   'w' (write)'''+\
147                                        '''   'r' (read)''' +\
148                                        '''   'a' (append)'''   
149
150        #Create filename
151        self.filename = create_filename(domain.get_name(), extension, len(domain))
152        self.timestep = 0
153        self.number_of_volumes = len(domain)
154        self.domain = domain
155   
156   
157    #FIXME: Should we have a general set_precision function?
158   
159
160
161#Class for storing output to e.g. visualisation
162class Data_format_sww(Data_format):
163    """Interface to native NetCDF format (.sww)
164    """
165
166       
167    def __init__(self, domain, mode = 'w'):
168        from Scientific.IO.NetCDF import NetCDFFile
169        from Numeric import Int, Float, Float32
170
171        self.precision = Float32 #Use single precision
172       
173        Data_format.__init__(self, domain, 'sww', mode)
174
175
176        # NetCDF file definition
177        fid = NetCDFFile(self.filename, mode)
178       
179        if mode == 'w':
180           
181            #Create new file
182            fid.institution = 'Geoscience Australia'
183            fid.description = 'Output from pyvolution suitable for plotting'
184           
185            if domain.smooth:
186                fid.smoothing = 'Yes'
187            else:   
188                fid.smoothing = 'No'
189
190            fid.order = domain.default_order   
191
192            # dimension definitions
193            fid.createDimension('number_of_volumes', self.number_of_volumes) 
194            fid.createDimension('number_of_vertices', 3)
195
196            if domain.smooth is True:
197                fid.createDimension('number_of_points', len(domain.vertexlist))
198            else:   
199                fid.createDimension('number_of_points', 3*self.number_of_volumes)
200               
201            fid.createDimension('number_of_timesteps', None) #extensible
202
203            # variable definitions
204            fid.createVariable('x', self.precision, ('number_of_points',))
205            fid.createVariable('y', self.precision, ('number_of_points',))
206            fid.createVariable('z', self.precision, ('number_of_points',))
207       
208            fid.createVariable('volumes', Int, ('number_of_volumes',
209                                                'number_of_vertices'))
210
211            fid.createVariable('time', self.precision,
212                               ('number_of_timesteps',))
213       
214            fid.createVariable('stage', self.precision,
215                               ('number_of_timesteps',
216                                'number_of_points'))
217
218
219        #Close
220        fid.close()               
221
222
223    def store_connectivity(self):
224        """Specialisation of store_connectivity for net CDF format
225
226        Writes x,y,z coordinates of triangles constituting
227        the bed elevation.
228        """
229
230        from Scientific.IO.NetCDF import NetCDFFile
231
232        from Numeric import concatenate
233
234        domain = self.domain
235
236        #Get NetCDF
237        fid = NetCDFFile(self.filename, 'a')  #Open existing file for append
238       
239        # Get the variables
240        x = fid.variables['x']
241        y = fid.variables['y']
242        z = fid.variables['z']
243       
244        volumes = fid.variables['volumes']
245
246        # Get X, Y and bed elevation Z
247        Q = domain.quantities['elevation']
248        X,Y,Z,V = Q.get_vertex_values(xy=True,
249                                      precision = self.precision)
250                                     
251
252
253        x[:] = X.astype(self.precision)
254        y[:] = Y.astype(self.precision)
255        z[:] = Z.astype(self.precision)
256
257        volumes[:] = V
258       
259        #Close
260        fid.close()
261
262
263    def store_timestep(self, name):
264        """Store time and named quantity to file
265        """
266        from Scientific.IO.NetCDF import NetCDFFile
267        from time import sleep
268               
269        #Get NetCDF
270        retries = 0
271        file_open = False
272        while not file_open and retries < 10:
273            try:       
274                fid = NetCDFFile(self.filename, 'a') #Open existing file
275            except IOError:
276                #This could happen if someone was reading the file.
277                #In that case, wait a while and try again
278                msg = 'Warning (store_timestep): File %s could not be opened'\
279                %self.filename
280                msg += ' - trying again'
281                print msg
282                retries += 1
283                sleep(1)
284            else:
285               file_open = True
286               
287        if not file_open:
288            msg = 'File %s could not be opened for append' %self.filename
289            raise msg
290               
291           
292        domain = self.domain
293       
294        # Get the variables
295        time = fid.variables['time']
296        stage = fid.variables['stage']       
297        i = len(time)
298
299        #Store stage
300        time[i] = self.domain.time
301
302        # Get quantity
303        Q = domain.quantities[name]
304        A,V = Q.get_vertex_values(xy=False,
305                                  precision = self.precision)
306       
307        stage[i,:] = A.astype(self.precision)
308
309        #Flush and close
310        fid.sync()
311        fid.close()
312
313
314#Class for handling checkpoints data
315class Data_format_cpt(Data_format):
316    """Interface to native NetCDF format (.cpt)
317    """
318
319       
320    def __init__(self, domain, mode = 'w'):
321        from Scientific.IO.NetCDF import NetCDFFile
322        from Numeric import Int, Float, Float
323
324        self.precision = Float #Use full precision
325       
326        Data_format.__init__(self, domain, 'sww', mode)
327
328
329        # NetCDF file definition
330        fid = NetCDFFile(self.filename, mode)
331       
332        if mode == 'w':
333            #Create new file
334            fid.institution = 'Geoscience Australia'
335            fid.description = 'Checkpoint data'
336            #fid.smooth = domain.smooth               
337            fid.order = domain.default_order   
338
339            # dimension definitions
340            fid.createDimension('number_of_volumes', self.number_of_volumes) 
341            fid.createDimension('number_of_vertices', 3)
342
343            #Store info at all vertices (no smoothing)
344            fid.createDimension('number_of_points', 3*self.number_of_volumes)
345            fid.createDimension('number_of_timesteps', None) #extensible
346
347            # variable definitions
348
349            #Mesh
350            fid.createVariable('x', self.precision, ('number_of_points',))
351            fid.createVariable('y', self.precision, ('number_of_points',))
352            #fid.createVariable('z', self.precision, ('number_of_points',))
353       
354            fid.createVariable('volumes', Int, ('number_of_volumes',
355                                                'number_of_vertices'))
356
357            fid.createVariable('time', self.precision,
358                               ('number_of_timesteps',))
359
360            #Allocate space for all quantities
361            for name in domain.quantities.keys():
362                fid.createVariable(name, self.precision,
363                                   ('number_of_timesteps',
364                                    'number_of_points'))
365
366        #Close
367        fid.close()               
368
369
370    def store_checkpoint(self):
371        """
372        Write x,y coordinates of triangles.
373        Write connectivity (
374        constituting
375        the bed elevation.
376        """
377
378        from Scientific.IO.NetCDF import NetCDFFile
379
380        from Numeric import concatenate
381
382        domain = self.domain
383
384        #Get NetCDF
385        fid = NetCDFFile(self.filename, 'a')  #Open existing file for append
386       
387        # Get the variables
388        x = fid.variables['x']
389        y = fid.variables['y']
390       
391        volumes = fid.variables['volumes']
392
393        # Get X, Y and bed elevation Z
394        Q = domain.quantities['elevation']
395        X,Y,Z,V = Q.get_vertex_values(xy=True,
396                                      precision = self.precision)
397                                     
398
399
400        x[:] = X.astype(self.precision)
401        y[:] = Y.astype(self.precision)
402        z[:] = Z.astype(self.precision)
403
404        volumes[:] = V
405       
406        #Close
407        fid.close()
408
409
410    def store_timestep(self, name):
411        """Store time and named quantity to file
412        """
413        from Scientific.IO.NetCDF import NetCDFFile
414        from time import sleep
415               
416        #Get NetCDF
417        retries = 0
418        file_open = False
419        while not file_open and retries < 10:
420            try:       
421                fid = NetCDFFile(self.filename, 'a') #Open existing file
422            except IOError:
423                #This could happen if someone was reading the file.
424                #In that case, wait a while and try again
425                msg = 'Warning (store_timestep): File %s could not be opened'\
426                %self.filename
427                msg += ' - trying again'
428                print msg
429                retries += 1
430                sleep(1)
431            else:
432               file_open = True
433               
434        if not file_open:
435            msg = 'File %s could not be opened for append' %self.filename
436            raise msg
437               
438           
439        domain = self.domain
440       
441        # Get the variables
442        time = fid.variables['time']
443        stage = fid.variables['stage']       
444        i = len(time)
445
446        #Store stage
447        time[i] = self.domain.time
448
449        # Get quantity
450        Q = domain.quantities[name]
451        A,V = Q.get_vertex_values(xy=False,
452                                  precision = self.precision)
453       
454        stage[i,:] = A.astype(self.precision)
455
456        #Flush and close
457        fid.sync()
458        fid.close()
459
460
461
462
463
464
465
466#Function for storing xya output
467#FIXME Not done yet for this version
468class Data_format_xya(Data_format):
469    """Generic interface to data formats
470    """
471
472
473    def __init__(self, domain, mode = 'w'):
474        from Scientific.IO.NetCDF import NetCDFFile
475        from Numeric import Int, Float, Float32
476
477        self.precision = Float32 #Use single precision
478       
479        Data_format.__init__(self, domain, 'xya', mode)
480       
481   
482   
483
484    def store_all(self):
485        """Specialisation of store all for xya format
486
487        Writes x,y,z coordinates of triangles constituting
488        the bed elevation.
489        """
490
491        from Numeric import concatenate
492
493        domain = self.domain
494       
495        fd = open(self.filename, 'w')
496
497
498        if domain.smooth is True:
499            number_of_points =  len(domain.vertexlist)
500        else:   
501            number_of_points = 3*self.number_of_volumes
502
503        numVertAttrib = 3 #Three attributes is what is assumed by the xya format
504
505        fd.write(str(number_of_points) + " " + str(numVertAttrib) +\
506                 " # <vertex #> <x> <y> [attributes]" + "\n")
507
508
509        # Get X, Y, bed elevation and friction (index=0,1)
510        X,Y,A,V = domain.get_vertex_values(xy=True, value_array='field_values',
511                                           indices = (0,1), precision = self.precision)
512
513        bed_eles = A[:,0]
514        fricts = A[:,1]
515
516        # Get stage (index=0)
517        B,V = domain.get_vertex_values(xy=False, value_array='conserved_quantities',
518                                       indices = (0,), precision = self.precision)
519
520        stages = B[:,0]
521
522        #<vertex #> <x> <y> [attributes]
523        for x, y, bed_ele, stage, frict in map(None, X, Y, bed_eles,
524                                               stages, fricts):
525
526            s = '%.6f %.6f %.6f %.6f %.6f\n' %(x, y, bed_ele, stage, frict)
527            fd.write(s)
528           
529        #close
530        fd.close()
531
532
533    def store_timestep(self, t, V0, V1, V2):
534        """Store time, water heights (and momentums) to file
535        """
536        pass
537   
538
539#Auxiliary
540def write_obj(filename,x,y,z):
541    """Store x,y,z vectors into filename (obj format)
542       Vectors are assumed to have dimension (M,3) where
543       M corresponds to the number elements.
544       triangles are assumed to be disconnected
545
546       The three numbers in each vector correspond to three vertices,
547
548       e.g. the x coordinate of vertex 1 of element i is in x[i,1]
549       
550    """
551    #print 'Writing obj to %s' % filename
552
553    import os.path
554
555    root, ext = os.path.splitext(filename)
556    if ext == '.obj':
557        FN = filename
558    else:
559        FN = filename + '.obj'
560       
561
562    outfile = open(FN, 'wb')
563    outfile.write("# Triangulation as an obj file\n")
564
565    M, N = x.shape
566    assert N==3  #Assuming three vertices per element
567
568    for i in range(M):
569        for j in range(N):
570            outfile.write("v %f %f %f\n" % (x[i,j],y[i,j],z[i,j]))
571
572    for i in range(M):
573        base = i*N
574        outfile.write("f %d %d %d\n" % (base+1,base+2,base+3))
575
576    outfile.close()
577
578
579
580#Conversion routines
581def sww2obj(basefilename, size):
582    """Convert netcdf based data output to obj
583    """
584    from Scientific.IO.NetCDF import NetCDFFile       
585
586    from Numeric import Float, zeros
587
588    #Get NetCDF
589    FN = create_filename(basefilename, 'sww', size)
590    print 'Reading from ', FN
591    fid = NetCDFFile(FN, 'r')  #Open existing file for read
592
593       
594    # Get the variables
595    x = fid.variables['x']
596    y = fid.variables['y']
597    z = fid.variables['z']       
598    time = fid.variables['time']
599    stage = fid.variables['stage']
600       
601    M = size  #Number of lines
602    xx = zeros((M,3), Float)
603    yy = zeros((M,3), Float)
604    zz = zeros((M,3), Float)
605
606    for i in range(M):
607        for j in range(3):
608            xx[i,j] = x[i+j*M]
609            yy[i,j] = y[i+j*M]
610            zz[i,j] = z[i+j*M]
611       
612    #Write obj for bathymetry
613    FN = create_filename(basefilename, 'obj', size)       
614    write_obj(FN,xx,yy,zz)
615
616
617    #Now read all the data with variable information, combine with
618    #x,y info and store as obj
619
620    for k in range(len(time)):
621        t = time[k]
622        print 'Processing timestep %f' %t
623
624        for i in range(M):
625            for j in range(3):
626                zz[i,j] = stage[k,i+j*M]
627
628
629        #Write obj for variable data
630        #FN = create_filename(basefilename, 'obj', size, time=t)
631        FN = create_filename(basefilename[:5], 'obj', size, time=t)   
632        write_obj(FN,xx,yy,zz)
633
634
635def dat2obj(basefilename):
636    """Convert line based data output to obj
637    """
638
639    import glob, os
640    from config import data_dir
641   
642   
643    #Get bathymetry and x,y's
644    lines = open(data_dir+os.sep+basefilename+'_geometry.dat', 'r').readlines()
645   
646    from Numeric import zeros, Float
647
648    M = len(lines)  #Number of lines
649    x = zeros((M,3), Float)
650    y = zeros((M,3), Float)
651    z = zeros((M,3), Float)
652
653    ##i = 0
654    for i, line in enumerate(lines):
655        tokens = line.split()
656        values = map(float,tokens)       
657
658        for j in range(3):
659            x[i,j] = values[j*3]
660            y[i,j] = values[j*3+1]
661            z[i,j] = values[j*3+2]       
662       
663        ##i += 1
664
665
666    #Write obj for bathymetry
667    write_obj(data_dir+os.sep+basefilename+'_geometry',x,y,z)
668
669
670    #Now read all the data files with variable information, combine with
671    #x,y info
672    #and store as obj
673
674    files = glob.glob(data_dir+os.sep+basefilename+'*.dat')
675
676    for filename in files:
677        print 'Processing %s' % filename
678
679        lines = open(data_dir+os.sep+filename,'r').readlines()
680        assert len(lines) == M
681        root, ext = os.path.splitext(filename)
682
683        #Get time from filename
684        i0 = filename.find('_time=')
685        if i0 == -1:
686            #Skip bathymetry file
687            continue
688       
689        i0 += 6  #Position where time starts           
690        i1 = filename.find('.dat')         
691
692        if i1 > i0:
693            t = float(filename[i0:i1])
694        else:
695            raise 'Hmmmm'
696               
697       
698       
699        ##i = 0
700        for i, line in enumerate(lines):
701            tokens = line.split()
702            values = map(float,tokens)       
703
704            for j in range(3):
705                z[i,j] = values[j]       
706       
707            ##i += 1
708
709        #Write obj for variable data
710        write_obj(data_dir+os.sep+basefilename+'_time=%.4f' %t,x,y,z)
711
712
713def filter_netcdf(filename1, filename2, first=0, last=None, step = 1):
714    """Read netcdf filename1, pick timesteps first:step:last and save to
715    nettcdf file filename2
716    """
717    from Scientific.IO.NetCDF import NetCDFFile       
718   
719    #Get NetCDF
720    infile = NetCDFFile(filename1, 'r')  #Open existing file for read
721    outfile = NetCDFFile(filename2, 'w')  #Open new file
722   
723
724    #Copy dimensions
725    for d in infile.dimensions:
726        outfile.createDimension(d, infile.dimensions[d])
727
728    for name in infile.variables:
729        var = infile.variables[name]
730        outfile.createVariable(name, var.typecode(), var.dimensions)
731       
732   
733    #Copy the static variables
734    for name in infile.variables:
735        if name == 'time' or name == 'stage':
736            pass
737        else:
738            #Copy
739            outfile.variables[name][:] = infile.variables[name][:]           
740
741    #Copy selected timesteps
742    time = infile.variables['time']
743    stage = infile.variables['stage']
744
745    newtime = outfile.variables['time']
746    newstage = outfile.variables['stage']   
747
748    if last is None:
749        last = len(time)
750
751    selection = range(first, last, step)
752    for i, j in enumerate(selection):
753        print 'Copying timestep %d of %d (%f)' %(j, last-first, time[j])
754        newtime[i] = time[j]
755        newstage[i,:] = stage[j,:]       
756       
757    #Close
758    infile.close()
759    outfile.close()
760
761
762#Get data objects
763def get_dataobject(domain, mode='w'):
764    """Return instance of class of given format using filename
765    """
766
767    cls = eval('Data_format_%s' %domain.format)
768    return cls(domain, mode)
769
770
771
772
773#OBSOLETE STUFF
774#Native checkpoint format.
775#Information needed to recreate a state is preserved
776#FIXME: Rethink and maybe use netcdf format
777def cpt_variable_writer(filename, t, v0, v1, v2):
778    """Store all conserved quantities to file
779    """
780
781    M, N = v0.shape
782       
783    FN = create_filename(filename, 'cpt', M, t)
784    #print 'Writing to %s' %FN
785   
786    fid = open(FN, 'w')
787    for i in range(M):
788        for j in range(N):       
789            fid.write('%.16e ' %v0[i,j])           
790        for j in range(N):
791            fid.write('%.16e ' %v1[i,j])                       
792        for j in range(N):
793            fid.write('%.16e ' %v2[i,j])                       
794           
795        fid.write('\n')
796    fid.close()
797
798
799def cpt_variable_reader(filename, t, v0, v1, v2):
800    """Store all conserved quantities to file
801    """
802
803    M, N = v0.shape
804       
805    FN = create_filename(filename, 'cpt', M, t)
806    #print 'Reading from %s' %FN
807   
808    fid = open(FN)
809
810
811    for i in range(M):
812        values = fid.readline().split() #Get one line
813
814        for j in range(N):
815            v0[i,j] = float(values[j])
816            v1[i,j] = float(values[3+j])
817            v2[i,j] = float(values[6+j])
818   
819    fid.close()
820
821def cpt_constant_writer(filename, X0, X1, X2, v0, v1, v2):
822    """Writes x,y,z,z,z coordinates of triangles constituting the bed
823    elevation.
824    Not in use pt
825    """
826
827    M, N = v0.shape
828
829    print X0
830    import sys; sys.exit() 
831    FN = create_filename(filename, 'cpt', M)
832    print 'Writing to %s' %FN
833   
834    fid = open(FN, 'w')
835    for i in range(M):
836        for j in range(2):       
837            fid.write('%.16e ' %X0[i,j])   #x, y
838        for j in range(N):                   
839            fid.write('%.16e ' %v0[i,j])       #z,z,z,
840           
841        for j in range(2):                   
842            fid.write('%.16e ' %X1[i,j])   #x, y
843        for j in range(N):                               
844            fid.write('%.16e ' %v1[i,j]) 
845       
846        for j in range(2):                               
847            fid.write('%.16e ' %X2[i,j])   #x, y
848        for j in range(N):                                           
849            fid.write('%.16e ' %v2[i,j]) 
850           
851        fid.write('\n')
852    fid.close()
853
854
855
856#Function for storing out to e.g. visualisation
857#FIXME: Do we want this?
858#FIXME: Not done yet for this version
859def dat_constant_writer(filename, X0, X1, X2, v0, v1, v2):
860    """Writes x,y,z coordinates of triangles constituting the bed elevation.
861    """
862
863    M, N = v0.shape
864       
865    FN = create_filename(filename, 'dat', M)
866    #print 'Writing to %s' %FN
867   
868    fid = open(FN, 'w')
869    for i in range(M):
870        for j in range(2):       
871            fid.write('%f ' %X0[i,j])   #x, y
872        fid.write('%f ' %v0[i,0])       #z
873           
874        for j in range(2):                   
875            fid.write('%f ' %X1[i,j])   #x, y
876        fid.write('%f ' %v1[i,0])       #z
877       
878        for j in range(2):                               
879            fid.write('%f ' %X2[i,j])   #x, y
880        fid.write('%f ' %v2[i,0])       #z
881           
882        fid.write('\n')
883    fid.close()
884
885
886
887def dat_variable_writer(filename, t, v0, v1, v2):
888    """Store water height to file
889    """
890
891    M, N = v0.shape
892       
893    FN = create_filename(filename, 'dat', M, t)
894    #print 'Writing to %s' %FN
895   
896    fid = open(FN, 'w')
897    for i in range(M):
898        fid.write('%.4f ' %v0[i,0])
899        fid.write('%.4f ' %v1[i,0])
900        fid.write('%.4f ' %v2[i,0])       
901
902        fid.write('\n')
903    fid.close()
904
Note: See TracBrowser for help on using the repository browser.