MITgcmutils/MITgcmutils/mnc.py

ae10cb9ac4 Oliv*0001 import sys
                0002 import glob
                0003 import numpy as np
34625de9eb Oliv*0004 from .netcdf import netcdf_file
ae10cb9ac4 Oliv*0005
                0006 _exclude_global = ['close',
                0007                    'createDimension',
                0008                    'createVariable',
                0009                    'dimensions',
                0010                    'filename',
                0011                    'flush',
                0012                    'fp',
                0013                    'mode',
                0014                    'sync',
                0015                    'use_mmap',
                0016                    'variables',
                0017                    'version_byte',
                0018                   ]
                0019
                0020 _exclude_var = ['assignValue',
                0021                 'data',
                0022                 'dimensions',
                0023                 'getValue',
                0024                 'isrec',
                0025                 'itemsize',
                0026                 'shape',
                0027                 'typecode',
                0028                ]
                0029
                0030 def getattributes(nc, exclude=[]):
                0031     # in order not to rely on implementation, provide fallback
                0032     try:
                0033         a = dict(nc._attributes)
                0034     except AttributeError:
                0035         a = dict((k, getattr(nc, k)) for k in dir(nc) if k[0] != '_' and k not in exclude)
                0036     return a
                0037
                0038
9d41c6fbc0 Oliv*0039 class MNC:
ae10cb9ac4 Oliv*0040     """
                0041     A file object for MNC (tiled NetCDF) data.
                0042
                0043     Should behave mostly like scipy.io.netcdf.netcdf_file in 'r' mode.
                0044
                0045     Parameters
                0046     ----------
7621b5d564 Oliv*0047     fpatt : string
                0048         glob pattern for tile files
                0049     layout : string
                0050         which global layout to use:
                0051
                0052         'model'
                0053             use layout implied by Nx, Ny
                0054         'exch2'
                0055             use exch2 global layout
                0056         'faces'
                0057             variables are lists of exch2 faces
                0058
                0059         default is to use exch2 layout if present, model otherwise
                0060
                0061     Example
                0062     -------
                0063     >>> nc = mnc_files('mnc_*/state.0000000000.t*.nc')
                0064     >>> temp = nc.variables['Temp'][:]
                0065     >>> salt = nv.variables['S'][:]
                0066     >>> nc.close()
ae10cb9ac4 Oliv*0067     temp and salt are now assembled (global) arrays of shape (Nt, Nr, Ny, Nx)
                0068     where Nt is the number iterations found in the file (in this case probably 1).
7621b5d564 Oliv*0069
                0070     Notes
                0071     -----
                0072     The multitime option is not implemented, i.e., MNC cannot read files split
                0073     in time.
ae10cb9ac4 Oliv*0074     """
                0075
                0076     # avoid problems with __del__
                0077     nc = []
                0078
                0079     def __init__(self, fpatt, layout=None, multitime=False):
                0080         fnames = glob.glob(fpatt)
                0081 #        if multitime:
                0082 #            iters = [ f[-18:-8] for f in fnames if f.endswith('.t001.nc') ]
                0083 #            iters.sort()
                0084 #            fnames_first = [ f for f in fnames if f[-18:-8] == iters[0] ]
                0085 #        else:
                0086 #            fnames_first = fnames
                0087         fnames.sort()
                0088
                0089         # open files
                0090         self.nc = [ netcdf_file(f,'r') for f in fnames ]
                0091
                0092         # global attributes
                0093         # get from first file, but remove/reset tile-specific ones
                0094         self._attributes = getattributes(self.nc[0], _exclude_global)
                0095         self._attributes['tile_number'] = 1
                0096         self._attributes['bi'] = 1
                0097         self._attributes['bj'] = 1
                0098         haveexch2 = False
