2025-09-05 09:42:21 +02:00

276 lines
8.3 KiB
Python
Executable File

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()