source: inundation/ga/storm_surge/pyvolution/data_manager.py @ 333

Last change on this file since 333 was 288, checked in by ole, 20 years ago

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