27d195aed6 Oliv*0099         for k in list(self._attributes):
ae10cb9ac4 Oliv*0100             if k.startswith('exch2_'):
                0101                 del self._attributes[k]
                0102                 haveexch2 = True
                0103
                0104         sNx = self.sNx
                0105         sNy = self.sNy
                0106         ntx = self.nSx*self.nPx
                0107         nty = self.nSy*self.nPy
                0108
                0109         if layout is None:
                0110             if haveexch2:
                0111                 layout = 'exch2'
                0112             else:
                0113                 layout = 'model'
                0114
                0115         self.layout = layout
                0116
                0117         # precompute indices
                0118         self._i0 = []
                0119         self._ie = []
                0120         self._j0 = []
                0121         self._je = []
9d41c6fbc0 Oliv*0122         self._fn = []
                0123         self._nf = 0
ae10cb9ac4 Oliv*0124         if layout == 'model':
                0125             self._nx = self.Nx
                0126             self._ny = self.Ny
                0127             for nc in self.nc:
                0128                 tn = nc.tile_number
                0129                 bj,bi = divmod(tn-1, ntx)
                0130                 ie = sNx*(bi+1-ntx)
                0131                 je = sNy*(bj+1-nty)
                0132                 self._i0.append(sNx*bi)
                0133                 self._j0.append(sNy*bj)
                0134                 self._ie.append(ie or None)
                0135                 self._je.append(je or None)
                0136         elif layout == 'exch2':
                0137             self._nx = 0
                0138             self._ny = 0
                0139             for nc in self.nc:
                0140                 i0 = nc.exch2_txGlobalo - 1
                0141                 j0 = nc.exch2_tyGlobalo - 1
                0142                 ie = i0 + sNx
                0143                 je = j0 + sNy
                0144                 self._i0.append(i0)
                0145                 self._j0.append(j0)
                0146                 self._ie.append(ie)
                0147                 self._je.append(je)
                0148                 self._nx = max(self._nx, ie)
                0149                 self._ny = max(self._ny, je)
                0150             # make ie, je relative to end (for non-tracer points)
                0151             for i in range(len(self._i0)):
                0152                 ie = self._ie[i] - self._nx
                0153                 je = self._je[i] - self._ny
                0154                 self._ie[i] = ie or None
                0155                 self._je[i] = je or None
                0156         elif layout == 'faces':
                0157             self._nx = {}
                0158             self._ny = {}
                0159             for nc in self.nc:
                0160                 fn = nc.exch2_myFace
                0161                 i0 = nc.exch2_tBasex
                0162                 j0 = nc.exch2_tBasey
                0163                 ie = i0 + sNx
                0164                 je = j0 + sNy
                0165                 self._fn.append(fn)
                0166                 self._i0.append(i0)
                0167                 self._j0.append(j0)
                0168                 self._ie.append(ie)
                0169                 self._je.append(je)
                0170                 self._nx[fn] = max(self._nx.get(fn, 0), ie)
                0171                 self._ny[fn] = max(self._ny.get(fn, 0), je)
                0172             # make ie, je relative to end (for non-tracer points)
                0173             for i in range(len(self._fn)):
                0174                 fn = self._fn[i]
                0175                 ie = self._ie[i] - self._nx[fn]
                0176                 je = self._je[i] - self._ny[fn]
                0177                 self._ie[i] = ie or None
                0178                 self._je[i] = je or None
                0179             self._fns = sorted(self._nx.keys())
                0180             self._nf = len(self._fns)
                0181             for i in range(len(self._fn)):
                0182                 self._fn[i] = self._fns.index(self._fn[i])
                0183             self._nx = np.array([self._nx[fn] for fn in self._fns])
                0184             self._ny = np.array([self._ny[fn] for fn in self._fns])
                0185         else:
                0186             raise ValueError('Unknown layout: {}'.format(layout))
                0187
                0188         # dimensions
                0189         self.dimensions = {}
                0190         for k,n in self.nc[0].dimensions.items():
                0191             # compute size of dimension in global array for X* and Y*
                0192             if k[0] == 'X':
                0193                 n += self._nx - sNx
                0194             if k[0] == 'Y':
                0195                 n += self._ny - sNy
                0196             self.dimensions[k] = n
                0197
                0198         # variables
                0199         var0 = self.nc[0].variables
                0200         # find size of record dimension first
                0201         if 'T' in self.dimensions and self.dimensions['T'] is None:
                0202             self.times = list(var0.get('T', [])[:])
                0203             self.iters = list(var0.get('iter', self.times)[:])
                0204             self.nrec = len(self.iters)
                0205
                0206         self.variables = dict((k, MNCVariable(self, k)) for k in var0)
                0207
                0208     def __getattr__(self, k):
