MITgcmutils/MITgcmutils/diagnostics.py

ae10cb9ac4 Oliv*0001 import re
                0002 import numpy as np
                0003
                0004 nstats = 5
                0005
                0006 def readstats(fname):
7621b5d564 Oliv*0007     '''
8fc3047e23 Oliv*0008     statsPerLayer,statsVertInt,itrs = readstats(fname)
7621b5d564 Oliv*0009
ae10cb9ac4 Oliv*0010     Read a diagstats text file into record arrays (or dictionaries).
                0011
7621b5d564 Oliv*0012     Parameters
                0013     ----------
                0014     fname : string
                0015         name of diagstats file to read
                0016
ae10cb9ac4 Oliv*0017     Returns
                0018     -------
8fc3047e23 Oliv*0019     statsPerLayer : record array or dict of arrays
                0020         statistics per layer, shape (len(itrs), len(nReg), Nr, 5)
                0021     statsVertInt  : record array or dict of arrays
                0022         column integrals, shape (len(itrs), len(nReg), 5)
7621b5d564 Oliv*0023     itrs : list of int
                0024         iteration numbers found in the file
                0025
                0026     Notes
                0027     -----
8fc3047e23 Oliv*0028     - The 5 columns of the resulting arrays are
                0029       average, std.dev, min, max and total volume.
7621b5d564 Oliv*0030     - There is a record (or dictionary key) for each field found in the file.
8fc3047e23 Oliv*0031     - Regional axis is omitted if nReg == 1
7621b5d564 Oliv*0032
ae10cb9ac4 Oliv*0033     '''
                0034     flds = []
8fc3047e23 Oliv*0035     regs = [0]
ae10cb9ac4 Oliv*0036     with open(fname) as f:
                0037         for line in f:
                0038             if line.startswith('# end of header'):
                0039                 break
                0040
8fc3047e23 Oliv*0041             m = re.match(r'^# ([^: ]*) *: *(.*)$', line.rstrip())
ae10cb9ac4 Oliv*0042             if m:
                0043                 var,val = m.groups()
a7e75041d1 Oliv*0044                 if var.startswith('Fields'):
8fc3047e23 Oliv*0045                     flds.extend(val.split())
                0046                 elif var == 'Regions':
                0047                     regs = val.split()
ae10cb9ac4 Oliv*0048
8fc3047e23 Oliv*0049         nreg = len(regs)
                0050         res = dict((fld,[[] for reg in regs]) for fld in flds)
                0051         itrs = dict((fld,[[] for reg in regs]) for fld in flds)
ae10cb9ac4 Oliv*0052
8fc3047e23 Oliv*0053         fieldline = None
ae10cb9ac4 Oliv*0054         for line in f:
                0055             if line.strip() == '':
                0056                 continue
                0057
                0058             if line.startswith('# records'):
                0059                 break
                0060
8fc3047e23 Oliv*0061             if fieldline is not None:
                0062                 assert line.startswith(' k')
                0063                 # parse field information from saved line and discard 'k' line
                0064                 line = fieldline
                0065
ae10cb9ac4 Oliv*0066             m = re.match(r' field : *([^ ]*) *; Iter = *([0-9]*) *; region # *([0-9]*) ; nb\.Lev = *([0-9]*)', line)
                0067             if m:
                0068                 fld,itr,reg,nlev = m.groups()
8fc3047e23 Oliv*0069                 ireg = regs.index(reg)
                0070                 itrs[fld][ireg].append(int(itr))
ae10cb9ac4 Oliv*0071                 tmp = np.zeros((int(nlev)+1,nstats))
8fc3047e23 Oliv*0072                 fieldline = None
                0073                 kmax = 0
ae10cb9ac4 Oliv*0074                 for line in f:
                0075                     if line.startswith(' k'):
                0076                         continue
                0077
                0078                     if line.strip() == '':
                0079                         break
                0080
8fc3047e23 Oliv*0081                     if line.startswith(' field :'):
                0082                         fieldline = line
                0083                         break
                0084
ae10cb9ac4 Oliv*0085                     cols = line.strip().split()
                0086                     k = int(cols[0])
                0087                     tmp[k] = [float(s) for s in cols[1:]]
8fc3047e23 Oliv*0088                     kmax = max(kmax, k)
ae10cb9ac4 Oliv*0089
8fc3047e23 Oliv*0090                 res[fld][ireg].append(tmp[:kmax+1])
ae10cb9ac4 Oliv*0091             else:
                0092                 raise ValueError('readstats: parse error: ' + line)
                0093
8fc3047e23 Oliv*0094     # assume all regions have the same iteration numbers
                0095     for fld in itrs:
                0096         itrs[fld] = itrs[fld][0]
                0097
                0098     # if shapes differ between fields, we return dictionaries instead
                0099     # of record array
                0100     asdict = False
                0101     shp = None
                0102     for fld in res:
                0103         res[fld] = np.array(res[fld])
                0104         if nreg == 1:
                0105             # remove region axis
                0106             res[fld].shape = res[fld].shape[1:]
                0107         else:
                0108             # iteration axis first, then region
                0109             res[fld] = res[fld].swapaxes(0,1)
                0110
                0111         if res[fld].size == 0:
                0112             # avoid indexing error later
                0113             res[fld].shape = res[fld].shape + (1,5)
                0114
                0115         if shp is None:
                0116             shp = res[fld].shape
                0117         else:
                0118             if res[fld].shape != shp:
                0119                 asdict = True
                0120
                0121     if asdict:
                0122         statsVertInt  = dict((fld,res[fld][...,0,:]) for fld in flds)
                0123         statsPerLayer = dict((fld,res[fld][...,1:,:]) for fld in flds)
                0124     else:
                0125         ra = np.rec.fromarrays([res[fld] for fld in flds], names=flds)
                0126         statsVertInt  = ra[...,0,:]
                0127         statsPerLayer = ra[...,1:,:]
                0128
                0129     return statsPerLayer,statsVertInt,itrs