[4891] | 1 | from netCDF4 import Dataset |
---|
| 2 | from argparse import ArgumentParser |
---|
| 3 | import numpy as np |
---|
| 4 | import sys |
---|
| 5 | |
---|
| 6 | # |
---|
| 7 | # Basic iceberg trajectory post-processing python script. |
---|
| 8 | # This script collates iceberg trajectories from the distributed datasets written |
---|
| 9 | # out by each processing region and rearranges the ragged arrays into contiguous |
---|
| 10 | # streams for each unique iceberg. The output arrays are 2D (ntraj, ntimes) arrays. |
---|
| 11 | # Note that some icebergs may only exist for a subset of the possible times. In these |
---|
| 12 | # cases the missing instances are filled with invalid (NaN) values. |
---|
| 13 | # |
---|
| 14 | |
---|
| 15 | parser = ArgumentParser(description='produce collated trajectory file from distributed output\ |
---|
| 16 | files, e.g. \n python ./icb_pp.py \ |
---|
| 17 | -t trajectory_icebergs_004248_ -n 296 -o trajsout.nc' ) |
---|
| 18 | |
---|
| 19 | parser.add_argument('-t',dest='froot',help='fileroot_of_distrbuted_data; root name of \ |
---|
| 20 | distributed trajectory output (usually completed with XXXX.nc, where \ |
---|
| 21 | XXXX is the 4 digit processor number)', |
---|
| 22 | default='trajectory_icebergs_004248_') |
---|
| 23 | |
---|
| 24 | parser.add_argument('-n',dest='fnum',help='number of distributed files to process', |
---|
| 25 | type=int, default=None) |
---|
| 26 | |
---|
| 27 | parser.add_argument('-o',dest='fout',help='collated_output_file; file name to receive the \ |
---|
| 28 | collated trajectory data', default='trajsout.nc') |
---|
| 29 | |
---|
| 30 | args = parser.parse_args() |
---|
| 31 | |
---|
| 32 | default_used = 0 |
---|
| 33 | if args.froot is None: |
---|
| 34 | pathstart = 'trajectory_icebergs_004248_' |
---|
| 35 | default_used = 1 |
---|
| 36 | else: |
---|
| 37 | pathstart = args.froot |
---|
| 38 | |
---|
| 39 | if args.fnum is None: |
---|
| 40 | procnum = 0 |
---|
| 41 | default_used = 1 |
---|
| 42 | else: |
---|
| 43 | procnum = args.fnum |
---|
| 44 | |
---|
| 45 | if args.fout is None: |
---|
| 46 | pathout = 'trajsout.nc' |
---|
| 47 | default_used = 1 |
---|
| 48 | else: |
---|
| 49 | pathout = args.fout |
---|
| 50 | |
---|
| 51 | if default_used == 1: |
---|
| 52 | print('At least one default value will be used; command executing is:') |
---|
| 53 | print('icb_pp.py -t ',pathstart,' -n ',procnum,' -o ',pathout) |
---|
| 54 | |
---|
| 55 | if procnum < 1: |
---|
| 56 | print('Need some files to collate! procnum = ',procnum) |
---|
[6423] | 57 | sys.exit(11) |
---|
[4891] | 58 | |
---|
| 59 | icu = [] |
---|
| 60 | times = [] |
---|
| 61 | # |
---|
| 62 | # Loop through all distributed datasets to obtain the complete list |
---|
| 63 | # of iceberg identification numbers and timesteps |
---|
| 64 | # |
---|
| 65 | for n in range(procnum): |
---|
| 66 | nn = '%4.4d' % n |
---|
| 67 | fw = Dataset(pathstart+nn+'.nc') |
---|
| 68 | if len(fw.dimensions['n']) > 0: |
---|
| 69 | print pathstart+nn+'.nc' |
---|
| 70 | ic = fw.variables['iceberg_number'][:,0] |
---|
| 71 | ts = fw.variables['timestep'][:] |
---|
| 72 | icv = np.unique(ic) |
---|
| 73 | ts = np.unique(ts) |
---|
| 74 | print('Min Max ts: ',ts.min(), ts.max()) |
---|
| 75 | print('Number unique icebergs= ',icv.shape[0]) |
---|
| 76 | icu.append(icv) |
---|
| 77 | times.append(ts) |
---|
| 78 | fw.close() |
---|
| 79 | # |
---|
| 80 | # Now flatten the lists and reduce to the unique spanning set |
---|
| 81 | # |
---|
| 82 | icu = np.concatenate(icu) |
---|
| 83 | icu = np.unique(icu) |
---|
| 84 | times = np.concatenate(times) |
---|
| 85 | times = np.unique(times) |
---|
| 86 | ntraj = icu.shape[0] |
---|
| 87 | print(ntraj, ' unique icebergs found across all datasets') |
---|
| 88 | print('Icebergs ids range from: ',icu.min(), 'to: ',icu.max()) |
---|
| 89 | print('times range from: ',times.min(), 'to: ', times.max()) |
---|
| 90 | # |
---|
| 91 | # Declare 2-D arrays to receive the data from all files |
---|
| 92 | # |
---|
| 93 | nt = times.shape[0] |
---|
| 94 | lons = np.zeros((ntraj, nt)) |
---|
| 95 | lats = np.zeros((ntraj, nt)) |
---|
| 96 | tims = np.zeros((ntraj, nt)) |
---|
| 97 | xis = np.zeros((ntraj, nt)) |
---|
| 98 | yjs = np.zeros((ntraj, nt)) |
---|
| 99 | # |
---|
| 100 | # initially fill with invalid data |
---|
| 101 | # |
---|
| 102 | lons.fill(np.nan) |
---|
| 103 | lats.fill(np.nan) |
---|
| 104 | xis.fill(np.nan) |
---|
| 105 | yjs.fill(np.nan) |
---|
| 106 | tims.fill(np.nan) |
---|
| 107 | # |
---|
| 108 | # loop through distributed datasets again, this time |
---|
| 109 | # checking indices against icu and times lists and |
---|
| 110 | # inserting data into the correct locations in the |
---|
| 111 | # 2-D collated sets. |
---|
| 112 | # |
---|
| 113 | for n in range(procnum): |
---|
| 114 | nn = '%4.4d' % n |
---|
| 115 | fw = Dataset(pathstart+nn+'.nc') |
---|
| 116 | # Note many distributed datafiles will contain no iceberg data |
---|
| 117 | # so skip quickly over these |
---|
| 118 | m = len(fw.dimensions['n']) |
---|
| 119 | if m > 0: |
---|
| 120 | inx = np.zeros(m, dtype=int) |
---|
| 121 | tsx = np.zeros(m, dtype=int) |
---|
| 122 | print pathstart+nn+'.nc' |
---|
| 123 | ic = fw.variables['iceberg_number'][:,0] |
---|
| 124 | ts = fw.variables['timestep'][:] |
---|
| 125 | lns = fw.variables['lon'][:] |
---|
| 126 | lts = fw.variables['lat'][:] |
---|
| 127 | xxs = fw.variables['xi'][:] |
---|
| 128 | yys = fw.variables['yj'][:] |
---|
| 129 | for k in range(m): |
---|
| 130 | inxx = np.where(icu == ic[k]) |
---|
| 131 | inx[k] = inxx[0] |
---|
| 132 | for k in range(m): |
---|
| 133 | inxx = np.where(times == ts[k]) |
---|
| 134 | tsx[k] = inxx[0] |
---|
| 135 | lons[inx[:],tsx[:]] = lns[:] |
---|
| 136 | lats[inx[:],tsx[:]] = lts[:] |
---|
| 137 | tims[inx[:],tsx[:]] = ts[:] |
---|
| 138 | xis[inx[:],tsx[:]] = xxs[:] |
---|
| 139 | yjs[inx[:],tsx[:]] = yys[:] |
---|
| 140 | fw.close() |
---|
| 141 | |
---|
| 142 | # Finally create the output file and write out the collated sets |
---|
| 143 | # |
---|
| 144 | fo = Dataset(pathout, 'w', format='NETCDF4') |
---|
| 145 | ntrj = fo.createDimension('ntraj', ntraj) |
---|
| 146 | nti = fo.createDimension('ntime', None) |
---|
| 147 | olon = fo.createVariable('lon', 'f4',('ntraj','ntime')) |
---|
| 148 | olat = fo.createVariable('lat', 'f4',('ntraj','ntime')) |
---|
| 149 | otim = fo.createVariable('ttim', 'f4',('ntraj','ntime')) |
---|
| 150 | oxis = fo.createVariable('xis', 'f4',('ntraj','ntime')) |
---|
| 151 | oyjs = fo.createVariable('yjs', 'f4',('ntraj','ntime')) |
---|
| 152 | icbn = fo.createVariable('icbn', 'f4',('ntraj')) |
---|
| 153 | olon[:,:] = lons |
---|
| 154 | olat[:,:] = lats |
---|
| 155 | otim[:,:] = tims |
---|
| 156 | oxis[:,:] = xis |
---|
| 157 | oyjs[:,:] = yjs |
---|
| 158 | icbn[:] = icu |
---|
| 159 | fo.close() |
---|