9340c6074d Oliv*0209         try:
                0210             return self._attributes[k]
                0211         except KeyError:
d75cbe6b61 Oliv*0212             raise AttributeError("'MNC' object has no attribute '" + k + "'")
ae10cb9ac4 Oliv*0213
                0214     def __dir__(self):
                0215         return self.__dict__.keys() + self._attributes.keys()
                0216
                0217     def close(self):
                0218         """Close tile files"""
                0219         for nc in self.nc:
                0220             nc.close()
                0221
                0222     __del__ = close
                0223
                0224     @property
                0225     def faces(self):
                0226         if self.layout == 'faces':
                0227             return self._fns
                0228         else:
                0229             return None
                0230
                0231
                0232 def calcstrides(slices, dims):
                0233     try:
                0234         slices[0]
                0235     except TypeError:
                0236         slices = (slices,)
                0237
                0238     if Ellipsis in slices:
                0239         cut = slices.index(Ellipsis)
                0240         slices = slices[:cut] + (len(dims)-len(slices)+1)*(slice(0,None,None),) + slices[cut+1:]
                0241     else:
                0242         slices = slices + (len(dims)-len(slices))*(slice(0,None,None),)
                0243
                0244 #    return tuple( hasattr(s,'indices') and s.indices(dim) or s for s,dim in zip(slices,dims) )
                0245     strides = []
                0246     shape = []
                0247     fullshape = []
                0248     for s,dim in zip(slices,dims):
                0249         try:
                0250             stride = s.indices(dim)
                0251         except AttributeError:
                0252             stride = (s, s+1, 1)
                0253             n = 1
                0254         else:
                0255             # real slice, will make a dimension
                0256             start,stop,step = stride
                0257             n = (stop-start+step-1)//step
                0258             shape.append(n)
                0259
                0260         fullshape.append(n)
                0261         strides.append(stride)
                0262
                0263     return tuple(strides), tuple(shape), tuple(fullshape)
                0264
                0265
d9a448a10c Oliv*0266 class MNCVariable(object):
ae10cb9ac4 Oliv*0267     def __init__(self, mnc, name):
                0268         self._name = name
9d41c6fbc0 Oliv*0269         self.nc = mnc.nc
                0270         self.layout = mnc.layout
                0271         self._i0 = mnc._i0
                0272         self._ie = mnc._ie
                0273         self._j0 = mnc._j0
                0274         self._je = mnc._je
                0275         self._nf = mnc._nf
                0276         self._fn = mnc._fn
ae10cb9ac4 Oliv*0277         v0 = mnc.nc[0].variables[name]
                0278         self._attributes = getattributes(v0, _exclude_var)
                0279         self.itemsize = v0.data.itemsize
                0280         self.typecode = v0.typecode
                0281         self.dtype = np.dtype(self.typecode())
                0282         self.dimensions = v0.dimensions
                0283         self.shape = tuple( mnc.dimensions[d] for d in self.dimensions )
                0284         self.isrec = self.shape[0] is None
                0285         if self.isrec:
                0286             self.shape = (mnc.nrec,) + self.shape[1:]
                0287
                0288         # which dimensions are tiled
                0289         self._Xdim = None
                0290         self._Ydim = None
                0291         for i,d in enumerate(self.dimensions):
                0292             if d[0] == 'X': self._Xdim = i
                0293             if d[0] == 'Y': self._Ydim = i
                0294
                0295     def __getattr__(self, k):
