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