def ocean_atlas_enhance(output_dir):
output_file = output_dir.parent / 'ocean_atlas.nc'
output_tmp_file = output_dir.parent / 'ocean_atlas_tmp.nc'
# change value that fits in a packed short var
nco = Nco()
options = ['-O', '-a _FillValue,,o,f,-127']
nco.ncatted(input=str(output_file),
output=str(output_tmp_file),
options=options)
# pack to use bytes
# - requires output file defined with -o option
options = ['-O', '-M flt_byt', f"-o {str(output_file)}"]
nco.ncpdq(input=str(output_tmp_file),
output=str(output_file),
options=options)
def combinelon(prefix, inputdir, fill_value='1e20', daily=False, year=None):
if daily:
files = [
inputdir + sep + f
for f in filter(listdir(inputdir), '%s*.%d.psims.nc' %
(prefix, year))
]
else:
files = [
inputdir + sep + f
for f in filter(listdir(inputdir), '%s*.psims.nc' % prefix)
]
# tile latitude and longitude indices
tlatidx = basename(files[0]).split('_')[1]
lonidx = [int(basename(f).split('_')[2][:4]) for f in files]
# get file information
with nc(files[0]) as f:
vars = setdiff1d(f.variables.keys(),
['time', 'scen', 'irr', 'lat', 'lon'])
nscen = f.variables['scen'].size
ldim = f.variables[vars[0]].dimensions[0]
vunits = [0] * len(vars)
vlnames = [0] * len(vars)
for i in range(len(vars)):
var = f.variables[vars[i]]
vunits[i] = var.units if 'units' in var.ncattrs() else ''
vlnames[i] = var.long_name if 'long_name' in var.ncattrs() else ''
# fill longitude gaps
for idx in setdiff1d(fulllonidx, lonidx):
if daily:
lonfile = inputdir + sep + '%s_%s_%04d.%d.psims.nc' % (
prefix, tlatidx, idx, year)
else:
lonfile = inputdir + sep + '%s_%s_%04d.psims.nc' % (prefix,
tlatidx, idx)
copyfile(files[0], lonfile)
lons = arange(-180 + tlond * (idx - 1) + lond / 2., -180 + tlond * idx,
lond)
with nc(lonfile, 'a') as f:
lonvar = f.variables['lon']
lonvar[:] = lons
for i in range(len(vars)):
var = f.variables[vars[i]]
var[:] = masked_array(zeros(var.shape), mask=ones(var.shape))
files.append(lonfile)
# output file
if daily:
outputfile = outputdir + sep + '%s_%s.%d.psims.nc' % (prefix, tlatidx,
year)
else:
outputfile = outputdir + sep + '%s_%s.psims.nc' % (prefix, tlatidx)
nco = Nco()
# make longitude lead dimension
for i in range(len(files)):
nco.ncpdq(input=files[i],
output=files[i],
options='-O -h -a lon,%s' % str(ldim))
# concatenate all files
if daily:
inputfiles = ' '.join([
inputdir + sep + '%s_%s_%04d.%d.psims.nc' %
(prefix, tlatidx, idx, year) for idx in fulllonidx
])
else:
inputfiles = ' '.join([
inputdir + sep + '%s_%s_%04d.psims.nc' % (prefix, tlatidx, idx)
for idx in fulllonidx
])
nco.ncrcat(input=inputfiles, output=outputfile, options='-h')
# make latitude lead dimension
nco.ncpdq(input=outputfile, output=outputfile, options='-O -h -a lat,lon')
# add new scenario dimension
nscennew = nscen / (1 + irrflag)
scen_range = ','.join([str(s) for s in range(1, nscennew + 1)])
scenopt = '-O -h -s \'defdim("scen_new",%d)\' -s "scen_new[scen_new]={%s}"' % (
nscennew, scen_range)
nco.ncap2(input=outputfile, output=outputfile, options=scenopt)
nco.ncatted(input=outputfile,
output=outputfile,
options='-O -h -a units,scen_new,c,c,"no"')
nco.ncatted(input=outputfile,
output=outputfile,
options='-O -h -a long_name,scen_new,c,c,"scenario"')
# add irr dimension
nirr = 1 + irrflag
irr_range = ','.join([str(i) for i in range(1, nirr + 1)])
irr_lname = ['ir', 'rf'][:1 + irrflag] if irr1st else ['rf', 'ir'][:1 +
irrflag]
irropt = '-O -h -s \'defdim("irr",%d)\' -s "irr[irr]={%s}"' % (nirr,
irr_range)
nco.ncap2(input=outputfile, output=outputfile, options=irropt)
nco.ncatted(input=outputfile,
output=outputfile,
options='-O -h -a units,irr,c,c,"mapping"')
nco.ncatted(input=outputfile,
output=outputfile,
options='-O -h -a long_name,irr,c,c,"%s"' %
','.join(irr_lname))
# refactor variables
for i in range(len(vars)):
var = str(vars[i])
# create new variable
opt = '-O -h -s "\'%s_new\'[lat,scen_new,irr,lon,time]=0.0f"' % var
nco.ncap2(input=outputfile, output=outputfile, options=opt)
# set attributes
opt = '-O -h -a _FillValue,%s_new,c,f,%s' % (var, fill_value)
nco.ncatted(input=outputfile, output=outputfile, options=opt)
if vunits[i]:
opt = '-O -h -a units,%s_new,c,c,"%s"' % (var, str(vunits[i]))
nco.ncatted(input=outputfile, output=outputfile, options=opt)
if vlnames[i]:
opt = '-O -h -a long_name,%s_new,c,c,"%s"' % (var, str(vlnames[i]))
nco.ncatted(input=outputfile, output=outputfile, options=opt)
# set value
opt = '-O -h -s "\'%s_new\'(:,:,:,:,:)=\'%s\'"' % (var, var)
nco.ncap2(input=outputfile, output=outputfile, options=opt)
# remove old variable
opt = '-O -h -x -v %s' % var
nco.ncks(input=outputfile, output=outputfile, options=opt)
# rename new variable
opt = '-O -h -v %s_new,%s' % (var, var)
nco.ncrename(input=outputfile, output=outputfile, options=opt)
# remove old scenario dimension
nco.ncks(input=outputfile, output=outputfile, options='-O -h -x -v scen')
nco.ncrename(input=outputfile,
output=outputfile,
options='-O -h -v scen_new,scen')
# limit spatial extent to sim grid
nco.ncks(input=outputfile,
output=outputfile,
options='-O -h -d lon,%f,%f' % (lon0, lon1))
(slatidx - 1) + latdelta * i) / latdelta)
filename = 'daily_%04d.%04d.psims.nc' % (latidx, year)
files_to_remove.append(filename)
pool.apply_async(combine_daily, [latidx, lonidx, year, filename])
pool.close()
pool.join()
# Concatenate along latitude and permute dimensions of final file
print "Running concat"
part_directory = os.path.join('parts', '%04d' % tslatidx)
finalfile = os.path.join(part_directory,
'output_%04d_%04d.psims.nc' % (tslatidx, tslonidx))
files_to_copy.append(finalfile)
mkdir_p(part_directory)
nco.ncrcat(input='output_*.psims.nc', output=finalfile, options='-h')
nco.ncpdq(input=finalfile, output=finalfile, options='-O -h -a lon,lat')
nco.ncpdq(input=finalfile, output=finalfile, options='-O -h -a time,lon')
# Concatenate along latitude and permute dimensions of final file
if daily and daily_combine:
print "Running concat daily"
for year in range(ref_year, ref_year + num_years):
part_directory = os.path.join('parts', '%04d' % tslatidx)
finalfile_daily = os.path.join(
part_directory,
'daily_%04d_%04d.%04d.psims.nc' % (tslatidx, tslonidx, year))
files_to_copy.append(finalfile_daily)
mkdir_p(part_directory)
nco.ncrcat(input='daily_*.%04d.psims.nc' % year,
output=finalfile_daily,
options='-h')
def combinelat(prefix, inputdir, outputfile, daily, year):
try:
if daily:
files = [
inputdir + sep + f
for f in filter(listdir(inputdir), '%s*.%d.psims.nc' %
(prefix, year))
]
else:
files = [
inputdir + sep + f
for f in filter(listdir(inputdir), '%s*.psims.nc' % prefix)
]
# tile latitude and longitude indices
latidx = [int(basename(f).split('_')[1][:4]) for f in files]
# variables
with nc(files[0]) as f:
vars = setdiff1d(f.variables.keys(),
['time', 'scen', 'irr', 'lat', 'lon'])
# fill longitude gaps
for idx in setdiff1d(fulllatidx, latidx):
if daily:
latfile = inputdir + sep + '%s_%04d.%d.psims.nc' % (prefix,
idx, year)
else:
latfile = inputdir + sep + '%s_%04d.psims.nc' % (prefix, idx)
copyfile(files[0], latfile)
lats = arange(90 - tlatd * (idx - 1) - latd / 2., 90 - tlatd * idx,
-latd)
with nc(latfile, 'a') as f:
latvar = f.variables['lat']
latvar[:] = lats
for i in range(len(vars)):
var = f.variables[vars[i]]
var[:] = masked_array(zeros(var.shape),
mask=ones(var.shape))
files.append(latfile)
# output file
nco = Nco()
# concatenate all files
if daily:
inputfiles = ' '.join([
inputdir + sep + '%s_%04d.%d.psims.nc' % (prefix, idx, year)
for idx in fulllatidx
])
else:
inputfiles = ' '.join([
inputdir + sep + '%s_%04d.psims.nc' % (prefix, idx)
for idx in fulllatidx
])
nco.ncrcat(input=inputfiles, output=outputfile, options='-h')
# make time lead dimension
if timelead:
nco.ncpdq(input=outputfile,
output=outputfile,
options='-O -h -a time,lat')
# limit spatial extent to sim grid
nco.ncks(input=outputfile,
output=outputfile,
options='-O -h -d lat,%f,%f' % (lat1, lat0))
# compress
system('nccopy -d9 -k4 %s %s' % (outputfile, outputfile + '.tmp'))
rename(outputfile + '.tmp', outputfile)
except:
print "[%s] %s" % (os.path.basename(__file__), traceback.format_exc())
