Index: branches/2017/dev_r8126_ROBUST08_no_ghost/NEMOGCM/TOOLS/MISCELLANEOUS/icb_pp.py
===================================================================
--- branches/2017/dev_r8126_ROBUST08_no_ghost/NEMOGCM/TOOLS/MISCELLANEOUS/icb_pp.py	(revision 8165)
+++ branches/2017/dev_r8126_ROBUST08_no_ghost/NEMOGCM/TOOLS/MISCELLANEOUS/icb_pp.py	(revision 8979)
@@ -12,19 +12,26 @@
 # cases the missing instances are filled with invalid (NaN) values.
 #
+# Version 2.0 August 2017. Adapted to process all variables and retain original
+#                          datatypes. (acc@noc.ac.uk)
 
-parser = ArgumentParser(description='produce collated trajectory file from distributed output\
-                                     files, e.g. \n python ./icb_pp.py \
-                                     -t  trajectory_icebergs_004248_ -n 296 -o trajsout.nc' )
+parser = ArgumentParser(description='produce collated trajectory file \
+                                     from distributed output files, e.g. \
+                                     \n python ./icb_pp.py \
+                                     -t  trajectory_icebergs_004248_ \
+                                     -n 296 -o trajsout.nc' )
 
-parser.add_argument('-t',dest='froot',help='fileroot_of_distrbuted_data; root name of \
-                     distributed trajectory output (usually completed with XXXX.nc, where \
-                     XXXX is the 4 digit processor number)', 
-                     default='trajectory_icebergs_004248_')
+parser.add_argument('-t',dest='froot',
+                         help='fileroot_of_distrbuted_data; root name \
+                               of  distributed trajectory output (usually \
+                               completed with XXXX.nc, where  XXXX is the \
+                               4 digit processor number)', 
+                      default='trajectory_icebergs_004248_')
 
 parser.add_argument('-n',dest='fnum',help='number of distributed files to process', 
-                     type=int, default=None)
+                         type=int, default=None)
 
-parser.add_argument('-o',dest='fout',help='collated_output_file; file name to receive the \
-                     collated trajectory data', default='trajsout.nc')
+parser.add_argument('-o',dest='fout',
+                         help='collated_output_file; file name to receive \
+                              the collated trajectory data', default='trajsout.nc')
 
 args = parser.parse_args()
@@ -64,17 +71,22 @@
 #
 for n in range(procnum):
- nn = '%4.4d' % n
- fw = Dataset(pathstart+nn+'.nc')
- if len(fw.dimensions['n']) > 0:
-   print pathstart+nn+'.nc'
-   ic = fw.variables['iceberg_number'][:,0]
-   ts = fw.variables['timestep'][:]
-   icv = np.unique(ic)
-   ts = np.unique(ts)
-   print('Min Max ts: ',ts.min(), ts.max())
-   print('Number unique icebergs= ',icv.shape[0])
-   icu.append(icv)
-   times.append(ts)
- fw.close()
+    nn = '%4.4d' % n
+    fw = Dataset(pathstart+nn+'.nc')
+    # keep a list of the variables in the first dataset
+    if n == 0:
+        varlist = fw.variables
+    #
+    # skip any files with no icebergs
+    if len(fw.dimensions['n']) > 0:
+        print pathstart+nn+'.nc'
+        ic = fw.variables['iceberg_number'][:,0]
+        ts = fw.variables['timestep'][:]
+        icv = np.unique(ic)
+        ts = np.unique(ts)
+        print('Min Max ts: ',ts.min(), ts.max())
+        print('Number unique icebergs= ',icv.shape[0])
+        icu.append(icv)
+        times.append(ts)
+    fw.close()
 #
 # Now flatten the lists and reduce to the unique spanning set
@@ -89,71 +101,87 @@
 print('times range from:        ',times.min(), 'to: ', times.max())
 #
-# Declare 2-D arrays to receive the data from all files
+# Declare array to receive data from all files
 #
 nt = times.shape[0]
-lons = np.zeros((ntraj, nt))
-lats = np.zeros((ntraj, nt))
-tims = np.zeros((ntraj, nt))
-xis  = np.zeros((ntraj, nt))
-yjs  = np.zeros((ntraj, nt))
+#
+n=0
+for key, value in varlist.iteritems() :
+    if key != "iceberg_number" :
+        n = n + 1
+inarr = np.zeros((n, ntraj, nt))
 #
 # initially fill with invalid data
 #