956cb4f62a Oliv*0296         try:
                0297             return self._attributes[k]
                0298         except KeyError:
d75cbe6b61 Oliv*0299             raise AttributeError("'MNCVariable' object has no attribute '" + k + "'")
ae10cb9ac4 Oliv*0300
                0301     def __dir__(self):
                0302         return self.__dict__.keys() + self._attributes.keys()
                0303
                0304     def __getitem__(self, ind):
9d41c6fbc0 Oliv*0305         if self.layout == 'faces':
ae10cb9ac4 Oliv*0306             return self._getfaces(ind)
                0307
                0308         if ind in [Ellipsis, slice(None)]:
                0309             # whole array
                0310             res = np.zeros(self.shape, self.typecode())
                0311             s = [slice(None) for d in self.shape]
9d41c6fbc0 Oliv*0312             for i,nc in enumerate(self.nc):
ae10cb9ac4 Oliv*0313                 if self._Xdim is not None:
9d41c6fbc0 Oliv*0314                     s[self._Xdim] = slice(self._i0[i], self._ie[i])
ae10cb9ac4 Oliv*0315                 if self._Ydim is not None:
9d41c6fbc0 Oliv*0316                     s[self._Ydim] = slice(self._j0[i], self._je[i])
7621b5d564 Oliv*0317                 res[tuple(s)] = nc.variables[self._name][:]
ae10cb9ac4 Oliv*0318
                0319             return res
                0320         else:
                0321             # read only required data
                0322             strides,resshape,fullshape = calcstrides(ind, self.shape)
                0323             res = np.zeros(fullshape, self.dtype)
                0324             s = [slice(*stride) for stride in strides]
                0325             sres = [slice(None) for d in fullshape]
                0326             if self._Xdim is not None: I0,Ie,Is = strides[self._Xdim]
                0327             if self._Ydim is not None: J0,Je,Js = strides[self._Ydim]
9d41c6fbc0 Oliv*0328             for i,nc in enumerate(self.nc):
ae10cb9ac4 Oliv*0329                 if self._Xdim is not None:
9d41c6fbc0 Oliv*0330                     i0 = self._i0[i]
                0331                     ie = self.shape[self._Xdim] + (self._ie[i] or 0)
