drifters/drifters_combine

   1 #!/usr/bin/env python
   2 #***********************************************************************
   3 #*                   GNU Lesser General Public License
   4 #*
   5 #* This file is part of the GFDL Flexible Modeling System (FMS).
   6 #*
   7 #* FMS is free software: you can redistribute it and/or modify it under
   8 #* the terms of the GNU Lesser General Public License as published by
   9 #* the Free Software Foundation, either version 3 of the License, or (at
  10 #* your option) any later version.
  11 #*
  12 #* FMS is distributed in the hope that it will be useful, but WITHOUT
  13 #* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  14 #* FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  15 #* for more details.
  16 #*
  17 #* You should have received a copy of the GNU Lesser General Public
  18 #* License along with the FV3 dynamical core.
  19 #* If not, see <http://www.gnu.org/licenses/>.
  20 #***********************************************************************
  21
  22 VERSION = "_FILE_VERSION"
  23
  24 from Scientific.IO.NetCDF import NetCDFFile
  25 import sys
  26 import re
  27 import Numeric as Num
  28 import glob
  29 from optparse import OptionParser
  30
  31
  32 class drifters_combine:
  33
  34     def __init__(self, filenames):
  35
  36         self.data = {}
  37         self.nd = None
  38         self.nf = None
  39
  40         self.global_atts = {}
  41         self.position_atts = {}
  42         self.field_atts = {}
  43
  44         for f in filenames:
  45             nc =  NetCDFFile(f, 'r')
  46             index_time= nc.variables['index_time'][:]
  47             time      = nc.variables['time'][:]
  48             ids       = nc.variables['ids'][:]
  49
  50             positions = nc.variables['positions'][:]
  51             nd = Num.shape(positions)[1]
  52
  53             # get global attributes
  54             for a in nc.__dict__:
  55                 self.global_atts[a] = nc.__dict__[a]
  56
  57             # get attributes
  58             for a in nc.variables['positions'].__dict__:
  59                 if not self.position_atts.has_key(a):
  60                     self.position_atts[a] = getattr(nc.variables['positions'], a)
  61
  62             fields    = None
  63             nf        = 0
  64             if 'fields' in nc.variables:
  65                 fields    = nc.variables['fields'][:]
  66                 nf = Num.shape(fields)[1]
  67                 for a in nc.variables['fields'].__dict__:
  68                     if not self.field_atts.has_key(a):
  69                         self.field_atts[a] = getattr(nc.variables['fields'], a)
  70
  71             if self.nf==None: self.nf = nf
  72             if self.nd==None: self.nd = nd
  73
  74             if nf != self.nf or nd != self.nd:
  75                 raise 'Incompatible number of fields (nf) or space dimensions (nd) in file'% f
  76
  77             for i in range(len(ids)):
  78                 id = ids[i]
  79                 it = index_time[i]
  80                 tim= time[i]
  81                 xyz= positions[i,:]
  82                 fld = ()
  83                 if nf>0: fld= fields[i,:]
  84
  85                 if not self.data.has_key(it):
  86                     self.data[it] = {}
  87
  88                 if not self.data[it].has_key(id):
  89                     self.data[it][id] = (tim,) + tuple(xyz) + tuple(fld)
  90
  91     def save(self, outfile):
  92
  93         nc = NetCDFFile(outfile,'w')
  94
  95         nc.createDimension('it_id', None)
  96         nc.createDimension('nd', self.nd)
  97         if self.nf>0:
  98             nc.createDimension('nf', self.nf)
  99
 100         nc_index_time = nc.createVariable('index_time', Num.Int, ('it_id',))
 101         nc_time       = nc.createVariable('time', Num.Float64, ('it_id',))
 102         nc_ids        = nc.createVariable('ids', Num.Int, ('it_id',))
 103         nc_positions  = nc.createVariable('positions', Num.Float64, \
 104                                           ('it_id', 'nd'))
 105
 106         for a in self.global_atts:
 107             setattr(nc, a, self.global_atts[a])
 108
 109         for a in self.position_atts:
 110             setattr(nc_positions, a, self.position_atts[a])
 111
 112         if self.nf>0:
 113             nc_fields     = nc.createVariable('fields', Num.Float64, \
 114                                           ('it_id', 'nf'))
 115             for a in self.field_atts:
 116                 setattr(nc_fields, a, self.field_atts[a])
 117
 118         k = 0
 119         for it in range(len(self.data)):
 120             print 'it=', it
 121             for id in self.data[it]:
 122                 print '\tid=', id
 123                 nc_index_time[k] = it
 124                 nc_time[k]       = self.data[it][id][0]
 125                 nc_ids[k]        = id
 126                 nc_positions[k:k+1,:]= self.data[it][id][1:1+self.nd]
 127                 if self.nf>0:
 128                     nc_fields[k:k+1,:]= self.data[it][id][1+self.nd:1+self.nd+self.nf]
 129                 k += 1
 130
 131         nc.close()
 132
 133
 134 ###############################################################################
 135
 136
 137
 138 def main():
 139
 140     parser = OptionParser(version=VERSION)
 141     parser.add_option('-p', '--pe_range', action='store', type="string",
 142                   dest="pe_range",
 143                   help='PE range (e.g. "10...20"). By default the output of all PEs will be combined.',
 144                   default='',
 145                   )
 146     parser.add_option('-f', '--file', action='store', type="string",
 147                   dest="filename",
 148                   help='Input file. Files FILE.XXXX to FILE.YYYY will be combined.',
 149                   default='drifters_out.nc',
 150                   )
 151     parser.add_option('-o', '--output', action='store', type="string",
 152                   dest="output",
 153                   help='The combined files will be saved as OUTPUT. By default OUTPUT=FILE.',
 154                   default='',
 155                   )
 156
 157     options, args = parser.parse_args(sys.argv)
 158     fname = options.filename
 159     if not fname:
 160         print 'ERROR: must supply a netcdf file name ([-f|--file] FILE).'
 161         sys.exit(1)
 162
 163     min_pe = 0
 164     if options.pe_range:
 165         p1 = re.search(r'^\s*(\d+)\s*\.', options.pe_range)
 166         p2 = re.search(r'\.\s*(\d+)\s*$', options.pe_range)
 167         if p1 and p2:
 168             min_pe = p1.group(1)
 169             max_pe = p2.group(1)
 170         else:
 171             print 'ERROR: pe range should be in format [-p|--pe_range] "min_pe..max_pe".'
 172             sys.exit(1)
 173
 174         fnames = [fname + '.' + str(pe).zfill(4) for pe in range(int(min_pe), int(max_pe)+1)]
 175
 176     else:
 177         fnames = glob.glob(fname + '.' + '????')
 178
 179     vfc = drifters_combine(fnames)
 180     if options.output:
 181         ofile = options.output
 182     else:
 183         ofile = fname
 184
 185     vfc.save(ofile)
 186
 187 if __name__=='__main__': main()