-lons.fill(np.nan)
-lats.fill(np.nan)
-xis.fill(np.nan)
-yjs.fill(np.nan)
-tims.fill(np.nan)
+inarr.fill(np.nan)
+#
+# Declare some lists to store variable names, types and long_name and units attributes
+# iceberg_number gets special treatment
+innam = []
+intyp = []
+inlngnam = []
+inunits = []
+for key, value in varlist.iteritems() :
+    if key != "iceberg_number" :
+        innam.append(key)
+#
+# reopen the first datset to collect variable attributes
+# (long_name and units only)
+#
+nn = '%4.4d' % 0
+fw = Dataset(pathstart+nn+'.nc')
+for key, value in varlist.iteritems() :
+    if key != "iceberg_number" :
+        intyp.append(fw.variables[key].dtype)
+        inlngnam.append(fw.variables[key].getncattr('long_name'))
+        inunits.append(fw.variables[key].getncattr('units'))
+fw.close()
 #
 # loop through distributed datasets again, this time
 # checking indices against icu and times lists and
 # inserting data into the correct locations in the 
-# 2-D collated sets.
+# collated sets.
 #
 for n in range(procnum):
- nn = '%4.4d' % n
- fw = Dataset(pathstart+nn+'.nc')
+    nn = '%4.4d' % n
+    fw = Dataset(pathstart+nn+'.nc')
+#
 # Note many distributed datafiles will contain no iceberg data
 # so skip quickly over these
- m  = len(fw.dimensions['n'])
- if m > 0:
-   inx = np.zeros(m, dtype=int)
-   tsx = np.zeros(m, dtype=int)
-   print pathstart+nn+'.nc'
-   ic = fw.variables['iceberg_number'][:,0]
-   ts = fw.variables['timestep'][:]
-   lns = fw.variables['lon'][:]
-   lts = fw.variables['lat'][:]
-   xxs = fw.variables['xi'][:]
-   yys = fw.variables['yj'][:]
-   for k in range(m):
-     inxx   = np.where(icu == ic[k])
-     inx[k] = inxx[0]
-   for k in range(m):
-     inxx   = np.where(times == ts[k])
-     tsx[k] = inxx[0]
-   lons[inx[:],tsx[:]] = lns[:]
-   lats[inx[:],tsx[:]] = lts[:]
-   tims[inx[:],tsx[:]] = ts[:]
-   xis[inx[:],tsx[:]] = xxs[:]
-   yjs[inx[:],tsx[:]] = yys[:]
- fw.close()
-
+    m  = len(fw.dimensions['n'])
+    if m > 0:
+        inx = np.zeros(m, dtype=int)
+        tsx = np.zeros(m, dtype=int)
+        #print pathstart+nn+'.nc'
+        ic = fw.variables['iceberg_number'][:,0]
+        ts = fw.variables['timestep'][:]
+        for k in range(m):
+            inxx   = np.where(icu == ic[k])
+            inx[k] = inxx[0]
+        for k in range(m):
+            inxx   = np.where(times == ts[k])
+            tsx[k] = inxx[0]
+        n = 0
+        for key, value in varlist.iteritems() :
+            if key != "iceberg_number" :
+                insmall = fw.variables[innam[n]][:]
+                inarr[n,inx[:],tsx[:]] = insmall[:]
+                n = n + 1
+    fw.close()
+#
 # Finally create the output file and write out the collated sets
 #
-fo = Dataset(pathout, 'w', format='NETCDF4')
+fo = Dataset(pathout, 'w', format='NETCDF4_CLASSIC')
 ntrj = fo.createDimension('ntraj', ntraj)
 nti  = fo.createDimension('ntime', None)
-olon = fo.createVariable('lon', 'f4',('ntraj','ntime'))
-olat = fo.createVariable('lat', 'f4',('ntraj','ntime'))
-otim = fo.createVariable('ttim', 'f4',('ntraj','ntime'))
-oxis = fo.createVariable('xis', 'f4',('ntraj','ntime'))
-oyjs = fo.createVariable('yjs', 'f4',('ntraj','ntime'))
-icbn = fo.createVariable('icbn', 'f4',('ntraj'))
-olon[:,:] = lons
-olat[:,:] = lats
-otim[:,:] = tims
-oxis[:,:] = xis
-oyjs[:,:] = yjs
+icbn = fo.createVariable('iceberg_number', 'i4',('ntraj'))
 icbn[:] = icu
+n = 0
+for key, value in varlist.iteritems() :
+    if key != "iceberg_number" :
+        oout = fo.createVariable(innam[n], intyp[n], ('ntraj','ntime'),
+                                 zlib=True, complevel=1, chunksizes=(1,nt))
+        oout[:,:] = inarr[n,:,:]
+        oout.long_name = inlngnam[n]
+        oout.units = inunits[n]
+        n = n + 1
 fo.close()