ae10cb9ac4 Oliv*0332                     a,b = divmod(I0 - i0, Is)
                0333                     e = np.clip(ie, I0, Ie)
                0334                     sres[self._Xdim] = slice(max(-a, 0), (e - I0)//Is)
                0335                     s[self._Xdim] = slice(max(I0 - i0, b), max(Ie - i0, 0), Is)
                0336                 if self._Ydim is not None:
9d41c6fbc0 Oliv*0337                     j0 = self._j0[i]
                0338                     je = self.shape[self._Ydim] + (self._je[i] or 0)
ae10cb9ac4 Oliv*0339                     a,b = divmod(J0 - j0, Js)
                0340                     e = np.clip(je, J0, Je)
                0341                     sres[self._Ydim] = slice(max(-a, 0), (e - J0)//Js)
                0342                     s[self._Ydim] = slice(max(J0 - j0, b), max(Je - j0, 0), Js)
7621b5d564 Oliv*0343                 res[tuple(sres)] = nc.variables[self._name][tuple(s)]
ae10cb9ac4 Oliv*0344
                0345             return res.reshape(resshape)
                0346
                0347     def _getfaces(self, ind):
                0348         res = []
9d41c6fbc0 Oliv*0349         for f in range(self._nf):
ae10cb9ac4 Oliv*0350             shape = tuple(np.isscalar(d) and d or d[f] for d in self.shape)
                0351             a = np.zeros(shape, self.typecode())
                0352             res.append(a)
                0353         s = [slice(None) for d in self.shape]
9d41c6fbc0 Oliv*0354         for i,nc in enumerate(self.nc):
                0355             fn = self._fn[i]
ae10cb9ac4 Oliv*0356             if self._Xdim is not None:
9d41c6fbc0 Oliv*0357                 s[self._Xdim] = slice(self._i0[i], self._ie[i])
ae10cb9ac4 Oliv*0358             if self._Ydim is not None:
9d41c6fbc0 Oliv*0359                 s[self._Ydim] = slice(self._j0[i], self._je[i])
7621b5d564 Oliv*0360             res[fn][tuple(s)] = nc.variables[self._name][:]
9d41c6fbc0 Oliv*0361         for f in range(self._nf):
ae10cb9ac4 Oliv*0362             res[f] = res[f][ind]
                0363
                0364         return res
                0365
                0366     def face(self, fn):
                0367         shape = tuple(np.isscalar(d) and d or d[fn] for d in self.shape)
                0368         res = np.zeros(shape, self.typecode())
                0369         s = [slice(None) for d in self.shape]
9d41c6fbc0 Oliv*0370         for i,nc in enumerate(self.nc):
                0371             if self._fn[i] == fn:
ae10cb9ac4 Oliv*0372                 if self._Xdim is not None:
9d41c6fbc0 Oliv*0373                     s[self._Xdim] = slice(self._i0[i], self._ie[i])
ae10cb9ac4 Oliv*0374                 if self._Ydim is not None:
9d41c6fbc0 Oliv*0375                     s[self._Ydim] = slice(self._j0[i], self._je[i])
7621b5d564 Oliv*0376                 res[tuple(s)] = nc.variables[self._name][:]
ae10cb9ac4 Oliv*0377
                0378         return res
                0379
                0380
                0381 def mnc_files(fpatt, layout=None):
                0382     return MNC(fpatt, layout)
                0383
                0384 mnc_files.__doc__ = MNC.__doc__
                0385
                0386
                0387 def rdmnc(fpatt, varnames=None, iters=None, slices=Ellipsis, layout=None):
7621b5d564 Oliv*0388     '''
                0389     Read one or more variables from an mnc file set.
                0390
ae10cb9ac4 Oliv*0391     Parameters
                0392     ----------
7621b5d564 Oliv*0393     fpatt : string
                0394         glob pattern for netcdf files comprising the set
                0395     varnames : list of strings, optional
                0396         list of variables to read (default all)
                0397     iters : list of int, optional
                0398         list of iterations (not time) to read
                0399     slices : tuple of slice objects
                0400         tuple of slices to read from each variable
                0401         (typically given as numpy.s_[...])
                0402
                0403     Returns
                0404     -------
                0405     dict of numpy arrays
                0406         dictionary of variable arrays
                0407
                0408     Example
                0409     -------
                0410     >>> S = rdmnc("mnc_*/state.0000000000.*', ['U', 'V'], slices=numpy.s_[..., 10:-10, 10:-10])
                0411     >>> u = S['U']
                0412     >>> v = S['V']
                0413
                0414     Notes
                0415     -----
ae10cb9ac4 Oliv*0416     Can currently read only one file set (i.e., 1 file per tile),
                0417     not several files split in time.
                0418
                0419     Consider using mnc_files for more control (and similar convenience).
                0420     The same restriction about multiple files applies, however.
                0421     '''
                0422     mnc = MNC(fpatt, layout)
                0423     if varnames is None:
                0424         varnames = mnc.variables.keys()
                0425     elif isinstance(varnames, str):
                0426         varnames = [varnames]
                0427
                0428     if iters is not None:
                0429         try:
                0430             iters[0]
                0431         except TypeError:
                0432             iters = [iters]
                0433
                0434         iits = [ mnc.iters.index(it) for it in iters ]
                0435
                0436         if not isinstance(slices, tuple):
                0437             slices = (slices,)
                0438
                0439     res = {}
                0440     for varname in varnames:
                0441         var = mnc.variables[varname]
                0442         if iters is not None and var.dimensions[0] == 'T':
                0443             res[varname] = np.array([var[(iit,)+slices] for iit in iits])
                0444         else:
                0445             res[varname] = var[slices]
                0446
                0447     mnc.close()
                0448     return res
                0449