import os from netCDF4 import Dataset from argparse import ArgumentParser import numpy as np import sys # # Basic iceberg trajectory restart post-processing python script. # This script collects iceberg information from the distributed restarts written # out by each processing region and writes the information into a global restart file. # The global restart file must already exist and contain the collated 2D spatial fields # as prepared by utilities such as rebuild_nemo. This python script simply adds the # iceberg position and state data that is held using the unlimited dimension 'n' which # has been ignored by rebuild_nemo. Each processing region that contains icebergs will # (probably) have a different number of icebergs (indicated by differing values for the # current size of the unlimited dimension). This script collects all icebergs into a # single unordered list. # parser = ArgumentParser(description='produce a global trajectory restart file from distributed output\ files, e.g. \n python ./icb_pp.py \ -f icebergs_00692992_restart_ -n 480 -o icebergs_restart.nc [-O]') parser.add_argument('-f',dest='froot',help='fileroot_of_distrbuted_data; root name of \ distributed trajectory restart file (usually completed with XXXX.nc, where \ XXXX is the 4 digit processor number)', default='icebergs_00692992_restart_') parser.add_argument('-n',dest='fnum',help='number of distributed files to process', type=int, default=None) parser.add_argument('-o',dest='fout',help='global_iceberg_restart; file name to append the \ global iceberg restart data to.', default='icebergs_restart.nc') parser.add_argument('-O',dest='fcre',help='Create the output file from scratch rather than \ append to an existing file.', \ action='store_true', default=False) args = parser.parse_args() default_used = 0 if args.froot is None: pathstart = 'icebergs_00692992_restart_' default_used = 1 else: pathstart = args.froot if args.fnum is None: procnum = 0 default_used = 1 else: procnum = args.fnum if args.fout is None: pathout = 'icebergs_restart.nc' default_used = 1 else: pathout = args.fout if default_used == 1: print('At least one default value will be used; command executing is:') print('icb_combrest.py -f ',pathstart,' -n ',procnum,' -o ',pathout) if procnum < 1: print('Need some files to collate! procnum = ',procnum) sys.exit(11) icu = [] times = [] ntraj = 0 nk = 0 # # Loop through all distributed datasets to obtain the total number # of icebergs to transfer # for n in range(procnum): nn = '%4.4d' % n try: fw = Dataset(pathstart+nn+'.nc') except: print 'Error: unable to open input file: ' + pathstart+nn+'.nc' sys.exit(12) for d in fw.dimensions : if d == 'n' : if len(fw.dimensions['n']) > 0: # print 'icebergs found in: ' + pathstart+nn+'.nc' if len(fw.dimensions['k']) > nk : nk = len(fw.dimensions['k']) ntraj = ntraj + len(fw.dimensions['n']) fw.close() # print(ntraj, ' icebergs found across all datasets') # # Declare 2-D arrays to receive the data from all files # lons = np.zeros(ntraj) lats = np.zeros(ntraj) xis = np.zeros(ntraj) yjs = np.zeros(ntraj) uvs = np.zeros(ntraj) vvs = np.zeros(ntraj) mas = np.zeros(ntraj) ths = np.zeros(ntraj) wis = np.zeros(ntraj) les = np.zeros(ntraj) dys = np.zeros(ntraj) mss = np.zeros(ntraj) msb = np.zeros(ntraj) hds = np.zeros(ntraj) yrs = np.zeros(ntraj , dtype=int) num = np.zeros((ntraj, nk), dtype=int) # # loop through distributed datasets again, this time # collecting all trajectory data # nt = 0 for n in range(procnum): nn = '%4.4d' % n fw = Dataset(pathstart+nn+'.nc') for d in fw.dimensions : if d == 'n' : # Note many distributed datafiles will contain no iceberg data # so skip quickly over these m = len(fw.dimensions['n']) if m > 0: # print pathstart+nn+'.nc' lons[nt:nt+m] = fw.variables['lon'][:] lats[nt:nt+m] = fw.variables['lat'][:] xis[nt:nt+m] = fw.variables['xi'][:] yjs[nt:nt+m] = fw.variables['yj'][:] uvs[nt:nt+m] = fw.variables['uvel'][:] vvs[nt:nt+m] = fw.variables['vvel'][:] mas[nt:nt+m] = fw.variables['mass'][:] ths[nt:nt+m] = fw.variables['thickness'][:] wis[nt:nt+m] = fw.variables['width'][:] les[nt:nt+m] = fw.variables['length'][:] dys[nt:nt+m] = fw.variables['day'][:] mss[nt:nt+m] = fw.variables['mass_scaling'][:] msb[nt:nt+m] = fw.variables['mass_of_bits'][:] hds[nt:nt+m] = fw.variables['heat_density'][:] yrs[nt:nt+m] = fw.variables['year'][:] num[nt:nt+m,:] = fw.variables['number'][:,:] nt = nt + m fw.close() # Finally create the output file and write out the collated sets # if args.fcre : try: fo = Dataset(pathout, 'w', format='NETCDF4') except: print 'Error accessing output file: ' + pathout print 'Check it is a writable location.' sys.exit(13) else : # Copy 2D variables across to output file from input file. This step avoids problems if rebuild_nemo # has created an "n" dimension in the prototype rebuilt file (ie. if there are icebergs on the zeroth # processor). try: os.rename(pathout,pathout.replace('.nc','_WORK.nc')) except OSError: print 'Error: unable to move icebergs restart file: '+pathout sys.exit(14) # try: fi = Dataset(pathout.replace('.nc','_WORK.nc'), 'r') except: print 'Error: unable to open icebergs restart file: '+pathout.replace('.nc','_WORK.nc') sys.exit(15) fo = Dataset(pathout, 'w') for dim in ['x','y','c','k']: indim = fi.dimensions[dim] fo.createDimension(dim, len(indim)) for var in ['kount','calving','calving_hflx','stored_ice','stored_heat']: invar = fi.variables[var] fo.createVariable(var, invar.datatype, invar.dimensions) fo.variables[var][:] = invar[:] if "long_name" in invar.ncattrs(): fo.variables[var].long_name = invar.long_name if "units" in invar.ncattrs(): fo.variables[var].units = invar.units os.remove(pathout.replace('.nc','_WORK.nc')) # add_k = 1 for d in fo.dimensions : if d == 'n' : print 'Error: dimension n already exists in output file' sys.exit(16) if d == 'k' : add_k = 0 onn = fo.createDimension('n', None) if add_k > 0 : onk = fo.createDimension('k', nk) olon = fo.createVariable('lon', 'f8',('n')) olat = fo.createVariable('lat', 'f8',('n')) oxis = fo.createVariable('xi', 'f8',('n')) oyjs = fo.createVariable('yj', 'f8',('n')) ouvs = fo.createVariable('uvel', 'f8',('n')) ovvs = fo.createVariable('vvel', 'f8',('n')) omas = fo.createVariable('mass', 'f8',('n')) oths = fo.createVariable('thickness', 'f8',('n')) owis = fo.createVariable('width', 'f8',('n')) oles = fo.createVariable('length', 'f8',('n')) odys = fo.createVariable('day', 'f8',('n')) omss = fo.createVariable('mass_scaling', 'f8',('n')) omsb = fo.createVariable('mass_of_bits', 'f8',('n')) ohds = fo.createVariable('heat_density', 'f8',('n')) oyrs = fo.createVariable('year', 'i4',('n')) onum = fo.createVariable('number', 'i4',('n','k')) # olon[:] = lons olon.long_name = "longitude" olon.units = "degrees_E" # olat[:] = lats olat.long_name = "latitude" olat.units = "degrees_N" # oxis[:] = xis oxis.long_name = "x grid box position" oxis.units = "fractional" # oyjs[:] = yjs oyjs.long_name = "y grid box position" oyjs.units = "fractional" # ouvs[:] = uvs ouvs.long_name = "zonal velocity" ouvs.units = "m/s" # ovvs[:] = vvs ovvs.long_name = "meridional velocity" ovvs.units = "m/s" # omas[:] = mas omas.long_name = "mass" omas.units = "kg" # oths[:] = ths oths.long_name = "thickness" oths.units = "m" # owis[:] = wis owis.long_name = "width" owis.units = "m" # oles[:] = les oles.long_name = "length" oles.units = "m" # odys[:] = dys odys.long_name = "year day of calving event" odys.units = "days" # omss[:] = mss omss.long_name = "scaling factor for mass of calving berg" omss.units = "none" # omsb[:] = msb omsb.long_name = "mass of bergy bits" omsb.units = "kg" # ohds[:] = hds ohds.long_name = "heat density" ohds.units = "J/kg" # oyrs[:] = yrs oyrs.long_name = "calendar year of calving event" oyrs.units = "years" # onum[:,:] = num onum.long_name = "iceberg number on this processor" onum.units = "count" # fo.close()