#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Script used by `reduce_output` job
Author: Michael Jaehn (jae), Empa, Switzerland
"""
from __future__ import print_function
import sys
import numpy as np
import netCDF4 as nc
from datetime import datetime
from datetime import timedelta
import glob
import os
import shutil
import csv
import helper
import logging
logging.basicConfig(level=logging.DEBUG)
def get_attrs(v):
return dict((k, v.getncattr(k)) for k in v.ncattrs())
def str2char(s):
return np.array(list(s), 'c')
def search_met(infile, fname_met):
"""
Checks which file contains the meteorology data
"""
time = None
tocheck = []
print(infile)
for f in infile:
basename = os.path.split(f)[1]
print(basename)
timestr = basename.split('lffd', 1)[1][:10]
mytime = datetime.strptime(timestr, '%Y%m%d%H')
if time is None:
time = mytime
if mytime == time:
tocheck.append(f)
for f in tocheck:
with nc.Dataset(f, 'r') as inf:
for var in ['T', 'P', 'PS', 'QV']:
if var in inf.variables:
basename = os.path.split(f)[1]
fileending = basename.split('lffd', 1)[1][10:-3]
fname_met[var] = fileending
return fname_met
def append_variable(nc, name, values, attrs=None):
"""\
Append variable (name, values) to netCDF file.
"""
from six import string_types
if attrs is None:
attrs = {}
var = nc.createVariable(name,
values.dtype, ('time', 'rlat', 'rlon'),
zlib=True)
var[0] = values
# add or update attributes
for key, value in attrs.items():
if isinstance(value, string_types):
var.setncattr(key, str2char(value))
else:
var.setncattr(key, value)
def reduce_output(infile, cfiles, h, nout_levels, output_path, fname_met, lsd,
convert):
dtime = infile.split('lffd', 1)[1][0:10]
"""Get path and filename for output file"""
path, output_filename = os.path.split(infile)
outfile = os.path.join(output_path, output_filename)
with nc.Dataset(infile, 'r') as inf, \
nc.Dataset(outfile, 'w') as outf:
# Copy global attributes all at once via dictionary
outf.setncatts(inf.__dict__)
# Copy dimensions
for name, dimension in inf.dimensions.items():
outf.createDimension(
name,
(len(dimension) if not dimension.isunlimited() else None))
# Get level information
level = len(inf.dimensions['level'])
level1 = len(inf.dimensions['level1'])
# Create new dimension level2
if nout_levels == -1:
nout_levels = level
else:
outf.createDimension('level2', nout_levels)
# Copy variables
for varname in inf.variables.keys():
logging.info('%s: %s' % (output_filename, varname))
var = inf.variables[varname]
attrs = get_attrs(var)
if len(var.dimensions) == 4:
if var.dimensions[1] == 'level':
var_dimensions = list(var.dimensions)
var_dimensions[1] = 'level' + '2' * (nout_levels != level)
elif var.dimensions[1] == 'level1':
var_dimensions = list(var.dimensions)
var_dimensions[1] = 'level' + '2' * (
nout_levels != level) + '1' * (nout_levels == level)
else:
var_dimensions = var.dimensions
else:
var_dimensions = var.dimensions
if len(var.dimensions) > 2:
if (varname in lsd):
if '_FillValue' in var.ncattrs():
fill_value = inf[varname].getncattr('_FillValue')
outf.createVariable(
varname,
var.dtype,
var_dimensions,
zlib=True,
fill_value=fill_value,
least_significant_digit=lsd[varname])
else:
outf.createVariable(
varname,
var.dtype,
var_dimensions,
zlib=True,
least_significant_digit=lsd[varname])
else:
if '_FillValue' in var.ncattrs():
fill_value = inf[varname].getncattr('_FillValue')
outf.createVariable(varname,
var.dtype,
var_dimensions,
zlib=True,
fill_value=fill_value)
else:
outf.createVariable(varname,
var.dtype,
var_dimensions,
zlib=True)
else:
outf.createVariable(varname,
var.dtype,
var_dimensions,
zlib=True)
for attr in var.ncattrs():
if attr != '_FillValue':
outf[varname].setncattr(attr, inf[varname].getncattr(attr))
gas = None
if varname != 'rotated_pole':
# Check for 3D data and extract only lower_levels
if 'level1' in var.dimensions and \
len(var.dimensions) == 4:
lstart = -nout_levels - (nout_levels == level)
elif 'level' in var.dimensions and \
len(var.dimensions) == 4:
lstart = -nout_levels
else:
outf[varname][:] = inf[varname][:]
if (varname.startswith('CO2_') or varname.startswith('CO_') or
varname.startswith('CH4_') or varname.startswith('C14_') or
varname.startswith('NOX_') or varname.startswith('NO2_')) and \
len(var.dimensions) == 4:
outvar = outf.variables[varname]
gas = varname.split('_')[0]
if gas == 'NOX':
attrs['standard_name'] = attrs['standard_name'].replace(
'NOX', 'NO2')
gas = 'NO2'
field = inf[varname][:, lstart:, :, :]
unit = attrs['units']
if convert:
outvar.units = helper.common_unit(gas)
if gas == 'C14':
outf[varname][:] = helper.convert_unit(
field, unit, outvar.units, 45.993e-3)
else:
outf[varname][:] = helper.convert_unit(
field, unit, outvar.units, gas)
else:
outvar.units = unit
outf[varname][:] = field
attrs['standard_name'] = attrs['standard_name'].replace(
'%s_mass_fraction' % gas, 'X%s' % gas)
attrs['units'] = outvar.units
attrs2 = attrs.copy()
attrs2['standard_name'] = attrs2['standard_name'].replace(
'X%s' % gas, 'Y%s' % gas)
elif varname != 'rotated_pole' and len(var.dimensions) == 4 and \
('level1' in var.dimensions or 'level' in var.dimensions):
outf[varname][:] = inf[varname][:, lstart:, :, :]
if gas:
fname_met_base = os.path.split(
infile)[0] + '/' + os.path.split(infile)[1][:14]
with nc.Dataset(fname_met_base+fname_met['QV']+'.nc', 'r') as inf_qv, \
nc.Dataset(fname_met_base+fname_met['T']+'.nc', 'r') as inf_t, \
nc.Dataset(fname_met_base+fname_met['PS']+'.nc', 'r') as inf_ps, \
nc.Dataset(fname_met_base+fname_met['P']+'.nc', 'r') as inf_p:
qv = np.array(inf_qv.variables['QV'][0])
t = np.array(inf_t.variables['T'][0])
ps = np.array(inf_ps.variables['PS'][0])
p = np.array(inf_p.variables['P'][0])
xm = np.array(inf.variables[varname][0])
mair = helper.calculate_mair(p, ps, h)
# Column-averaged dry-air mole fraction
column = helper.calculate_xgas(xm, mair, gas, qv)
column = column.astype('f4')
logging.info('%s: X%s' % (output_filename, varname))
append_variable(outf, 'X%s' % varname, column, attrs=attrs)
# Column-averaged moist-air mole fraction
column2 = helper.calculate_xgas(xm, mair, gas, 0.0)
column2 = column2.astype('f4')
logging.info('%s: Y%s' % (output_filename, varname))
append_variable(outf,
'Y%s' % varname,
column2,
attrs=attrs2)
return fname_met
[docs]def main(indir, outdir, strdate_start, strdate_end, nout_levels, csvfile,
convert_gas):
"""
Script to reduce output.
Output: TODO
Parameters (TODO)
----------
opath : str
Output path where the processed data is going to be written
"""
"""Get list of constant files"""
cfiles = []
read_cfile = False
for infile in sorted(glob.glob(os.path.join(indir, "lffd*[0-9]c*.nc"))):
cfiles.append(infile)
if not read_cfile:
# Read the first constant file and store height value
logging.info(infile)
with nc.Dataset(infile) as inf:
h = np.array(inf.variables['HHL'][0])
read_cfile = True
if not read_cfile:
logging.error('Constant file could not be read')
"""Get list of files"""
infiles = sorted(glob.glob(os.path.join(indir, "lffd*.nc")))
"""Remove constant files from file list"""
for cfile in cfiles:
if cfile in infiles:
infiles.remove(cfile)
date_start = datetime.strptime(strdate_start, '%Y%m%d%H')
date_end = datetime.strptime(strdate_end, '%Y%m%d%H')
"""Only take relevant date period into account"""
for infile in list(infiles):
basename = os.path.split(infile)[1]
timestr = basename.split('lffd', 1)[1][:10]
mytime = datetime.strptime(timestr, '%Y%m%d%H')
if mytime < date_start or mytime > date_end:
infiles.remove(infile)
"""Initialize dict to remember file name for meteo variables"""
fname_met = {
'T': None,
'P': None,
'PS': None,
'QV': None,
}
fname_met = search_met(infiles, fname_met)
"""Translate csv file to dict"""
variables = helper.find_variables_file(csvfile)
lsd = variables['lsd'].to_dict()
"""Loop over all input files and apply output reduction"""
for infile in infiles:
fname_met = reduce_output(infile, cfiles, h, int(nout_levels), outdir,
fname_met, lsd, convert_gas == 'True')
if __name__ == '__main__':
indir = sys.argv[1]
outdir = sys.argv[2]
strdate_start = sys.argv[3]
strdate_end = sys.argv[4]
nout_levels = sys.argv[5]
csvfile = sys.argv[6]
convert_gas = sys.argv[7]
main(indir, outdir, strdate_start, strdate_end, nout_levels, csvfile,
convert_gas)