Add capability to read several regions

- read field names from multiple lines
- avoid array of zeros for k-dep part of 2-D variables

utils/python/MITgcmutils/MITgcmutils/diagnostics.py

@@ -30,53 +30,97 @@ def readstats(fname):
 
     '''
     flds = []
+    regs = [0]
     with open(fname) as f:
         for line in f:
             if line.startswith('# end of header'):
                 break
 
-            m = re.match(r'^# ([^:]*) *: *(.*)$', line.rstrip())
+            m = re.match(r'^# ([^: ]*) *: *(.*)$', line.rstrip())
             if m:
                 var,val = m.groups()
                 if var.startswith('Fields'):
-                    flds = val.split()
+                    flds.extend(val.split())
+                elif var == 'Regions':
+                    regs = val.split()
 
-        res = dict((fld,[]) for fld in flds)
-        itrs = dict((fld,[]) for fld in flds)
+        nreg = len(regs)
+        res = dict((fld,[[] for reg in regs]) for fld in flds)
+        itrs = dict((fld,[[] for reg in regs]) for fld in flds)
 
+        fieldline = None
         for line in f:
             if line.strip() == '':
                 continue
 
             if line.startswith('# records'):
                 break
 
+            if fieldline is not None:
+                assert line.startswith(' k')
+                # parse field information from saved line and discard 'k' line
+                line = fieldline
+
             m = re.match(r' field : *([^ ]*) *; Iter = *([0-9]*) *; region # *([0-9]*) ; nb\.Lev = *([0-9]*)', line)
             if m:
                 fld,itr,reg,nlev = m.groups()
-                itrs[fld].append(int(itr))
+                ireg = regs.index(reg)
+                itrs[fld][ireg].append(int(itr))
                 tmp = np.zeros((int(nlev)+1,nstats))
+                fieldline = None
+                kmax = 0
                 for line in f:
                     if line.startswith(' k'):
                         continue
 
                     if line.strip() == '':
                         break
 
+                    if line.startswith(' field :'):
+                        fieldline = line
+                        break
+
                     cols = line.strip().split()
                     k = int(cols[0])
                     tmp[k] = [float(s) for s in cols[1:]]
+                    kmax = max(kmax, k)
 
-                res[fld].append(tmp)
-
+                res[fld][ireg].append(tmp[:kmax+1])
             else:
                 raise ValueError('readstats: parse error: ' + line)
 
-    try:
-        all = np.rec.fromarrays([np.array(res[fld]) for fld in flds], names=flds)
-        return all[:,1:],all[:,0],itrs
-    except:
-        totals = dict((fld,np.array(res[fld])[:,0]) for fld in flds)
-        locals = dict((fld,np.array(res[fld])[:,1:]) for fld in flds)
-        return locals,totals,itrs
-
+    # assume all regions have the same iteration numbers
+    for fld in itrs:
+        itrs[fld] = itrs[fld][0]
+
+    # if shapes differ between fields, we return dictionaries instead of record array
+    asdict = False
+    shp = None
+    for fld in res:
+        res[fld] = np.array(res[fld])
+        if nreg == 1:
+            # remove region axis
+            res[fld].shape = res[fld].shape[1:]
+        else:
+            # iteration axis first, then region
+            res[fld] = res[fld].swapaxes(0,1)
+
+        if res[fld].size == 0:
+            # avoid indexing error later
+            res[fld].shape = res[fld].shape + (1,5)
+
+        if shp is None:
+            shp = res[fld].shape
+        else:
+            if res[fld].shape != shp:
+                asdict = True
+
+    if asdict:
+        tots = dict((fld,res[fld][...,0,:]) for fld in flds)
+        locs = dict((fld,res[fld][...,1:,:]) for fld in flds)
+    else:
+        ra = np.rec.fromarrays([res[fld] for fld in flds], names=flds)
+        tots = ra[...,0,:]
+        locs = ra[...,1:,:]
+
+    return locs,tots,itrs