def __init__(self, source, target=None, varlist=None, ignorelist=None, tmp=True, feedback=True):
''' Initialize processor and pass input and output datasets. '''
# check varlist
if varlist is None: varlist = source.variables.keys() # all source variables
elif not isinstance(varlist,(list,tuple)): raise TypeError
self.varlist = varlist # list of variable to be processed
# ignore list (e.g. variables that will cause errors)
if ignorelist is None: ignorelist = [] # an empty list
elif not isinstance(ignorelist,(list,tuple)): raise TypeError
self.ignorelist = ignorelist # list of variable *not* to be processed
# check input
if not isinstance(source,Dataset): raise TypeError
if isinstance(source,DatasetNetCDF) and not 'r' in source.mode: raise PermissionError
self.input = source
self.source = source
# check output
if target is not None:
if not isinstance(target,Dataset): raise TypeError
if isinstance(target,DatasetNetCDF) and not 'w' in target.mode: raise PermissionError
else:
if not tmp: raise DatasetError, "Need target location, if temporary storage is disables (tmp=False)."
self.output = target
# temporary dataset
self.tmp = tmp
if tmp: self.tmpput = Dataset(name='tmp', title='Temporary Dataset', varlist=[], atts={})
else: self.tmpput = None
# determine if temporary storage is used and assign target dataset
if self.tmp: self.target = self.tmpput
else: self.target = self.output
# whether or not to print status output
self.feedback = feedback
def performExtraction(dataset,
mode,
stnfct,
dataargs,
loverwrite=False,
varlist=None,
lwrite=True,
lreturn=False,
ldebug=False,
lparallel=False,
pidstr='',
logger=None):
''' worker function to extract point data from gridded dataset '''
# input checking
if not isinstance(dataset, basestring): raise TypeError
if not isinstance(dataargs, dict):
raise TypeError # all dataset arguments are kwargs
if not callable(stnfct):
raise TypeError # function to load station dataset
if lparallel:
if not lwrite:
raise IOError, 'In parallel mode we can only write to disk (i.e. lwrite = True).'
if lreturn:
raise IOError, 'Can not return datasets in parallel mode (i.e. lreturn = False).'
# logging
if logger is None: # make new logger
logger = logging.getLogger() # new logger
logger.addHandler(logging.StreamHandler())
else:
if isinstance(logger, basestring):
logger = logging.getLogger(name=logger) # connect to existing one
elif not isinstance(logger, logging.Logger):
raise TypeError, 'Expected logger ID/handle in logger KW; got {}'.format(
str(logger))
lclim = False
lts = False
if mode == 'climatology': lclim = True
elif mode == 'time-series': lts = True
else: raise NotImplementedError, "Unrecognized Mode: '{:s}'".format(mode)
## extract meta data from arguments
dataargs, loadfct, srcage, datamsgstr = getMetaData(
dataset, mode, dataargs)
dataset_name = dataargs.dataset_name
periodstr = dataargs.periodstr
avgfolder = dataargs.avgfolder
# load template dataset
stndata = stnfct() # load station dataset from function
if not isinstance(stndata, Dataset): raise TypeError
# N.B.: the loading function is necessary, because DataseNetCDF instances do not pickle well
# get filename for target dataset and do some checks
filename = getTargetFile(dataset=dataset,
mode=mode,
dataargs=dataargs,
lwrite=lwrite,
station=stndata.name)
if ldebug: filename = 'test_' + filename
if not os.path.exists(avgfolder):
raise IOError, "Dataset folder '{:s}' does not exist!".format(
avgfolder)
lskip = False # else just go ahead
if lwrite:
if lreturn:
tmpfilename = filename # no temporary file if dataset is passed on (can't rename the file while it is open!)
else:
if lparallel: tmppfx = 'tmp_exstns_{:s}_'.format(pidstr[1:-1])
else: tmppfx = 'tmp_exstns_'.format(pidstr[1:-1])
tmpfilename = tmppfx + filename
filepath = avgfolder + filename
tmpfilepath = avgfolder + tmpfilename
if os.path.exists(filepath):
if not loverwrite:
age = datetime.fromtimestamp(os.path.getmtime(filepath))
# if source file is newer than sink file or if sink file is a stub, recompute, otherwise skip
if age > srcage and os.path.getsize(filepath) > 1e5:
lskip = True
# N.B.: NetCDF files smaller than 100kB are usually incomplete header fragments from a previous crashed
# depending on last modification time of file or overwrite setting, start computation, or skip
if lskip:
# print message
skipmsg = "\n{:s} >>> Skipping: file '{:s}' in dataset '{:s}' already exists and is newer than source file.".format(
pidstr, filename, dataset_name)
skipmsg += "\n{:s} >>> ('{:s}')\n".format(pidstr, filepath)
logger.info(skipmsg)
else:
## actually load datasets
source = loadfct() # load source
# check period
if 'period' in source.atts and dataargs.periodstr != source.atts.period: # a NetCDF attribute
raise DateError, "Specifed period is inconsistent with netcdf records: '{:s}' != '{:s}'".format(
periodstr, source.atts.period)
# print message
if lclim:
opmsgstr = "Extracting '{:s}'-type Point Data from Climatology ({:s})".format(
stndata.name, periodstr)
elif lts:
opmsgstr = "Extracting '{:s}'-type Point Data from Time-series".format(
stndata.name)
else:
raise NotImplementedError, "Unrecognized Mode: '{:s}'".format(mode)
# print feedback to logger
logger.info(
'\n{0:s} *** {1:^65s} *** \n{0:s} *** {2:^65s} *** \n'
.format(pidstr, datamsgstr, opmsgstr))
if not lparallel and ldebug: logger.info('\n' + str(source) + '\n')
## create new sink/target file
# set attributes
atts = source.atts.copy()
atts[
'period'] = dataargs.periodstr if dataargs.periodstr else 'time-series'
atts['name'] = dataset_name
atts['station'] = stndata.name
atts['title'] = '{:s} (Stations) from {:s} {:s}'.format(
stndata.title, dataset_name, mode.title())
# make new dataset
if lwrite: # write to NetCDF file
if os.path.exists(tmpfilepath):
os.remove(tmpfilepath) # remove old temp files
sink = DatasetNetCDF(folder=avgfolder,
filelist=[tmpfilename],
atts=atts,
mode='w')
else:
sink = Dataset(atts=atts) # ony create dataset in memory
# initialize processing
CPU = CentralProcessingUnit(source,
sink,
varlist=varlist,
tmp=False,
feedback=ldebug)
# extract data at station locations
CPU.Extract(template=stndata, flush=True)
# get results
CPU.sync(flush=True)
# print dataset
if not lparallel and ldebug:
logger.info('\n' + str(sink) + '\n')
# write results to file
if lwrite:
sink.sync()
writemsg = "\n{:s} >>> Writing to file '{:s}' in dataset {:s}".format(
pidstr, filename, dataset_name)
writemsg += "\n{:s} >>> ('{:s}')\n".format(pidstr, filepath)
logger.info(writemsg)
# rename file to proper name
if not lreturn:
sink.unload()
sink.close()
del sink # destroy all references
if os.path.exists(filepath):
os.remove(filepath) # remove old file
os.rename(tmpfilepath, filepath)
# N.B.: there is no temporary file if the dataset is returned, because an open file can't be renamed
# clean up and return
source.unload()
del source #, CPU
if lreturn:
return sink # return dataset for further use (netcdf file still open!)
else:
return 0 # "exit code"
def loadKister_StnTS(station=None,
well=None,
folder=None,
varlist='default',
varatts=None,
name='observations',
title=None,
basin=None,
start_date=None,
end_date=None,
sampling=None,
period=None,
date_range=None,
llastIncl=True,
WSC_station=None,
basin_list=None,
filenames=None,
time_axis='datetime',
scalefactors=None,
metadata=None,
lkgs=False,
ntime=None,
**kwargs):
''' load EnKF ensemble data as formatted GeoPy Dataset '''
if folder and not os.path.exists(folder): raise IOError(folder)
# default values
if isinstance(varlist, str) and varlist == 'default':
varlist = []
if station: varlist += ['discharge']
if well: varlist += ['head']
if varatts is None: varatts = variable_attributes.copy()
# figure out time axis
if date_range: start_date, end_date, sampling = date_range
time = timeAxis(start_date=start_date,
end_date=end_date,
sampling=sampling,
date_range=date_range,
time_axis=time_axis,
llastIncl=llastIncl,
ntime=ntime,
varatts=varatts)
ntime = len(time)
# load WSC station meta data
pass
# initialize Dataset
dataset = Dataset(name=name,
title=title if title else name.title(),
atts=metadata)
# load well data
if 'head' in varlist:
if not well: raise ArgumentError
if folder:
filepath = os.path.join(folder, well) # default output folder
else:
filepath = station
data = readKister(filepath=filepath,
period=(start_date, end_date),
resample=sampling,
lvalues=True)
assert ntime == len(data), data.shape
atts = varatts['head']
dataset += Variable(atts=atts, data=data, axes=(time, ))
# load discharge/hydrograph data
if 'discharge' in varlist:
if not station: raise ArgumentError
if folder:
filepath = os.path.join(folder, station) # default output folder
else:
filepath = station
data = readKister(filepath=filepath,
period=(start_date, end_date),
resample=sampling,
lvalues=True)
assert ntime == len(data), data.shape
atts = varatts['discharge']
if lkgs:
data *= 1000.
if atts['units'] == 'm^3/s': atts['units'] = 'kg/s'
dataset += Variable(atts=atts, data=data, axes=(time, ))
# return formatted Dataset
if scalefactors is not None and scalefactors != 1:
raise NotImplementedError
return dataset
def loadEnKF_StnTS(folder=None,
varlist='all',
varatts=None,
name='enkf',
title='EnKF',
basin=None,
start_date=None,
end_date=None,
sampling=None,
period=None,
date_range=None,
llastIncl=True,
WSC_station=None,
basin_list=None,
filenames=None,
prefix=None,
time_axis='datetime',
scalefactors=None,
metadata=None,
lkgs=False,
out_dir='out/',
yaml_file='../input_data/obs_meta.yaml',
lYAML=True,
nreal=None,
ntime=None,
**kwargs):
''' load EnKF ensemble data as formatted GeoPy Dataset '''
out_folder = os.path.join(folder, 'out/') # default output folder
if not os.path.exists(out_folder): raise IOError(out_folder)
# default values
if isinstance(varlist, str) and varlist == 'hydro':
varlist = Hydro.varlist
elif isinstance(varlist, str) and varlist == 'obs':
varlist = Obs.varlist
elif isinstance(varlist, str) and varlist == 'all':
varlist = Hydro.varlist + Obs.varlist
elif not isinstance(varlist, (tuple, list)):
raise TypeError(varlist)
if varatts is None: varatts = variable_attributes.copy()
varmap = {
varatt['name']: enkf_name
for enkf_name, varatt in list(varatts.items())
}
varlist = [varmap[var] for var in varlist]
# load WSC station meta data
pass
# initialize Dataset
dataset = Dataset(name=name,
title=title if title else name.title(),
atts=metadata)
ensemble = None
time = None
observation = None
# load observation/innovation data
if any([var in Obs.atts for var in varlist]):
# load data
vardata = loadObs(varlist=[var for var in varlist if var in Obs.atts],
folder=out_folder,
lpandas=False)
ntime, nobs, nreal = list(vardata.values())[0].shape
# create Axes
if time is None:
# figure out time axis
time = timeAxis(start_date=start_date,
end_date=end_date,
sampling=sampling,
date_range=date_range,
time_axis=time_axis,
llastIncl=llastIncl,
ntime=ntime,
varatts=varatts)
elif len(time) != ntime:
raise AxisError(time)
if ensemble is None:
# construct ensemble axis
ensemble = Axis(atts=varatts['ensemble'],
coord=np.arange(1, nreal + 1))
elif len(ensemble) != nreal:
raise AxisError(ensemble)
if observation is None:
# construct ensemble axis
observation = Axis(atts=varatts['observation'],
coord=np.arange(1, nobs + 1))
elif len(observation) != nobs:
raise AxisError(observation)
# create variables
for varname, data in list(vardata.items()):
dataset += Variable(atts=varatts[varname],
data=data,
axes=(time, observation, ensemble))
# load YAML data, if available
if lYAML:
# load YAML file
yaml_path = os.path.join(out_folder, yaml_file)
if not os.path.exists(yaml_path): raise IOError(yaml_path)
with open(yaml_path, 'r') as yf:
obs_meta = yaml.load(yf)
if obs_meta is None: raise IOError(yaml_path) # not a YAML file?
# constant create variables
for cvar, cval in list(obs_meta[0].items()):
if isinstance(cval, str): dtype, missing = np.string_, ''
elif isinstance(cval, (np.integer, int)):
dtype, missing = np.int_, 0
elif isinstance(cval, (np.inexact, float)):
dtype, missing = np.float_, np.NaN
else:
dtype = None # skip
if dtype:
data = np.asarray([
missing if obs[cvar] is None else obs[cvar]
for obs in obs_meta
],
dtype=dtype)
if cvar in varatts: atts = varatts[cvar]
else: atts = dict(name=cvar, units='')
dataset += Variable(atts=atts,
data=data,
axes=(observation, ))
elif ntime is None:
# try to infer time dimension from backup.info file
backup_info = os.path.join(folder, 'backup.info')
if os.path.exists(backup_info):
with open(backup_info, 'r') as bf:
ntime = int(bf.readline())
# load discharge/hydrograph data
if 'discharge' in varlist:
data = loadHydro(folder=out_folder, nreal=nreal, ntime=ntime)
ntime, nreal = data.shape
if time is None:
# figure out time axis
time = timeAxis(start_date=start_date,
end_date=end_date,
sampling=sampling,
date_range=date_range,
time_axis=time_axis,
llastIncl=llastIncl,
ntime=ntime,
varatts=varatts)
elif len(time) != ntime:
raise AxisError(time)
if ensemble is None:
# construct ensemble axis
ensemble = Axis(atts=varatts['ensemble'],
coord=np.arange(1, nreal + 1))
elif len(ensemble) != nreal:
raise AxisError(ensemble)
atts = varatts['discharge']
if lkgs:
data *= 1000.
if atts['units'] == 'm^3/s': atts['units'] = 'kg/s'
dataset += Variable(atts=atts, data=data, axes=(time, ensemble))
# return formatted Dataset
if scalefactors is not None and scalefactors != 1:
raise NotImplementedError
return dataset
def performExport(dataset, mode, dataargs, expargs, bcargs, loverwrite=False,
ldebug=False, lparallel=False, pidstr='', logger=None):
''' worker function to export ASCII rasters for a given dataset '''
# input checking
if not isinstance(dataset,basestring): raise TypeError
if not isinstance(dataargs,dict): raise TypeError # all dataset arguments are kwargs
# logging
if logger is None: # make new logger
logger = logging.getLogger() # new logger
logger.addHandler(logging.StreamHandler())
else:
if isinstance(logger,basestring):
logger = logging.getLogger(name=logger) # connect to existing one
elif not isinstance(logger,logging.Logger):
raise TypeError, 'Expected logger ID/handle in logger KW; got {}'.format(str(logger))
## extract meta data from arguments
dataargs, loadfct, srcage, datamsgstr = getMetaData(dataset, mode, dataargs, lone=False)
dataset_name = dataargs.dataset_name; periodstr = dataargs.periodstr; domain = dataargs.domain
# figure out bias correction parameters
if bcargs:
bcargs = bcargs.copy() # first copy, then modify...
bc_method = bcargs.pop('method',None)
if bc_method is None: raise ArgumentError("Need to specify bias-correction method to use bias correction!")
bc_obs = bcargs.pop('obs_dataset',None)
if bc_obs is None: raise ArgumentError("Need to specify observational dataset to use bias correction!")
bc_reference = bcargs.pop('reference',None)
if bc_reference is None: # infer from experiment name
if dataset_name[-5:] in ('-2050','-2100'): bc_reference = dataset_name[:-5] # cut of period indicator and hope for the best
else: bc_reference = dataset_name
bc_grid = bcargs.pop('grid',None)
if bc_grid is None: bc_grid = dataargs.grid
bc_domain = bcargs.pop('domain',None)
if bc_domain is None: bc_domain = domain
bc_varlist = bcargs.pop('varlist',None)
bc_varmap = bcargs.pop('varmap',None)
bc_tag = bcargs.pop('tag',None) # an optional name extension/tag
bc_pattern = bcargs.pop('file_pattern',None) # usually default in getPickleFile
lgzip = bcargs.pop('lgzip',None) # if pickle is gzipped (None: auto-detect based on file name extension)
# get name of pickle file (and folder)
picklefolder = dataargs.avgfolder.replace(dataset_name,bc_reference)
picklefile = getPickleFileName(method=bc_method, obs_name=bc_obs, gridstr=bc_grid, domain=bc_domain,
tag=bc_tag, pattern=bc_pattern)
picklepath = '{:s}/{:s}'.format(picklefolder,picklefile)
if lgzip:
picklepath += '.gz' # add extension
if not os.path.exists(picklepath): raise IOError(picklepath)
elif lgzip is None:
lgzip = False
if not os.path.exists(picklepath):
lgzip = True # assume gzipped file
picklepath += '.gz' # try with extension...
if not os.path.exists(picklepath): raise IOError(picklepath)
elif not os.path.exists(picklepath): raise IOError(picklepath)
pickleage = datetime.fromtimestamp(os.path.getmtime(picklepath))
# determine age of pickle file and compare against source age
else:
bc_method = False
pickleage = srcage
# parse export options
expargs = expargs.copy() # first copy, then modify...
lm3 = expargs.pop('lm3') # convert kg/m^2/s to m^3/m^2/s (water flux)
expformat = expargs.pop('format') # needed to get FileFormat object
exp_list= expargs.pop('exp_list') # this handled outside of export
compute_list = expargs.pop('compute_list', []) # variables to be (re-)computed - by default all
# initialize FileFormat class instance
fileFormat = getFileFormat(expformat, bc_method=bc_method, **expargs)
# get folder for target dataset and do some checks
expname = '{:s}_d{:02d}'.format(dataset_name,domain) if domain else dataset_name
expfolder = fileFormat.defineDataset(dataset=dataset, mode=mode, dataargs=dataargs, lwrite=True, ldebug=ldebug)
# prepare destination for new dataset
lskip = fileFormat.prepareDestination(srcage=max(srcage,pickleage), loverwrite=loverwrite)
# depending on last modification time of file or overwrite setting, start computation, or skip
if lskip:
# print message
skipmsg = "\n{:s} >>> Skipping: Format '{:s} for dataset '{:s}' already exists and is newer than source file.".format(pidstr,expformat,dataset_name)
skipmsg += "\n{:s} >>> ('{:s}')\n".format(pidstr,expfolder)
logger.info(skipmsg)
else:
## actually load datasets
source = loadfct() # load source data
# check period
if 'period' in source.atts and dataargs.periodstr != source.atts.period: # a NetCDF attribute
raise DateError, "Specifed period is inconsistent with netcdf records: '{:s}' != '{:s}'".format(periodstr,source.atts.period)
# load BiasCorrection object from pickle
if bc_method:
op = gzip.open if lgzip else open
with op(picklepath, 'r') as filehandle:
BC = pickle.load(filehandle)
# assemble logger entry
bcmsgstr = "(performing bias-correction using {:s} from {:s} towards {:s})".format(BC.long_name,bc_reference,bc_obs)
# print message
if mode == 'climatology': opmsgstr = 'Exporting Climatology ({:s}) to {:s} Format'.format(periodstr, expformat)
elif mode == 'time-series': opmsgstr = 'Exporting Time-series to {:s} Format'.format(expformat)
elif mode[-5:] == '-mean': opmsgstr = 'Exporting {:s}-Mean ({:s}) to {:s} Format'.format(mode[:-5], periodstr, expformat)
else: raise NotImplementedError, "Unrecognized Mode: '{:s}'".format(mode)
# print feedback to logger
logmsg = '\n{0:s} *** {1:^65s} *** \n{0:s} *** {2:^65s} *** \n'.format(pidstr,datamsgstr,opmsgstr)
if bc_method:
logmsg += "{0:s} *** {1:^65s} *** \n".format(pidstr,bcmsgstr)
logger.info(logmsg)
if not lparallel and ldebug: logger.info('\n'+str(source)+'\n')
# create GDAL-enabled target dataset
sink = Dataset(axes=(source.xlon,source.ylat), name=expname, title=source.title, atts=source.atts.copy())
addGDALtoDataset(dataset=sink, griddef=source.griddef)
assert sink.gdal, sink
# apply bias-correction
if bc_method:
source = BC.correct(source, asNC=False, varlist=bc_varlist, varmap=bc_varmap) # load bias-corrected variables into memory
# N.B.: for variables that are not bias-corrected, data are not loaded immediately but on demand; this way
# I/O and computing can be further disentangled and not all variables are always needed
# compute intermediate variables, if necessary
for varname in exp_list:
variables = None # variable list
var = None
# (re-)compute variable, if desired...
if varname in compute_list:
if varname == 'precip': var = newvars.computeTotalPrecip(source)
elif varname == 'waterflx': var = newvars.computeWaterFlux(source)
elif varname == 'liqwatflx': var = newvars.computeLiquidWaterFlux(source)
elif varname == 'netrad': var = newvars.computeNetRadiation(source, asVar=True)
elif varname == 'netrad_bb': var = newvars.computeNetRadiation(source, asVar=True, lrad=False, name='netrad_bb')
elif varname == 'netrad_bb0': var = newvars.computeNetRadiation(source, asVar=True, lrad=False, lA=False, name='netrad_bb0')
elif varname == 'vapdef': var = newvars.computeVaporDeficit(source)
elif varname in ('pet','pet_pm','petrad','petwnd') and 'pet' not in sink:
if 'petrad' in exp_list or 'petwnd' in exp_list:
variables = newvars.computePotEvapPM(source, lterms=True) # default; returns mutliple PET terms
else: var = newvars.computePotEvapPM(source, lterms=False) # returns only PET
elif varname == 'pet_th': var = None # skip for now
#var = computePotEvapTh(source) # simplified formula (less prerequisites)
# ... otherwise load from source file
if var is None and variables is None and varname in source:
var = source[varname].load() # load data (may not have to load all)
#else: raise VariableError, "Unsupported Variable '{:s}'.".format(varname)
# for now, skip variables that are None
if var or variables:
# handle lists as well
if var and variables: raise VariableError, (var,variables)
elif var: variables = (var,)
for var in variables:
addGDALtoVar(var=var, griddef=sink.griddef)
if not var.gdal and isinstance(fileFormat,ASCII_raster):
raise GDALError, "Exporting to ASCII_raster format requires GDAL-enabled variables."
# add to new dataset
sink += var
# convert units
if lm3:
for var in sink:
if var.units == 'kg/m^2/s':
var /= 1000. # divide to get m^3/m^2/s
var.units = 'm^3/m^2/s' # update units
# compute seasonal mean if we are in mean-mode
if mode[-5:] == '-mean':
sink = sink.seasonalMean(season=mode[:-5], lclim=True)
# N.B.: to remain consistent with other output modes,
# we need to prevent renaming of the time axis
sink = concatDatasets([sink,sink], axis='time', lensembleAxis=True)
sink.squeeze() # we need the year-axis until now to distinguish constant fields; now remove
# print dataset
if not lparallel and ldebug:
logger.info('\n'+str(sink)+'\n')
# export new dataset to selected format
fileFormat.exportDataset(sink)
# write results to file
writemsg = "\n{:s} >>> Export of Dataset '{:s}' to Format '{:s}' complete.".format(pidstr,expname, expformat)
writemsg += "\n{:s} >>> ('{:s}')\n".format(pidstr,expfolder)
logger.info(writemsg)
# clean up and return
source.unload(); #del source
return 0 # "exit code"
# imports
from glob import glob
from geodata.base import Dataset, Axis, Variable
from geodata.netcdf import writeNetCDF
# load list if well files and generate list of wells
well_files = glob(os.path.join(data_folder,'W*.xlsx'))
well_files.sort()
wells = [os.path.basename(name[:-5]) for name in well_files]
print(wells)
# dataset
time_ax = Axis(coord=np.arange(12*(period[1]-period[0]))+252, **varatts['time']) # origin: 1979-01
well_ax = Axis(coord=np.arange(len(wells))+1, name='well', units='')
dataset = Dataset(name=conservation_authority, title=conservation_authority+' Observation Wells')
# add meta data
meta_dicts = [loadMetadata(well, conservation_authority=conservation_authority) for well in wells]
for key in meta_dicts[0].keys():
if key in varatts: atts = varatts[key]
elif key.lower() in varatts: atts = varatts[key.lower()]
else: atts = dict(name=key, units='')
if atts['units']: data = np.asarray([wmd[key] for wmd in meta_dicts], dtype=np.float64)
else: data = np.asarray([wmd[key] for wmd in meta_dicts])
try:
dataset += Variable(data=data, axes=(well_ax,), **atts)
except:
pass
# add names
dataset += Variable(data=wells, axes=(well_ax,), name='well_name', units='',
atts=dict(long_name='Short Well Name'))
def performExport(dataset, mode, dataargs, expargs, loverwrite=False,
ldebug=False, lparallel=False, pidstr='', logger=None):
''' worker function to perform regridding for a given dataset and target grid '''
# input checking
if not isinstance(dataset,basestring): raise TypeError
if not isinstance(dataargs,dict): raise TypeError # all dataset arguments are kwargs
# logging
if logger is None: # make new logger
logger = logging.getLogger() # new logger
logger.addHandler(logging.StreamHandler())
else:
if isinstance(logger,basestring):
logger = logging.getLogger(name=logger) # connect to existing one
elif not isinstance(logger,logging.Logger):
raise TypeError, 'Expected logger ID/handle in logger KW; got {}'.format(str(logger))
## extract meta data from arguments
dataargs, loadfct, srcage, datamsgstr = getMetaData(dataset, mode, dataargs, lone=False)
dataset_name = dataargs.dataset_name; periodstr = dataargs.periodstr; domain = dataargs.domain
# parse export options
expargs = expargs.copy() # first copy, then modify...
lm3 = expargs.pop('lm3') # convert kg/m^2/s to m^3/m^2/s (water flux)
expformat = expargs.pop('format') # needed to get FileFormat object
varlist = expargs.pop('varlist') # this handled outside of export
# initialize FileFormat class instance
fileFormat = getFileFormat(expformat, **expargs)
# get folder for target dataset and do some checks
expname = '{:s}_d{:02d}'.format(dataset_name,domain) if domain else dataset_name
expfolder = fileFormat.defineDataset(dataset=dataset, mode=mode, dataargs=dataargs, lwrite=True, ldebug=ldebug)
# prepare destination for new dataset
lskip = fileFormat.prepareDestination(srcage=srcage, loverwrite=loverwrite)
# depending on last modification time of file or overwrite setting, start computation, or skip
if lskip:
# print message
skipmsg = "\n{:s} >>> Skipping: Format '{:s} for dataset '{:s}' already exists and is newer than source file.".format(pidstr,expformat,dataset_name)
skipmsg += "\n{:s} >>> ('{:s}')\n".format(pidstr,expfolder)
logger.info(skipmsg)
else:
## actually load datasets
source = loadfct() # load source data
# check period
if 'period' in source.atts and dataargs.periodstr != source.atts.period: # a NetCDF attribute
raise DateError, "Specifed period is inconsistent with netcdf records: '{:s}' != '{:s}'".format(periodstr,source.atts.period)
# print message
if mode == 'climatology': opmsgstr = 'Exporting Climatology ({:s}) to {:s} Format'.format(periodstr, expformat)
elif mode == 'time-series': opmsgstr = 'Exporting Time-series to {:s} Format'.format(expformat)
elif mode[-5:] == '-mean': opmsgstr = 'Exporting {:s}-Mean ({:s}) to {:s} Format'.format(mode[:-5], periodstr, expformat)
else: raise NotImplementedError, "Unrecognized Mode: '{:s}'".format(mode)
# print feedback to logger
logger.info('\n{0:s} *** {1:^65s} *** \n{0:s} *** {2:^65s} *** \n'.format(pidstr,datamsgstr,opmsgstr))
if not lparallel and ldebug: logger.info('\n'+str(source)+'\n')
# create GDAL-enabled target dataset
sink = Dataset(axes=(source.xlon,source.ylat), name=expname, title=source.title)
addGDALtoDataset(dataset=sink, griddef=source.griddef)
assert sink.gdal, sink
# N.B.: data are not loaded immediately but on demand; this way I/O and computing are further
# disentangled and not all variables are always needed
# Compute intermediate variables, if necessary
for varname in varlist:
vars = None # variable list
if varname in source:
var = source[varname].load() # load data (may not have to load all)
else:
var = None
if varname == 'waterflx': var = newvars.computeWaterFlux(source)
elif varname == 'liqwatflx': var = newvars.computeLiquidWaterFlux(source)
elif varname == 'netrad': var = newvars.computeNetRadiation(source, asVar=True)
elif varname == 'netrad_0': var = newvars.computeNetRadiation(source, asVar=True, lA=False, name='netrad_0')
elif varname == 'netrad_bb': var = newvars.computeNetRadiation(source, asVar=True, lrad=False, name='netrad_bb')
elif varname == 'vapdef': var = newvars.computeVaporDeficit(source)
elif varname == 'pet' or varname == 'pet_pm':
if 'petrad' in varlist or 'petwnd' in varlist:
vars = newvars.computePotEvapPM(source, lterms=True) # default; returns mutliple PET terms
else: var = newvars.computePotEvapPM(source, lterms=False) # returns only PET
elif varname == 'pet_th': var = None # skip for now
#var = computePotEvapTh(source) # simplified formula (less prerequisites)
else: raise VariableError, "Unsupported Variable '{:s}'.".format(varname)
# for now, skip variables that are None
if var or vars:
# handle lists as well
if var and vars: raise VariableError, (var,vars)
elif var: vars = (var,)
for var in vars:
addGDALtoVar(var=var, griddef=sink.griddef)
if not var.gdal and isinstance(fileFormat,ASCII_raster):
raise GDALError, "Exporting to ASCII_raster format requires GDAL-enabled variables."
# add to new dataset
sink += var
# convert units
if lm3:
for var in sink:
if var.units == 'kg/m^2/s':
var /= 1000. # divide to get m^3/m^2/s
var.units = 'm^3/m^2/s' # update units
# compute seasonal mean if we are in mean-mode
if mode[-5:] == '-mean':
sink = sink.seasonalMean(season=mode[:-5], taxatts=dict(name='time'))
# N.B.: to remain consistent with other output modes,
# we need to prevent renaming of the time axis
# print dataset
if not lparallel and ldebug:
logger.info('\n'+str(sink)+'\n')
# export new dataset to selected format
fileFormat.exportDataset(sink)
# write results to file
writemsg = "\n{:s} >>> Export of Dataset '{:s}' to Format '{:s}' complete.".format(pidstr,expname, expformat)
writemsg += "\n{:s} >>> ('{:s}')\n".format(pidstr,expfolder)
logger.info(writemsg)
# clean up and return
source.unload(); #del source
return 0 # "exit code"
def rasterDataset(name=None,
title=None,
vardefs=None,
axdefs=None,
atts=None,
projection=None,
griddef=None,
lgzip=None,
lgdal=True,
lmask=True,
fillValue=None,
lskipMissing=True,
lgeolocator=True,
file_pattern=None,
lfeedback=True,
**kwargs):
''' function to load a set of variables that are stored in raster format in a systematic directory tree into a Dataset
Variables and Axis are defined as follows:
vardefs[varname] = dict(name=string, units=string, axes=tuple of strings, atts=dict, plot=dict, dtype=np.dtype, fillValue=value)
axdefs[axname] = dict(name=string, units=string, atts=dict, coord=array or list) or None
The path to raster files is constructed as variable_pattern+axes_pattern, where axes_pattern is defined through the axes,
(as in rasterVarialbe) and variable_pattern takes the special keywords VAR, which is the variable key in vardefs.
'''
## prepare input data and axes
if griddef:
xlon, ylat = griddef.xlon, griddef.ylat
if projection is None:
projection = griddef.projection
elif projection != griddef.projection:
raise ArgumentError("Conflicting projection and GridDef!")
geotransform = griddef.geotransform
isProjected = griddef.isProjected
else:
xlon = ylat = geotransform = None
isProjected = False if projection is None else True
# construct axes dict
axes = dict()
for axname, axdef in axdefs.items():
assert 'coord' in axdef, axdef
assert ('name' in axdef and 'units' in axdef) or 'atts' in axdef, axdef
if axdef is None:
axes[axname] = None
else:
ax = Axis(**axdef)
axes[ax.name] = ax
# check for map Axis
if isProjected:
if 'x' not in axes: axes['x'] = xlon
if 'y' not in axes: axes['y'] = ylat
else:
if 'lon' not in axes: axes['lon'] = xlon
if 'lat' not in axes: axes['lat'] = ylat
## load raster data into Variable objects
varlist = []
for varname, vardef in vardefs.items():
# check definitions
assert 'axes' in vardef and 'dtype' in vardef, vardef
assert ('name' in vardef
and 'units' in vardef) or 'atts' in vardef, vardef
# determine relevant axes
vardef = vardef.copy()
axes_list = [
None if ax is None else axes[ax] for ax in vardef.pop('axes')
]
# define path parameters (with varname)
path_params = vardef.pop('path_params', None)
path_params = dict() if path_params is None else path_params.copy()
if 'VAR' not in path_params:
path_params['VAR'] = varname # a special key
# add kwargs and relevant axis indices
relaxes = [ax.name for ax in axes_list
if ax is not None] # relevant axes
for key, value in kwargs.items():
if key not in axes or key in relaxes:
vardef[key] = value
# create Variable object
var = rasterVariable(projection=projection,
griddef=griddef,
file_pattern=file_pattern,
lgzip=lgzip,
lgdal=lgdal,
lmask=lmask,
lskipMissing=lskipMissing,
axes=axes_list,
path_params=path_params,
lfeedback=lfeedback,
**vardef)
# vardef components: name, units, atts, plot, dtype, fillValue
varlist.append(var)
# check that map axes are correct
for ax in var.xlon, var.ylat:
if axes[ax.name] is None: axes[ax.name] = ax
elif axes[ax.name] != ax:
raise AxisError("{} axes are incompatible.".format(ax.name))
if griddef is None: griddef = var.griddef
elif griddef != var.griddef:
raise AxisError("GridDefs are inconsistent.")
if geotransform is None: geotransform = var.geotransform
elif geotransform != var.geotransform:
raise AxisError(
"Conflicting geotransform (from Variable) and GridDef!\n {} != {}"
.format(var.geotransform, geotransform))
## create Dataset
# create dataset
dataset = Dataset(name=name,
title=title,
varlist=varlist,
axes=axes,
atts=atts)
# add GDAL functionality
dataset = addGDALtoDataset(dataset,
griddef=griddef,
projection=projection,
geotransform=geotransform,
gridfolder=None,
lwrap360=None,
geolocator=lgeolocator,
lforce=False)
# N.B.: for some reason we also need to pass the geotransform, otherwise it is recomputed internally and some consistency
# checks fail due to machine-precision differences
# return GDAL-enabled Dataset
return dataset
def Regrid(self, griddef=None, projection=None, geotransform=None, size=None, xlon=None, ylat=None,
lmask=True, int_interp=None, float_interp=None, **kwargs):
''' Setup climatology and start computation; calls processClimatology. '''
# make temporary gdal dataset
if self.source is self.target:
if self.tmp: assert self.source == self.tmpput and self.target == self.tmpput
# the operation can not be performed "in-place"!
self.target = Dataset(name='tmptoo', title='Temporary target dataset for non-in-place operations', varlist=[], atts={})
ltmptoo = True
else: ltmptoo = False
# make sure the target dataset is a GDAL-enabled dataset
if 'gdal' in self.target.__dict__:
# gdal info alread present
if griddef is not None or projection is not None or geotransform is not None:
raise AttributeError, "Target Dataset '%s' is already GDAL enabled - cannot overwrite settings!"%self.target.name
if self.target.xlon is None: raise GDALError, "Map axis 'xlon' not found!"
if self.target.ylat is None: raise GDALError, "Map axis 'ylat' not found!"
xlon = self.target.xlon; ylat = self.target.ylat
else:
# need to set GDAL parameters
if self.tmp and 'gdal' in self.output.__dict__:
# transfer gdal settings from output to temporary dataset
assert self.target is not self.output
projection = self.output.projection; geotransform = self.output.geotransform
xlon = self.output.xlon; ylat = self.output.ylat
else:
# figure out grid definition from input
if griddef is None:
griddef = GridDefinition(projection=projection, geotransform=geotransform, size=size, xlon=xlon, ylat=ylat)
# pass arguments through GridDefinition, if not provided
projection=griddef.projection; geotransform=griddef.geotransform
xlon=griddef.xlon; ylat=griddef.ylat
# apply GDAL settings target dataset
for ax in (xlon,ylat): self.target.addAxis(ax, loverwrite=True) # i.e. replace if already present
self.target = addGDALtoDataset(self.target, projection=projection, geotransform=geotransform)
# use these map axes
xlon = self.target.xlon; ylat = self.target.ylat
assert isinstance(xlon,Axis) and isinstance(ylat,Axis)
# determine source dataset grid definition
if self.source.griddef is None:
srcgrd = GridDefinition(projection=self.source.projection, geotransform=self.source.geotransform,
size=self.source.mapSize, xlon=self.source.xlon, ylat=self.source.ylat)
else: srcgrd = self.source.griddef
srcres = srcgrd.scale; tgtres = griddef.scale
# determine if shift is necessary to insure correct wrapping
if not srcgrd.isProjected and not griddef.isProjected:
lwrapSrc = srcgrd.wrap360
lwrapTgt = griddef.wrap360
# check grids
for grd in (srcgrd,griddef):
if grd.wrap360:
assert grd.geotransform[0] + grd.geotransform[1]*(len(grd.xlon)-1) > 180
assert np.round(grd.geotransform[1]*len(grd.xlon), decimals=2) == 360 # require 360 deg. to some accuracy...
assert any( grd.xlon.getArray() > 180 ) # need to wrap around
assert all( grd.xlon.getArray() >= 0 )
assert all( grd.xlon.getArray() <= 360 )
else:
assert grd.geotransform[0] + grd.geotransform[1]*(len(grd.xlon)-1) < 180
assert all( grd.xlon.getArray() >= -180 )
assert all( grd.xlon.getArray() <= 180 )
else:
lwrapSrc = False # no need to shift, if a projected grid is involved!
lwrapTgt = False # no need to shift, if a projected grid is involved!
# determine GDAL interpolation
if int_interp is None: int_interp = gdalInterp('nearest')
else: int_interp = gdalInterp(int_interp)
if float_interp is None:
if srcres < tgtres: float_interp = gdalInterp('convolution') # down-sampling: 'convolution'
else: float_interp = gdalInterp('cubicspline') # up-sampling
else: float_interp = gdalInterp(float_interp)
# prepare function call
function = functools.partial(self.processRegrid, ylat=ylat, xlon=xlon, lwrapSrc=lwrapSrc, lwrapTgt=lwrapTgt, # already set parameters
lmask=lmask, int_interp=int_interp, float_interp=float_interp)
# start process
if self.feedback: print('\n +++ processing regridding +++ ')
self.process(function, **kwargs) # currently 'flush' is the only kwarg
# now make sure we have a GDAL dataset!
self.target = addGDALtoDataset(self.target, griddef=griddef)
if self.feedback: print('\n')
if self.tmp: self.tmpput = self.target
if ltmptoo: assert self.tmpput.name == 'tmptoo' # set above, when temp. dataset is created
def Extract(self, template=None, stnax=None, xlon=None, ylat=None, laltcorr=True, **kwargs):
''' Extract station data points from gridded datasets; calls processExtract.
A station dataset can be passed as template (must have station coordinates. '''
if not self.source.gdal: raise DatasetError, "Source dataset must be GDAL enabled! {:s} is not.".format(self.source.name)
if template is None: raise NotImplementedError
elif isinstance(template, Dataset):
if not template.hasAxis('station'): raise DatasetError, "Template station dataset needs to have a station axis."
if not ( (template.hasVariable('lat') or template.hasVariable('stn_lat')) and
(template.hasVariable('lon') or template.hasVariable('stn_lon')) ):
raise DatasetError, "Template station dataset needs to have lat/lon arrays for the stations."
else: raise TypeError
# make temporary dataset
if self.source is self.target:
if self.tmp: assert self.source == self.tmpput and self.target == self.tmpput
# the operation can not be performed "in-place"!
self.target = Dataset(name='tmptoo', title='Temporary target dataset for non-in-place operations', varlist=[], atts={})
ltmptoo = True
else: ltmptoo = False
src = self.source; tgt = self.target # short-cuts
# determine source dataset grid definition
if src.griddef is None:
srcgrd = GridDefinition(projection=self.source.projection, geotransform=self.source.geotransform,
size=self.source.mapSize, xlon=self.source.xlon, ylat=self.source.ylat)
else: srcgrd = src.griddef
# figure out horizontal axes (will be replaced with station axis)
if isinstance(xlon,Axis):
if not src.hasAxis(xlon, check=True): raise DatasetError
elif isinstance(xlon,basestring): xlon = src.getAxis(xlon)
else: xlon = src.x if srcgrd.isProjected else src.lon
if isinstance(ylat,Axis):
if not src.hasAxis(ylat, check=True): raise DatasetError
elif isinstance(ylat,basestring): ylat = src.getAxis(ylat)
else: ylat = src.y if srcgrd.isProjected else src.lat
if stnax: # not in source dataset!
if src.hasAxis(stnax, check=True): raise DatasetError, "Source dataset must not have a 'station' axis!"
elif template: stnax = template.station # station axis
else: raise ArgumentError, "A station axis needs to be supplied."
assert isinstance(xlon,Axis) and isinstance(ylat,Axis) and isinstance(stnax,Axis)
# transform to dataset-native coordinate system
if template:
if template.hasVariable('lat'): lats = template.lat.getArray()
else: lats = template.stn_lat.getArray()
if template.hasVariable('lon'): lons = template.lon.getArray()
else: lons = template.stn_lon.getArray()
else: raise NotImplementedError, "Cannot extract station data without a station template Dataset"
# adjust longitudes
if srcgrd.isProjected:
if lons.max() > 180.: lons = np.where(lons > 180., 360.-lons, lons)
# reproject coordinate
latlon = osr.SpatialReference()
latlon.SetWellKnownGeogCS('WGS84') # a normal lat/lon coordinate system
tx = osr.CoordinateTransformation(latlon,srcgrd.projection)
xs = []; ys = []
for i in xrange(len(lons)):
x,y,z = tx.TransformPoint(lons[i].astype(np.float64),lats[i].astype(np.float64))
xs.append(x); ys.append(y); del z
lons = np.array(xs); lats = np.array(ys)
#lons,lats = tx.TransformPoints(lons,lats) # doesn't seem to work...
else:
if lons.min() < 0. and xlon.coord.max() > 180.: lons = np.where(lons < 0., lons + 360., lons)
elif lons.max() > 180. and xlon.coord.min() < 0.: lons = np.where(lons > 180., 360.-lons, lons)
else: pass # source and template do not conflict
# generate index list
ixlon = []; iylat = []; istn = []; zs_err = [] # also record elevation error
lzs = src.hasVariable('zs')
lstnzs = template.hasVariable('zs') or template.hasVariable('stn_zs')
if laltcorr and lzs and lstnzs:
if src.zs.ndim > 2: src.zs = src.zs(time=0, lidx=True) # first time-slice (for CESM)
if src.zs.ndim != 2 or not src.gdal or src.zs.units != 'm': raise VariableError
# consider altidue of surrounding points as well
zs = src.zs.getArray(unmask=True,fillValue=-300)
if template.hasVariable('zs'): stn_zs = template.zs.getArray(unmask=True,fillValue=-300)
else: stn_zs = template.stn_zs.getArray(unmask=True,fillValue=-300)
if src.zs.axisIndex(xlon.name) == 0: zs.transpose() # assuming lat,lon or y,x order is more common
ye,xe = zs.shape # assuming order lat,lon or y,x
xe -= 1; ye -= 1 # last valid index, not length
for n,lon,lat in zip(xrange(len(stnax)),lons,lats):
ip = xlon.getIndex(lon, mode='left', outOfBounds=True)
jp = ylat.getIndex(lat, mode='left', outOfBounds=True)
if ip is not None and jp is not None:
# find neighboring point with smallest altitude error
# ip = im+1 if im < xe else im
# jp = jm+1 if jm < ye else jm
im = ip-1 if ip > 0 else ip
jm = jp-1 if jp > 0 else jp
zdiff = np.Infinity # initialize, so that it triggers at least once
# check four closest grid points
for i in im,ip:
for j in jm,jp:
ze = zs[j,i]-stn_zs[n]
zd = np.abs(ze) # compute elevation error
if zd < zdiff: ii,jj,zdiff,zerr = i,j,zd,ze # preliminary selection, triggers at least once
ixlon.append(ii); iylat.append(jj); istn.append(n); zs_err.append(zerr) # final selection
else:
# just choose horizontally closest point
for n,lon,lat in zip(xrange(len(stnax)),lons,lats):
i = xlon.getIndex(lon, mode='closest', outOfBounds=True)
j = ylat.getIndex(lat, mode='closest', outOfBounds=True)
if i is not None and j is not None:
if lzs: # compute elevation error
zs_err.append(zs[j,i]-stn_zs[n])
ixlon.append(i); iylat.append(j); istn.append(n)
# N.B.: it is necessary to append, because we don't know the number of valid points
ixlon = np.array(ixlon); iylat = np.array(iylat); istn = np.array(istn); zs_err = np.array(zs_err)
# prepare target dataset
# N.B.: attributes should already be set in target dataset (by caller module)
# we are also assuming the new dataset has no axes yet
assert len(tgt.axes) == 0
# add axes from source data
for axname,ax in src.axes.iteritems():
if axname not in (xlon.name,ylat.name):
tgt.addAxis(ax, asNC=True, copy=True)
# add station axis (trim to valid coordinates)
newstnax = stnax.copy(coord=stnax.coord[istn]) # same but with trimmed coordinate array
tgt.addAxis(newstnax, asNC=True, copy=True) # already new copy
# create variable for elevation error
if lzs:
assert len(zs_err) > 0
zs_err = Variable(name='zs_err', units='m', data=zs_err, axes=(newstnax,),
atts=dict(long_name='Station Elevation Error'))
tgt.addVariable(zs_err, asNC=True, copy=True); del zs_err # need to copy to make NC var
# add a bunch of other variables with station meta data
for var in template.variables.itervalues():
if var.ndim == 1 and var.hasAxis(stnax): # station attributes
if var.name[-4:] != '_len' or var.name == 'stn_rec_len': # exclude certain attributes
newvar = var.copy(data=var.getArray()[istn], axes=(newstnax,))
if newvar.name[:4] != 'stn_' and newvar.name[:8] != 'station_' and newvar.name[:8] != 'cluster_':
newvar.name = 'stn_'+newvar.name # copy cluster_* as they are!
# N.B.: we need to rename, or name collisions will happen!
tgt.addVariable(newvar, asNC=True, copy=True); del newvar # need to copy to make NC var
# save all the meta data
tgt.sync()
# prepare function call
function = functools.partial(self.processExtract, ixlon=ixlon, iylat=iylat, ylat=ylat, xlon=xlon, stnax=stnax) # already set parameters
# start process
if self.feedback: print('\n +++ processing point-data extraction +++ ')
self.process(function, **kwargs) # currently 'flush' is the only kwarg
if self.feedback: print('\n')
if self.tmp: self.tmpput = self.target
if ltmptoo: assert self.tmpput.name == 'tmptoo' # set above, when temp. dataset is created
class CentralProcessingUnit(object):
def __init__(self, source, target=None, varlist=None, ignorelist=None, tmp=True, feedback=True):
''' Initialize processor and pass input and output datasets. '''
# check varlist
if varlist is None: varlist = source.variables.keys() # all source variables
elif not isinstance(varlist,(list,tuple)): raise TypeError
self.varlist = varlist # list of variable to be processed
# ignore list (e.g. variables that will cause errors)
if ignorelist is None: ignorelist = [] # an empty list
elif not isinstance(ignorelist,(list,tuple)): raise TypeError
self.ignorelist = ignorelist # list of variable *not* to be processed
# check input
if not isinstance(source,Dataset): raise TypeError
if isinstance(source,DatasetNetCDF) and not 'r' in source.mode: raise PermissionError
self.input = source
self.source = source
# check output
if target is not None:
if not isinstance(target,Dataset): raise TypeError
if isinstance(target,DatasetNetCDF) and not 'w' in target.mode: raise PermissionError
else:
if not tmp: raise DatasetError, "Need target location, if temporary storage is disables (tmp=False)."
self.output = target
# temporary dataset
self.tmp = tmp
if tmp: self.tmpput = Dataset(name='tmp', title='Temporary Dataset', varlist=[], atts={})
else: self.tmpput = None
# determine if temporary storage is used and assign target dataset
if self.tmp: self.target = self.tmpput
else: self.target = self.output
# whether or not to print status output
self.feedback = feedback
def getTmp(self, asNC=False, filename=None, deepcopy=False, **kwargs):
''' Get a copy of the temporary data in dataset format. '''
if not self.tmp: raise DatasetError
# make new dataset (name and title should transfer in atts dict)
if asNC:
if not isinstance(filename,basestring): raise TypeError
writeData = kwargs.pop('writeData',False)
ncformat = kwargs.pop('ncformat','NETCDF4')
zlib = kwargs.pop('zlib',True)
dataset = asDatasetNC(self.tmpput, ncfile=filename, mode='wr', deepcopy=deepcopy,
writeData=writeData, ncformat=ncformat, zlib=zlib, **kwargs)
else:
dataset = self.tmpput.copy(varsdeep=deepcopy, atts=self.input.atts.copy(), **kwargs)
# return dataset
return dataset
def sync(self, varlist=None, flush=False, gdal=True, copydata=True):
''' Transfer contents of temporary storage to output/target dataset. '''
if not isinstance(self.output,Dataset): raise DatasetError, "Cannot sync without target Dataset!"
if self.tmp:
if varlist is None: varlist = self.tmpput.variables.keys()
for varname in varlist:
if varname in self.tmpput.variables:
var = self.tmpput.variables[varname]
self.output.addVariable(var, loverwrite=True, deepcopy=copydata)
# N.B.: without copydata/deepcopy, only the variable header is created but no data is written
if flush: var.unload() # remove unnecessary references (unlink data)
if gdal and 'gdal' in self.tmpput.__dict__:
if self.tmpput.gdal:
projection = self.tmpput.projection; geotransform = self.tmpput.geotransform
#xlon = self.tmpput.xlon; ylat = self.tmpput.ylat
else:
projection=None; geotransform=None; #xlon = None; ylat = None
self.output = addGDALtoDataset(self.output, projection=projection, geotransform=geotransform)
# self.source = self.output # future operations will write to the output dataset directly
# self.target = self.output # future operations will write to the output dataset directly
def writeNetCDF(self, filename=None, folder=None, ncformat='NETCDF4', zlib=True, writeData=True, close=False, flush=False):
''' Write current temporary storage to a NetCDF file. '''
if self.tmp:
if not isinstance(filename,basestring): raise TypeError
if folder is not None: filename = folder + filename
output = writeNetCDF(self.tmpput, filename, ncformat=ncformat, zlib=zlib, writeData=writeData, close=False)
if flush: self.tmpput.unload()
if self.feedback: print('\nOutput written to {0:s}\n'.format(filename))
else:
self.output.sync()
output = self.output.dataset # get (primary) NetCDF file
if self.feedback: print('\nSynchronized dataset {0:s} with temporary storage.\n'.format(output.name))
# flush?
if flush: self.output.unload()
# close file or return file handle
if close: output.close()
else: return output
def process(self, function, flush=False):
''' This method applies the desired operation/function to each variable in varlist. '''
if flush: # this function is to save RAM by flushing results to disk immediately
if not isinstance(self.output,DatasetNetCDF):
raise ProcessError, "Flush can only be used with NetCDF Datasets (and not with temporary storage)."
if self.tmp: # flush requires output to be target
if self.source.gdal and not self.tmpput.gdal:
self.tmpput = addGDALtoDataset(self.tmpput, projection=self.source.projection, geotransform=self.source.geotransform)
self.source = self.tmpput
self.target = self.output
self.tmp = False # not using temporary storage anymore
# loop over input variables
for varname in self.varlist:
# check agaisnt ignore list
if varname not in self.ignorelist:
# check if variable already exists
if self.target.hasVariable(varname):
# "in-place" operations
var = self.target.variables[varname]
newvar = function(var) # perform actual processing
if newvar.ndim != var.ndim or newvar.shape != var.shape: raise VariableError
if newvar is not var: self.target.replaceVariable(var,newvar)
elif self.source.hasVariable(varname):
var = self.source.variables[varname]
ldata = var.data # whether data was pre-loaded
# perform operation from source and copy results to target
newvar = function(var) # perform actual processing
if not ldata: var.unload() # if it was already loaded, don't unload
self.target.addVariable(newvar, copy=True) # copy=True allows recasting as, e.g., a NC variable
else:
raise DatasetError, "Variable '%s' not found in input dataset."%varname
assert varname == newvar.name
# flush data to disk immediately
if flush:
self.output.variables[varname].unload() # again, free memory
newvar.unload(); del var, newvar # free space; already added to new dataset
# after everything is said and done:
self.source = self.target # set target to source for next time
## functions (or function pairs, rather) that perform operations on the data
# every function pair needs to have a setup function and a processing function
# the former sets up the target dataset and the latter operates on the variables
# function pair to average data over a given collection of shapes
def ShapeAverage(self, shape_dict=None, shape_name=None, shpax=None, xlon=None, ylat=None,
memory=500, **kwargs):
''' Average over a limited area of a gridded datasets; calls processAverageShape.
A dictionary of NamedShape objects is expected to define the averaging areas.
'memory' controls the garbage collection interval and approximately corresponds
to MB in temporary (it does not include loading the variable into RAM, though). '''
if not self.source.gdal: raise DatasetError, "Source dataset must be GDAL enabled! {:s} is not.".format(self.source.name)
if not isinstance(shape_dict,OrderedDict): raise TypeError
if not all(isinstance(shape,NamedShape) for shape in shape_dict.itervalues()): raise TypeError
# make temporary dataset
if self.source is self.target:
if self.tmp: assert self.source == self.tmpput and self.target == self.tmpput
# the operation can not be performed "in-place"!
self.target = Dataset(name='tmptoo', title='Temporary target dataset for non-in-place operations', varlist=[], atts={})
ltmptoo = True
else: ltmptoo = False
src = self.source; tgt = self.target # short-cuts
# determine source dataset grid definition
if src.griddef is None:
srcgrd = GridDefinition(projection=self.source.projection, geotransform=self.source.geotransform,
size=self.source.mapSize, xlon=self.source.xlon, ylat=self.source.ylat)
else: srcgrd = src.griddef
# figure out horizontal axes (will be replaced with station axis)
if isinstance(xlon,Axis):
if not src.hasAxis(xlon, check=True): raise DatasetError
elif isinstance(xlon,basestring): xlon = src.getAxis(xlon)
else: xlon = src.x if srcgrd.isProjected else src.lon
if isinstance(ylat,Axis):
if not src.hasAxis(ylat, check=True): raise DatasetError
elif isinstance(ylat,basestring): ylat = src.getAxis(ylat)
else: ylat = src.y if srcgrd.isProjected else src.lat
# check/create shapes axis
if shpax: # not in source dataset!
# if shape axis supplied
if src.hasAxis(shpax, check=True): raise DatasetError, "Source dataset must not have a 'shape' axis!"
if len(shpax) != len(shape_dict): raise AxisError
else:
# creat shape axis, if not supplied
shpatts = dict(name='shape', long_name='Ordinal Number of Shape', units='#')
shpax = Axis(coord=np.arange(1,len(shape_dict)+1), atts=shpatts) # starting at 1
assert isinstance(xlon,Axis) and isinstance(ylat,Axis) and isinstance(shpax,Axis)
# prepare target dataset
# N.B.: attributes should already be set in target dataset (by caller module)
# we are also assuming the new dataset has no axes yet
assert len(tgt.axes) == 0
# add station axis (trim to valid coordinates)
tgt.addAxis(shpax, asNC=True, copy=True) # already new copy
# add axes from source data
for axname,ax in src.axes.iteritems():
if axname not in (xlon.name,ylat.name):
tgt.addAxis(ax, asNC=True, copy=True)
# add shape names
shape_names = [shape.name for shape in shape_dict.itervalues()] # can construct Variable from list!
atts = dict(name='shape_name', long_name='Name of Shape', units='')
tgt.addVariable(Variable(data=shape_names, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add proper names
shape_long_names = [shape.long_name for shape in shape_dict.itervalues()] # can construct Variable from list!
atts = dict(name='shp_long_name', long_name='Proper Name of Shape', units='')
tgt.addVariable(Variable(data=shape_long_names, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add shape category
shape_type = [shape.shapetype for shape in shape_dict.itervalues()] # can construct Variable from list!
atts = dict(name='shp_type', long_name='Type of Shape', units='')
tgt.addVariable(Variable(data=shape_type, axes=(shpax,), atts=atts), asNC=True, copy=True)
# collect rasterized masks from shape files
mask_array = np.zeros((len(shpax),)+srcgrd.size[::-1], dtype=np.bool)
# N.B.: rasterize() returns mask in (y,x) shape, size is ordered as (x,y)
shape_masks = []; shp_full = []; shp_empty = []; shp_encl = []
for i,shape in enumerate(shape_dict.itervalues()):
mask = shape.rasterize(griddef=srcgrd, asVar=False)
mask_array[i,:] = mask
masksum = mask.sum()
lfull = masksum == 0; shp_full.append( lfull )
lempty = masksum == mask.size; shp_empty.append( lempty )
shape_masks.append( mask if not lempty else None )
if lempty: shp_encl.append( False )
else:
shp_encl.append( np.all( mask[[0,-1],:] == True ) and np.all( mask[:,[0,-1]] == True ) )
# i.e. if boundaries are masked
# N.B.: shapes that have no overlap with grid will be skipped and filled with NaN
# add rasterized masks to new dataset
atts = dict(name='shp_mask', long_name='Rasterized Shape Mask', units='')
tgt.addVariable(Variable(data=mask_array, atts=atts, axes=(shpax,srcgrd.ylat.copy(),srcgrd.xlon.copy())),
asNC=True, copy=True)
# add area enclosed by shape
da = srcgrd.geotransform[1]*srcgrd.geotransform[5]
mask_area = (1-mask_array).mean(axis=2).mean(axis=1)*da
atts = dict(name='shp_area', long_name='Area Contained in the Shape',
units= 'm^2' if srcgrd.isProjected else 'deg^2' )
tgt.addVariable(Variable(data=mask_area, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add flag to indicate if shape is fully enclosed by domain
atts = dict(name='shp_encl', long_name='If Shape is fully included in Domain', units= '')
tgt.addVariable(Variable(data=shp_encl, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add flag to indicate if shape fully covers domain
atts = dict(name='shp_full', long_name='If Shape fully covers Domain', units= '')
tgt.addVariable(Variable(data=shp_full, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add flag to indicate if shape and domain have no overlap
atts = dict(name='shp_empty', long_name='If Shape and Domain have no Overlap', units= '')
tgt.addVariable(Variable(data=shp_empty, axes=(shpax,), atts=atts), asNC=True, copy=True)
# save all the meta data
tgt.sync()
# prepare function call
function = functools.partial(self.processShapeAverage, masks=shape_masks, ylat=ylat, xlon=xlon,
shpax=shpax, memory=memory) # already set parameters
# start process
if self.feedback: print('\n +++ processing shape/area averaging +++ ')
self.process(function, **kwargs) # currently 'flush' is the only kwarg
if self.feedback: print('\n')
if self.tmp: self.tmpput = self.target
if ltmptoo: assert self.tmpput.name == 'tmptoo' # set above, when temp. dataset is created
# the previous method sets up the process, the next method performs the computation
def processShapeAverage(self, var, masks=None, ylat=None, xlon=None, shpax=None, memory=500):
''' Compute masked area averages from variable data. 'memory' controls the garbage collection
interval approximately corresponds to MB in RAM.'''
# process gdal variables (if a variable has a horiontal grid, it should be GDAL enabled)
if var.gdal and ( np.issubdtype(var.dtype,np.integer) or np.issubdtype(var.dtype,np.inexact) ):
if self.feedback: print('\n'+var.name),
assert var.hasAxis(xlon) and var.hasAxis(ylat)
assert len(masks) == len(shpax)
tgt = self.target
assert tgt.hasAxis(shpax, strict=False) and shpax not in var.axes
# assemble new axes
axes = [tgt.getAxis(shpax.name)]
for ax in var.axes:
if ax not in (xlon,ylat) and ax.name != shpax.name: # these axes are just transferred
axes.append(tgt.getAxis(ax.name))
# N.B.: shape axis well be outer axis
axes = tuple(axes)
# pre-allocate
shape = tuple(len(ax) for ax in axes)
tgtdata = np.zeros(shape, dtype=np.float32)
# now we loop over all shapes/masks
if self.feedback:
varname = var.name
print '\n ... loading ',varname
var.load()
if self.feedback:
varname = var.name
print '\n ... averaging ',varname
## compute shape averages for each time step
# Basically, for each shape the entire array is masked, using the shape and the broadcasting
# functionality for horizontal masks of the Variable class (which creates a lot of overhead);
# then the masked average is taken and the process is repeated for the next shape/mask.
# Using mapMean creates a lot of overhead and there are probably more efficient ways to do it.
if var.ndim == 2:
for i,mask in enumerate(masks):
if mask is None: tgtdata[i] = np.NaN # NaN for missing values (i.e. no overlap)
else: tgtdata[i] = var.mapMean(mask=mask, asVar=False, squeeze=True) # compute the averages
if self.feedback: print varname, i
# garbage collection is typically not necessary for 2D fields
elif var.ndim > 2:
cnt = float(0.); inc = float(var.data_array.nbytes) / (1024.*1024.) # counter/increment to estimate memory useage (in MB)
for i,mask in enumerate(masks):
if mask is None: tgtdata[i,:] = np.NaN # NaN for missing values (i.e. no overlap)
else:
tgtdata[i,:] = var.mapMean(mask=mask, asVar=False, squeeze=True).filled(np.NaN) # mapMean returns a masked array
# N.B.: this is necessary, because sometimes shapes only contain invalid values
cnt += inc # keep track of memory
# N.B.: mapMean creates a lot of temporary arrays that don't get garbage-collected
if self.feedback: print varname, i, cnt
if cnt > memory:
cnt = 0 # reset counter
gc.collect() # collect garbage in certain itnervals
if self.feedback: print 'garbage collected'
else: raise AxisError
# create new Variable
assert shape == tgtdata.shape
newvar = var.copy(axes=axes, data=tgtdata) # new axes and data
del tgtdata, mask # clean up (just to make sure)
gc.collect() # clean
else:
var.load() # need to load variables into memory to copy it (and we are not doing anything else...)
newvar = var # just pass over the variable to the new dataset
# return variable
return newvar
# function pair to extract station data from a time-series (or climatology)
def Extract(self, template=None, stnax=None, xlon=None, ylat=None, laltcorr=True, **kwargs):
''' Extract station data points from gridded datasets; calls processExtract.
A station dataset can be passed as template (must have station coordinates. '''
if not self.source.gdal: raise DatasetError, "Source dataset must be GDAL enabled! {:s} is not.".format(self.source.name)
if template is None: raise NotImplementedError
elif isinstance(template, Dataset):
if not template.hasAxis('station'): raise DatasetError, "Template station dataset needs to have a station axis."
if not ( (template.hasVariable('lat') or template.hasVariable('stn_lat')) and
(template.hasVariable('lon') or template.hasVariable('stn_lon')) ):
raise DatasetError, "Template station dataset needs to have lat/lon arrays for the stations."
else: raise TypeError
# make temporary dataset
if self.source is self.target:
if self.tmp: assert self.source == self.tmpput and self.target == self.tmpput
# the operation can not be performed "in-place"!
self.target = Dataset(name='tmptoo', title='Temporary target dataset for non-in-place operations', varlist=[], atts={})
ltmptoo = True
else: ltmptoo = False
src = self.source; tgt = self.target # short-cuts
# determine source dataset grid definition
if src.griddef is None:
srcgrd = GridDefinition(projection=self.source.projection, geotransform=self.source.geotransform,
size=self.source.mapSize, xlon=self.source.xlon, ylat=self.source.ylat)
else: srcgrd = src.griddef
# figure out horizontal axes (will be replaced with station axis)
if isinstance(xlon,Axis):
if not src.hasAxis(xlon, check=True): raise DatasetError
elif isinstance(xlon,basestring): xlon = src.getAxis(xlon)
else: xlon = src.x if srcgrd.isProjected else src.lon
if isinstance(ylat,Axis):
if not src.hasAxis(ylat, check=True): raise DatasetError
elif isinstance(ylat,basestring): ylat = src.getAxis(ylat)
else: ylat = src.y if srcgrd.isProjected else src.lat
if stnax: # not in source dataset!
if src.hasAxis(stnax, check=True): raise DatasetError, "Source dataset must not have a 'station' axis!"
elif template: stnax = template.station # station axis
else: raise ArgumentError, "A station axis needs to be supplied."
assert isinstance(xlon,Axis) and isinstance(ylat,Axis) and isinstance(stnax,Axis)
# transform to dataset-native coordinate system
if template:
if template.hasVariable('lat'): lats = template.lat.getArray()
else: lats = template.stn_lat.getArray()
if template.hasVariable('lon'): lons = template.lon.getArray()
else: lons = template.stn_lon.getArray()
else: raise NotImplementedError, "Cannot extract station data without a station template Dataset"
# adjust longitudes
if srcgrd.isProjected:
if lons.max() > 180.: lons = np.where(lons > 180., 360.-lons, lons)
# reproject coordinate
latlon = osr.SpatialReference()
latlon.SetWellKnownGeogCS('WGS84') # a normal lat/lon coordinate system
tx = osr.CoordinateTransformation(latlon,srcgrd.projection)
xs = []; ys = []
for i in xrange(len(lons)):
x,y,z = tx.TransformPoint(lons[i].astype(np.float64),lats[i].astype(np.float64))
xs.append(x); ys.append(y); del z
lons = np.array(xs); lats = np.array(ys)
#lons,lats = tx.TransformPoints(lons,lats) # doesn't seem to work...
else:
if lons.min() < 0. and xlon.coord.max() > 180.: lons = np.where(lons < 0., lons + 360., lons)
elif lons.max() > 180. and xlon.coord.min() < 0.: lons = np.where(lons > 180., 360.-lons, lons)
else: pass # source and template do not conflict
# generate index list
ixlon = []; iylat = []; istn = []; zs_err = [] # also record elevation error
lzs = src.hasVariable('zs')
lstnzs = template.hasVariable('zs') or template.hasVariable('stn_zs')
if laltcorr and lzs and lstnzs:
if src.zs.ndim > 2: src.zs = src.zs(time=0, lidx=True) # first time-slice (for CESM)
if src.zs.ndim != 2 or not src.gdal or src.zs.units != 'm': raise VariableError
# consider altidue of surrounding points as well
zs = src.zs.getArray(unmask=True,fillValue=-300)
if template.hasVariable('zs'): stn_zs = template.zs.getArray(unmask=True,fillValue=-300)
else: stn_zs = template.stn_zs.getArray(unmask=True,fillValue=-300)
if src.zs.axisIndex(xlon.name) == 0: zs.transpose() # assuming lat,lon or y,x order is more common
ye,xe = zs.shape # assuming order lat,lon or y,x
xe -= 1; ye -= 1 # last valid index, not length
for n,lon,lat in zip(xrange(len(stnax)),lons,lats):
ip = xlon.getIndex(lon, mode='left', outOfBounds=True)
jp = ylat.getIndex(lat, mode='left', outOfBounds=True)
if ip is not None and jp is not None:
# find neighboring point with smallest altitude error
# ip = im+1 if im < xe else im
# jp = jm+1 if jm < ye else jm
im = ip-1 if ip > 0 else ip
jm = jp-1 if jp > 0 else jp
zdiff = np.Infinity # initialize, so that it triggers at least once
# check four closest grid points
for i in im,ip:
for j in jm,jp:
ze = zs[j,i]-stn_zs[n]
zd = np.abs(ze) # compute elevation error
if zd < zdiff: ii,jj,zdiff,zerr = i,j,zd,ze # preliminary selection, triggers at least once
ixlon.append(ii); iylat.append(jj); istn.append(n); zs_err.append(zerr) # final selection
else:
# just choose horizontally closest point
for n,lon,lat in zip(xrange(len(stnax)),lons,lats):
i = xlon.getIndex(lon, mode='closest', outOfBounds=True)
j = ylat.getIndex(lat, mode='closest', outOfBounds=True)
if i is not None and j is not None:
if lzs: # compute elevation error
zs_err.append(zs[j,i]-stn_zs[n])
ixlon.append(i); iylat.append(j); istn.append(n)
# N.B.: it is necessary to append, because we don't know the number of valid points
ixlon = np.array(ixlon); iylat = np.array(iylat); istn = np.array(istn); zs_err = np.array(zs_err)
# prepare target dataset
# N.B.: attributes should already be set in target dataset (by caller module)
# we are also assuming the new dataset has no axes yet
assert len(tgt.axes) == 0
# add axes from source data
for axname,ax in src.axes.iteritems():
if axname not in (xlon.name,ylat.name):
tgt.addAxis(ax, asNC=True, copy=True)
# add station axis (trim to valid coordinates)
newstnax = stnax.copy(coord=stnax.coord[istn]) # same but with trimmed coordinate array
tgt.addAxis(newstnax, asNC=True, copy=True) # already new copy
# create variable for elevation error
if lzs:
assert len(zs_err) > 0
zs_err = Variable(name='zs_err', units='m', data=zs_err, axes=(newstnax,),
atts=dict(long_name='Station Elevation Error'))
tgt.addVariable(zs_err, asNC=True, copy=True); del zs_err # need to copy to make NC var
# add a bunch of other variables with station meta data
for var in template.variables.itervalues():
if var.ndim == 1 and var.hasAxis(stnax): # station attributes
if var.name[-4:] != '_len' or var.name == 'stn_rec_len': # exclude certain attributes
newvar = var.copy(data=var.getArray()[istn], axes=(newstnax,))
if newvar.name[:4] != 'stn_' and newvar.name[:8] != 'station_' and newvar.name[:8] != 'cluster_':
newvar.name = 'stn_'+newvar.name # copy cluster_* as they are!
# N.B.: we need to rename, or name collisions will happen!
tgt.addVariable(newvar, asNC=True, copy=True); del newvar # need to copy to make NC var
# save all the meta data
tgt.sync()
# prepare function call
function = functools.partial(self.processExtract, ixlon=ixlon, iylat=iylat, ylat=ylat, xlon=xlon, stnax=stnax) # already set parameters
# start process
if self.feedback: print('\n +++ processing point-data extraction +++ ')
self.process(function, **kwargs) # currently 'flush' is the only kwarg
if self.feedback: print('\n')
if self.tmp: self.tmpput = self.target
if ltmptoo: assert self.tmpput.name == 'tmptoo' # set above, when temp. dataset is created
# the previous method sets up the process, the next method performs the computation
def processExtract(self, var, ixlon=None, iylat=None, ylat=None, xlon=None, stnax=None):
''' Extract grid poitns corresponding to stations. '''
# process gdal variables (if a variable has a horiontal grid, it should be GDAL enabled)
if var.gdal:
if self.feedback: print('\n'+var.name),
tgt = self.target
assert xlon in var.axes and ylat in var.axes
assert tgt.hasAxis(stnax, strict=False) and stnax not in var.axes
# assemble new axes
axes = [tgt.getAxis(stnax.name)]
for ax in var.axes:
if ax.name not in (xlon.name,ylat.name) and ax.name != stnax.name: # these axes are just transferred
axes.append(tgt.getAxis(ax.name))
axes = tuple(axes)
shape = tuple(len(ax) for ax in axes)
srcdata = var.getArray(copy=False) # don't make extra copy
# roll x & y axes to the front (xlon first, then ylat, then the rest)
srcdata = np.rollaxis(srcdata, axis=var.axisIndex(ylat.name), start=0)
srcdata = np.rollaxis(srcdata, axis=var.axisIndex(xlon.name), start=0)
assert srcdata.shape == (len(xlon),len(ylat))+shape[1:]
# here we extract the data points
if srcdata.ndim == 2:
tgtdata = srcdata[ixlon,iylat] # constructed above
elif srcdata.ndim > 2:
tgtdata = srcdata[ixlon,iylat,:] # constructed above
else: raise AxisError
#try: except: print srcdata.shape, [slc.max() for slc in slices]
# create new Variable
assert shape == tgtdata.shape
newvar = var.copy(axes=axes, data=tgtdata) # new axes and data
del srcdata, tgtdata # clean up (just to make sure)
else:
var.load() # need to load variables into memory, because we are not doing anything else...
newvar = var # just pass over the variable to the new dataset
# return variable
return newvar
# function pair to compute a climatology from a time-series
def Regrid(self, griddef=None, projection=None, geotransform=None, size=None, xlon=None, ylat=None,
lmask=True, int_interp=None, float_interp=None, **kwargs):
''' Setup climatology and start computation; calls processClimatology. '''
# make temporary gdal dataset
if self.source is self.target:
if self.tmp: assert self.source == self.tmpput and self.target == self.tmpput
# the operation can not be performed "in-place"!
self.target = Dataset(name='tmptoo', title='Temporary target dataset for non-in-place operations', varlist=[], atts={})
ltmptoo = True
else: ltmptoo = False
# make sure the target dataset is a GDAL-enabled dataset
if 'gdal' in self.target.__dict__:
# gdal info alread present
if griddef is not None or projection is not None or geotransform is not None:
raise AttributeError, "Target Dataset '%s' is already GDAL enabled - cannot overwrite settings!"%self.target.name
if self.target.xlon is None: raise GDALError, "Map axis 'xlon' not found!"
if self.target.ylat is None: raise GDALError, "Map axis 'ylat' not found!"
xlon = self.target.xlon; ylat = self.target.ylat
else:
# need to set GDAL parameters
if self.tmp and 'gdal' in self.output.__dict__:
# transfer gdal settings from output to temporary dataset
assert self.target is not self.output
projection = self.output.projection; geotransform = self.output.geotransform
xlon = self.output.xlon; ylat = self.output.ylat
else:
# figure out grid definition from input
if griddef is None:
griddef = GridDefinition(projection=projection, geotransform=geotransform, size=size, xlon=xlon, ylat=ylat)
# pass arguments through GridDefinition, if not provided
projection=griddef.projection; geotransform=griddef.geotransform
xlon=griddef.xlon; ylat=griddef.ylat
# apply GDAL settings target dataset
for ax in (xlon,ylat): self.target.addAxis(ax, loverwrite=True) # i.e. replace if already present
self.target = addGDALtoDataset(self.target, projection=projection, geotransform=geotransform)
# use these map axes
xlon = self.target.xlon; ylat = self.target.ylat
assert isinstance(xlon,Axis) and isinstance(ylat,Axis)
# determine source dataset grid definition
if self.source.griddef is None:
srcgrd = GridDefinition(projection=self.source.projection, geotransform=self.source.geotransform,
size=self.source.mapSize, xlon=self.source.xlon, ylat=self.source.ylat)
else: srcgrd = self.source.griddef
srcres = srcgrd.scale; tgtres = griddef.scale
# determine if shift is necessary to insure correct wrapping
if not srcgrd.isProjected and not griddef.isProjected:
lwrapSrc = srcgrd.wrap360
lwrapTgt = griddef.wrap360
# check grids
for grd in (srcgrd,griddef):
if grd.wrap360:
assert grd.geotransform[0] + grd.geotransform[1]*(len(grd.xlon)-1) > 180
assert np.round(grd.geotransform[1]*len(grd.xlon), decimals=2) == 360 # require 360 deg. to some accuracy...
assert any( grd.xlon.getArray() > 180 ) # need to wrap around
assert all( grd.xlon.getArray() >= 0 )
assert all( grd.xlon.getArray() <= 360 )
else:
assert grd.geotransform[0] + grd.geotransform[1]*(len(grd.xlon)-1) < 180
assert all( grd.xlon.getArray() >= -180 )
assert all( grd.xlon.getArray() <= 180 )
else:
lwrapSrc = False # no need to shift, if a projected grid is involved!
lwrapTgt = False # no need to shift, if a projected grid is involved!
# determine GDAL interpolation
if int_interp is None: int_interp = gdalInterp('nearest')
else: int_interp = gdalInterp(int_interp)
if float_interp is None:
if srcres < tgtres: float_interp = gdalInterp('convolution') # down-sampling: 'convolution'
else: float_interp = gdalInterp('cubicspline') # up-sampling
else: float_interp = gdalInterp(float_interp)
# prepare function call
function = functools.partial(self.processRegrid, ylat=ylat, xlon=xlon, lwrapSrc=lwrapSrc, lwrapTgt=lwrapTgt, # already set parameters
lmask=lmask, int_interp=int_interp, float_interp=float_interp)
# start process
if self.feedback: print('\n +++ processing regridding +++ ')
self.process(function, **kwargs) # currently 'flush' is the only kwarg
# now make sure we have a GDAL dataset!
self.target = addGDALtoDataset(self.target, griddef=griddef)
if self.feedback: print('\n')
if self.tmp: self.tmpput = self.target
if ltmptoo: assert self.tmpput.name == 'tmptoo' # set above, when temp. dataset is created
# the previous method sets up the process, the next method performs the computation
def processRegrid(self, var, ylat=None, xlon=None, lwrapSrc=False, lwrapTgt=False, lmask=True, int_interp=None, float_interp=None):
''' Compute a climatology from a variable time-series. '''
# process gdal variables
if var.gdal:
if self.feedback: print('\n'+var.name),
# replace axes
axes = list(var.axes)
axes[var.axisIndex(var.ylat)] = ylat
axes[var.axisIndex(var.xlon)] = xlon
# create new Variable
var.load() # most rebust way to determine the dtype! and we need it later anyway
newvar = var.copy(axes=axes, data=None, projection=self.target.projection) # and, of course, load new data
# if necessary, shift array back, to ensure proper wrapping of coordinates
# prepare regridding
# get GDAL dataset instances
srcdata = var.getGDAL(load=True, wrap360=lwrapSrc)
tgtdata = newvar.getGDAL(load=False, wrap360=lwrapTgt, allocate=True, fillValue=var.fillValue)
# determine GDAL interpolation
if 'gdal_interp' in var.__dict__: gdal_interp = var.gdal_interp
elif 'gdal_interp' in var.atts: gdal_interp = var.atts['gdal_interp']
else: # use default based on variable type
if np.issubdtype(var.dtype, np.integer): gdal_interp = int_interp # can't process logicals anyway...
else: gdal_interp = float_interp
# perform regridding
err = gdal.ReprojectImage(srcdata, tgtdata, var.projection.ExportToWkt(), newvar.projection.ExportToWkt(), gdal_interp)
#print srcdata.ReadAsArray().std(), tgtdata.ReadAsArray().std()
#print var.projection.ExportToWkt()
#print newvar.projection.ExportToWkt()
del srcdata # clean up (just to make sure)
# N.B.: the target array should be allocated and prefilled with missing values, otherwise ReprojectImage
# will just fill missing values with zeros!
if err != 0: raise GDALError, 'ERROR CODE %i'%err
#tgtdata.FlushCash()
# load data into new variable
newvar.loadGDAL(tgtdata, mask=lmask, wrap360=lwrapTgt, fillValue=var.fillValue)
del tgtdata # clean up (just to make sure)
else:
var.load() # need to load variables into memory, because we are not doing anything else...
newvar = var # just pass over the variable to the new dataset
# return variable
return newvar
# function pair to compute a climatology from a time-series
def Climatology(self, timeAxis='time', climAxis=None, period=None, offset=0, shift=0, timeSlice=None, **kwargs):
''' Setup climatology and start computation; calls processClimatology. '''
if period is not None and not isinstance(period,(np.integer,int)): raise TypeError # period in years
if not isinstance(offset,(np.integer,int)): raise TypeError # offset in years (from start of record)
if not isinstance(shift,(np.integer,int)): raise TypeError # shift in month (if first month is not January)
# construct new time axis for climatology
if climAxis is None:
climAxis = Axis(name=timeAxis, units='month', length=12, coord=np.arange(1,13,1), dtype=dtype_int) # monthly climatology
else:
if not isinstance(climAxis,Axis): raise TypeError
# add axis to output dataset
if self.target.hasAxis(climAxis.name):
self.target.repalceAxis(climAxis, check=False) # will have different shape
else:
self.target.addAxis(climAxis, copy=True) # copy=True allows recasting as, e.g., a NC variable
climAxis = self.target.axes[timeAxis] # make sure we have exactly that instance
# figure out time slice
if period is not None:
start = offset * len(climAxis); end = start + period * len(climAxis)
timeSlice = slice(start,end,None)
else:
if not isinstance(timeSlice,slice): raise TypeError, timeSlice
# add variables that will cause errors to ignorelist (e.g. strings)
for varname,var in self.source.variables.iteritems():
if var.hasAxis(timeAxis) and var.dtype.kind == 'S': self.ignorelist.append(varname)
# prepare function call
function = functools.partial(self.processClimatology, # already set parameters
timeAxis=timeAxis, climAxis=climAxis, timeSlice=timeSlice, shift=shift)
# start process
if self.feedback: print('\n +++ processing climatology +++ ')
if self.source.gdal: griddef = self.source.griddef
else: griddef = None
self.process(function, **kwargs) # currently 'flush' is the only kwarg
# add GDAL to target
if griddef is not None:
self.target = addGDALtoDataset(self.target, griddef=griddef)
# N.B.: if the dataset is empty, it wont do anything, hence we do it now
if self.feedback: print('\n')
# the previous method sets up the process, the next method performs the computation
def processClimatology(self, var, timeAxis='time', climAxis=None, timeSlice=None, shift=0):
''' Compute a climatology from a variable time-series. '''
# process variable that have a time axis
if var.hasAxis(timeAxis):
if self.feedback: print('\n'+var.name),
# prepare averaging
tidx = var.axisIndex(timeAxis)
interval = len(climAxis)
newshape = list(var.shape)
newshape[tidx] = interval # shape of the climatology field
if not (interval == 12): raise NotImplementedError
# load data
if timeSlice is not None:
idx = tuple([timeSlice if ax.name == timeAxis else slice(None) for ax in var.axes])
else: idx = None
dataarray = var.getArray(idx=idx, unmask=False, copy=False)
if var.masked: avgdata = ma.zeros(newshape, dtype=var.dtype) # allocate array
else: avgdata = np.zeros(newshape, dtype=var.dtype) # allocate array
# average data
timelength = dataarray.shape[tidx]
if timelength % interval == 0:
# use array indexing
climelts = np.arange(interval, dtype=dtype_int)
for t in xrange(0,timelength,interval):
if self.feedback: print('.'), # t/interval+1
avgdata += dataarray.take(t+climelts, axis=tidx)
del dataarray # clean up
# normalize
avgdata /= (timelength/interval)
else:
# simple indexing
climcnt = np.zeros(interval, dtype=dtype_int)
for t in xrange(timelength):
if self.feedback and t%interval == 0: print('.'), # t/interval+1
idx = int(t%interval)
climcnt[idx] += 1
if dataarray.ndim == 1:
avgdata[idx] = avgdata[idx] + dataarray[t]
else:
avgdata[idx,:] = avgdata[idx,:] + dataarray[t,:]
del dataarray # clean up
# normalize
for i in xrange(interval):
if avgdata.ndim == 1:
if climcnt[i] > 0: avgdata[i] /= climcnt[i]
else: avgdata[i] = 0 if np.issubdtype(var.dtype, np.integer) else np.NaN
else:
if climcnt[i] > 0: avgdata[i,:] /= climcnt[i]
else: avgdata[i,:] = 0 if np.issubdtype(var.dtype, np.integer) else np.NaN
# shift data (if first month was not January)
if shift != 0: avgdata = np.roll(avgdata, shift, axis=tidx)
# create new Variable
axes = tuple([climAxis if ax.name == timeAxis else ax for ax in var.axes]) # exchange time axis
newvar = var.copy(axes=axes, data=avgdata, dtype=var.dtype) # and, of course, load new data
del avgdata # clean up - just to make sure
# print newvar.name, newvar.masked
# print newvar.fillValue
# print newvar.data_array.__class__
else:
var.load() # need to load variables into memory, because we are not doing anything else...
newvar = var.copy()
# return variable
return newvar
def Shift(self, shift=0, axis=None, byteShift=False, **kwargs):
''' Method to initialize shift along a coordinate axis. '''
# kwarg input
if shift == 0 and axis == None:
for key,value in kwargs.iteritems():
if self.target.hasAxis(key) or self.input.hasAxis(key):
if axis is None: axis = key; shift = value
else: raise ProcessError, "Can only process one coordinate shift at a time."
del kwargs[axis] # remove entry
# check input
if isinstance(axis,basestring):
if self.target.hasAxis(axis): axis = self.target.axes[axis]
elif self.input.hasAxis(axis): axis = self.input.axes[axis].copy()
else: raise AxisError, "Axis '%s' not found in Dataset."%axis
else:
if not isinstance(axis,Axis): raise TypeError
# apply shift to new axis
if byteShift:
# shift coordinate vector like data
coord = np.roll(axis.getArray(unmask=False), shift) # 1-D
else:
coord = axis.getArray(unmask=False) + shift # shift coordinates
# transform coordinate shifts into index shifts (linear scaling)
shift = int( shift / (axis[1] - axis[0]) )
axis.coord = coord
# add axis to output dataset
if self.target.hasAxis(axis, strict=True): pass
elif self.target.hasAxis(axis.name): self.target.repalceAxis(axis)
else: self.target.addAxis(axis, copy=True) # copy=True allows recasting as, e.g., a NC variable
axis = self.target.axes[axis.name] # make sure we have the right version!
# prepare function call
function = functools.partial(self.processShift, # already set parameters
shift=shift, axis=axis)
# start process
if self.feedback: print('\n +++ processing shift/roll +++ ')
self.process(function, **kwargs) # currently 'flush' is the only kwarg
if self.feedback: print('\n')
# the previous method sets up the process, the next method performs the computation
def processShift(self, var, shift=None, axis=None):
''' Method that shifts a data array along a given axis. '''
# only process variables that have the specified axis
if var.hasAxis(axis.name):
if self.feedback: print('\n'+var.name), # put line break before test, instead of after
# shift data array
newdata = np.roll(var.getArray(unmask=False), shift, axis=var.axisIndex(axis))
# create new Variable
axes = tuple([axis if ax.name == axis.name else ax for ax in var.axes]) # replace axis with shifted version
newvar = var.copy(axes=axes, data=newdata) # and, of course, load new data
var.unload(); del var, newdata
else:
var.load() # need to load variables into memory, because we are not doing anything else...
newvar = var
var.unload(); del var
# return variable
return newvar
def ShapeAverage(self, shape_dict=None, shape_name=None, shpax=None, xlon=None, ylat=None,
memory=500, **kwargs):
''' Average over a limited area of a gridded datasets; calls processAverageShape.
A dictionary of NamedShape objects is expected to define the averaging areas.
'memory' controls the garbage collection interval and approximately corresponds
to MB in temporary (it does not include loading the variable into RAM, though). '''
if not self.source.gdal: raise DatasetError, "Source dataset must be GDAL enabled! {:s} is not.".format(self.source.name)
if not isinstance(shape_dict,OrderedDict): raise TypeError
if not all(isinstance(shape,NamedShape) for shape in shape_dict.itervalues()): raise TypeError
# make temporary dataset
if self.source is self.target:
if self.tmp: assert self.source == self.tmpput and self.target == self.tmpput
# the operation can not be performed "in-place"!
self.target = Dataset(name='tmptoo', title='Temporary target dataset for non-in-place operations', varlist=[], atts={})
ltmptoo = True
else: ltmptoo = False
src = self.source; tgt = self.target # short-cuts
# determine source dataset grid definition
if src.griddef is None:
srcgrd = GridDefinition(projection=self.source.projection, geotransform=self.source.geotransform,
size=self.source.mapSize, xlon=self.source.xlon, ylat=self.source.ylat)
else: srcgrd = src.griddef
# figure out horizontal axes (will be replaced with station axis)
if isinstance(xlon,Axis):
if not src.hasAxis(xlon, check=True): raise DatasetError
elif isinstance(xlon,basestring): xlon = src.getAxis(xlon)
else: xlon = src.x if srcgrd.isProjected else src.lon
if isinstance(ylat,Axis):
if not src.hasAxis(ylat, check=True): raise DatasetError
elif isinstance(ylat,basestring): ylat = src.getAxis(ylat)
else: ylat = src.y if srcgrd.isProjected else src.lat
# check/create shapes axis
if shpax: # not in source dataset!
# if shape axis supplied
if src.hasAxis(shpax, check=True): raise DatasetError, "Source dataset must not have a 'shape' axis!"
if len(shpax) != len(shape_dict): raise AxisError
else:
# creat shape axis, if not supplied
shpatts = dict(name='shape', long_name='Ordinal Number of Shape', units='#')
shpax = Axis(coord=np.arange(1,len(shape_dict)+1), atts=shpatts) # starting at 1
assert isinstance(xlon,Axis) and isinstance(ylat,Axis) and isinstance(shpax,Axis)
# prepare target dataset
# N.B.: attributes should already be set in target dataset (by caller module)
# we are also assuming the new dataset has no axes yet
assert len(tgt.axes) == 0
# add station axis (trim to valid coordinates)
tgt.addAxis(shpax, asNC=True, copy=True) # already new copy
# add axes from source data
for axname,ax in src.axes.iteritems():
if axname not in (xlon.name,ylat.name):
tgt.addAxis(ax, asNC=True, copy=True)
# add shape names
shape_names = [shape.name for shape in shape_dict.itervalues()] # can construct Variable from list!
atts = dict(name='shape_name', long_name='Name of Shape', units='')
tgt.addVariable(Variable(data=shape_names, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add proper names
shape_long_names = [shape.long_name for shape in shape_dict.itervalues()] # can construct Variable from list!
atts = dict(name='shp_long_name', long_name='Proper Name of Shape', units='')
tgt.addVariable(Variable(data=shape_long_names, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add shape category
shape_type = [shape.shapetype for shape in shape_dict.itervalues()] # can construct Variable from list!
atts = dict(name='shp_type', long_name='Type of Shape', units='')
tgt.addVariable(Variable(data=shape_type, axes=(shpax,), atts=atts), asNC=True, copy=True)
# collect rasterized masks from shape files
mask_array = np.zeros((len(shpax),)+srcgrd.size[::-1], dtype=np.bool)
# N.B.: rasterize() returns mask in (y,x) shape, size is ordered as (x,y)
shape_masks = []; shp_full = []; shp_empty = []; shp_encl = []
for i,shape in enumerate(shape_dict.itervalues()):
mask = shape.rasterize(griddef=srcgrd, asVar=False)
mask_array[i,:] = mask
masksum = mask.sum()
lfull = masksum == 0; shp_full.append( lfull )
lempty = masksum == mask.size; shp_empty.append( lempty )
shape_masks.append( mask if not lempty else None )
if lempty: shp_encl.append( False )
else:
shp_encl.append( np.all( mask[[0,-1],:] == True ) and np.all( mask[:,[0,-1]] == True ) )
# i.e. if boundaries are masked
# N.B.: shapes that have no overlap with grid will be skipped and filled with NaN
# add rasterized masks to new dataset
atts = dict(name='shp_mask', long_name='Rasterized Shape Mask', units='')
tgt.addVariable(Variable(data=mask_array, atts=atts, axes=(shpax,srcgrd.ylat.copy(),srcgrd.xlon.copy())),
asNC=True, copy=True)
# add area enclosed by shape
da = srcgrd.geotransform[1]*srcgrd.geotransform[5]
mask_area = (1-mask_array).mean(axis=2).mean(axis=1)*da
atts = dict(name='shp_area', long_name='Area Contained in the Shape',
units= 'm^2' if srcgrd.isProjected else 'deg^2' )
tgt.addVariable(Variable(data=mask_area, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add flag to indicate if shape is fully enclosed by domain
atts = dict(name='shp_encl', long_name='If Shape is fully included in Domain', units= '')
tgt.addVariable(Variable(data=shp_encl, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add flag to indicate if shape fully covers domain
atts = dict(name='shp_full', long_name='If Shape fully covers Domain', units= '')
tgt.addVariable(Variable(data=shp_full, axes=(shpax,), atts=atts), asNC=True, copy=True)
# add flag to indicate if shape and domain have no overlap
atts = dict(name='shp_empty', long_name='If Shape and Domain have no Overlap', units= '')
tgt.addVariable(Variable(data=shp_empty, axes=(shpax,), atts=atts), asNC=True, copy=True)
# save all the meta data
tgt.sync()
# prepare function call
function = functools.partial(self.processShapeAverage, masks=shape_masks, ylat=ylat, xlon=xlon,
shpax=shpax, memory=memory) # already set parameters
# start process
if self.feedback: print('\n +++ processing shape/area averaging +++ ')
self.process(function, **kwargs) # currently 'flush' is the only kwarg
if self.feedback: print('\n')
if self.tmp: self.tmpput = self.target
if ltmptoo: assert self.tmpput.name == 'tmptoo' # set above, when temp. dataset is created
# imports
from glob import glob
from geodata.base import Dataset, Axis, Variable
from geodata.netcdf import writeNetCDF
# load list if well files and generate list of wells
well_files = glob(os.path.join(data_folder, 'W*.xlsx'))
well_files.sort()
wells = [os.path.basename(name[:-5]) for name in well_files]
print(wells)
# dataset
time_ax = Axis(coord=np.arange(12 * (period[1] - period[0])) + 252,
**varatts['time']) # origin: 1979-01
well_ax = Axis(coord=np.arange(len(wells)) + 1, name='well', units='')
dataset = Dataset(name=conservation_authority,
title=conservation_authority + ' Observation Wells')
# add meta data
meta_dicts = [
loadMetadata(well, conservation_authority=conservation_authority)
for well in wells
]
for key in meta_dicts[0].keys():
if key in varatts: atts = varatts[key]
elif key.lower() in varatts: atts = varatts[key.lower()]
else: atts = dict(name=key, units='')
if atts['units']:
data = np.asarray([wmd[key] for wmd in meta_dicts],
dtype=np.float64)
else:
data = np.asarray([wmd[key] for wmd in meta_dicts])
try:
def performRegridding(dataset,
mode,
griddef,
dataargs,
loverwrite=False,
varlist=None,
lwrite=True,
lreturn=False,
ldebug=False,
lparallel=False,
pidstr='',
logger=None):
''' worker function to perform regridding for a given dataset and target grid '''
# input checking
if not isinstance(dataset, basestring): raise TypeError
if not isinstance(dataargs, dict):
raise TypeError # all dataset arguments are kwargs
if not isinstance(griddef, GridDefinition): raise TypeError
if lparallel:
if not lwrite:
raise IOError, 'Can only write to disk in parallel mode (i.e. lwrite = True).'
if lreturn:
raise IOError, 'Can not return datasets in parallel mode (i.e. lreturn = False).'
# logging
if logger is None: # make new logger
logger = logging.getLogger() # new logger
logger.addHandler(logging.StreamHandler())
else:
if isinstance(logger, basestring):
logger = logging.getLogger(name=logger) # connect to existing one
elif not isinstance(logger, logging.Logger):
raise TypeError, 'Expected logger ID/handle in logger KW; got {}'.format(
str(logger))
## extract meta data from arguments
dataargs, loadfct, srcage, datamsgstr = getMetaData(
dataset, mode, dataargs)
dataset_name = dataargs.dataset_name
periodstr = dataargs.periodstr
avgfolder = dataargs.avgfolder
# get filename for target dataset and do some checks
filename = getTargetFile(
dataset=dataset,
mode=mode,
dataargs=dataargs,
lwrite=lwrite,
grid=griddef.name.lower(),
)
# prepare target dataset
if ldebug: filename = 'test_' + filename
if not os.path.exists(avgfolder):
raise IOError, "Dataset folder '{:s}' does not exist!".format(
avgfolder)
lskip = False # else just go ahead
if lwrite:
if lreturn:
tmpfilename = filename # no temporary file if dataset is passed on (can't rename the file while it is open!)
else:
if lparallel: tmppfx = 'tmp_regrid_{:s}_'.format(pidstr[1:-1])
else: tmppfx = 'tmp_regrid_'.format(pidstr[1:-1])
tmpfilename = tmppfx + filename
filepath = avgfolder + filename
tmpfilepath = avgfolder + tmpfilename
if os.path.exists(filepath):
if not loverwrite:
age = datetime.fromtimestamp(os.path.getmtime(filepath))
# if source file is newer than sink file or if sink file is a stub, recompute, otherwise skip
if age > srcage and os.path.getsize(filepath) > 1e6:
lskip = True
if hasattr(griddef,
'filepath') and griddef.filepath is not None:
gridage = datetime.fromtimestamp(
os.path.getmtime(griddef.filepath))
if age < gridage: lskip = False
# N.B.: NetCDF files smaller than 1MB are usually incomplete header fragments from a previous crashed
# depending on last modification time of file or overwrite setting, start computation, or skip
if lskip:
# print message
skipmsg = "\n{:s} >>> Skipping: file '{:s}' in dataset '{:s}' already exists and is newer than source file.".format(
pidstr, filename, dataset_name)
skipmsg += "\n{:s} >>> ('{:s}')\n".format(pidstr, filepath)
logger.info(skipmsg)
else:
## actually load datasets
source = loadfct() # load source
# check period
if 'period' in source.atts and dataargs.periodstr != source.atts.period: # a NetCDF attribute
raise DateError, "Specifed period is inconsistent with netcdf records: '{:s}' != '{:s}'".format(
periodstr, source.atts.period)
# print message
if mode == 'climatology':
opmsgstr = 'Regridding Climatology ({:s}) to {:s} Grid'.format(
periodstr, griddef.name)
elif mode == 'time-series':
opmsgstr = 'Regridding Time-series to {:s} Grid'.format(
griddef.name)
else:
raise NotImplementedError, "Unrecognized Mode: '{:s}'".format(mode)
# print feedback to logger
logger.info(
'\n{0:s} *** {1:^65s} *** \n{0:s} *** {2:^65s} *** \n'
.format(pidstr, datamsgstr, opmsgstr))
if not lparallel and ldebug: logger.info('\n' + str(source) + '\n')
## create new sink/target file
# set attributes
atts = source.atts.copy()
atts['period'] = periodstr
atts['name'] = dataset_name
atts['grid'] = griddef.name
if mode == 'climatology':
atts['title'] = '{:s} Climatology on {:s} Grid'.format(
dataset_name, griddef.name)
elif mode == 'time-series':
atts['title'] = '{:s} Time-series on {:s} Grid'.format(
dataset_name, griddef.name)
# make new dataset
if lwrite: # write to NetCDF file
if os.path.exists(tmpfilepath):
os.remove(tmpfilepath) # remove old temp files
sink = DatasetNetCDF(folder=avgfolder,
filelist=[tmpfilename],
atts=atts,
mode='w')
else:
sink = Dataset(atts=atts) # ony create dataset in memory
# initialize processing
CPU = CentralProcessingUnit(source,
sink,
varlist=varlist,
tmp=False,
feedback=ldebug)
# perform regridding (if target grid is different from native grid!)
if griddef.name != dataset:
# reproject and resample (regrid) dataset
CPU.Regrid(griddef=griddef, flush=True)
# get results
CPU.sync(flush=True)
# add geolocators
sink = addGeoLocator(sink,
griddef=griddef,
lgdal=True,
lreplace=True,
lcheck=True)
# N.B.: WRF datasets come with their own geolocator arrays - we need to replace those!
# add length and names of month
if mode == 'climatology' and not sink.hasVariable(
'length_of_month') and sink.hasVariable('time'):
addLengthAndNamesOfMonth(
sink,
noleap=True if dataset.upper() in ('WRF', 'CESM') else False)
# print dataset
if not lparallel and ldebug:
logger.info('\n' + str(sink) + '\n')
# write results to file
if lwrite:
sink.sync()
writemsg = "\n{:s} >>> Writing to file '{:s}' in dataset {:s}".format(
pidstr, filename, dataset_name)
writemsg += "\n{:s} >>> ('{:s}')\n".format(pidstr, filepath)
logger.info(writemsg)
# rename file to proper name
if not lreturn:
sink.unload()
sink.close()
del sink # destroy all references
if os.path.exists(filepath):
os.remove(filepath) # remove old file
os.rename(
tmpfilepath,
filepath) # this would also overwrite the old file...
# N.B.: there is no temporary file if the dataset is returned, because an open file can't be renamed
# clean up and return
source.unload()
del source, CPU
if lreturn:
return sink # return dataset for further use (netcdf file still open!)
else:
return 0 # "exit code"
def performExport(dataset, mode, dataargs, expargs, bcargs, loverwrite=False,
ldebug=False, lparallel=False, pidstr='', logger=None):
''' worker function to export ASCII rasters for a given dataset '''
# input checking
if not isinstance(dataset,basestring): raise TypeError
if not isinstance(dataargs,dict): raise TypeError # all dataset arguments are kwargs
# logging
if logger is None: # make new logger
logger = logging.getLogger() # new logger
logger.addHandler(logging.StreamHandler())
else:
if isinstance(logger,basestring):
logger = logging.getLogger(name=logger) # connect to existing one
elif not isinstance(logger,logging.Logger):
raise TypeError, 'Expected logger ID/handle in logger KW; got {}'.format(str(logger))
## extract meta data from arguments
dataargs, loadfct, srcage, datamsgstr = getMetaData(dataset, mode, dataargs, lone=False)
dataset_name = dataargs.dataset_name; periodstr = dataargs.periodstr; domain = dataargs.domain
# figure out bias correction parameters
if bcargs:
bcargs = bcargs.copy() # first copy, then modify...
bc_method = bcargs.pop('method',None)
if bc_method is None: raise ArgumentError("Need to specify bias-correction method to use bias correction!")
bc_obs = bcargs.pop('obs_dataset',None)
if bc_obs is None: raise ArgumentError("Need to specify observational dataset to use bias correction!")
bc_reference = bcargs.pop('reference',None)
if bc_reference is None: # infer from experiment name
if dataset_name[-5:] in ('-2050','-2100'): bc_reference = dataset_name[:-5] # cut of period indicator and hope for the best
else: bc_reference = dataset_name
bc_grid = bcargs.pop('grid',None)
if bc_grid is None: bc_grid = dataargs.grid
bc_domain = bcargs.pop('domain',None)
if bc_domain is None: bc_domain = domain
bc_varlist = bcargs.pop('varlist',None)
bc_varmap = bcargs.pop('varmap',None)
bc_tag = bcargs.pop('tag',None) # an optional name extension/tag
bc_pattern = bcargs.pop('file_pattern',None) # usually default in getPickleFile
lgzip = bcargs.pop('lgzip',None) # if pickle is gzipped (None: auto-detect based on file name extension)
# get name of pickle file (and folder)
picklefolder = dataargs.avgfolder.replace(dataset_name,bc_reference)
picklefile = getPickleFileName(method=bc_method, obs_name=bc_obs, gridstr=bc_grid, domain=bc_domain,
tag=bc_tag, pattern=bc_pattern)
picklepath = '{:s}/{:s}'.format(picklefolder,picklefile)
if lgzip:
picklepath += '.gz' # add extension
if not os.path.exists(picklepath): raise IOError(picklepath)
elif lgzip is None:
lgzip = False
if not os.path.exists(picklepath):
lgzip = True # assume gzipped file
picklepath += '.gz' # try with extension...
if not os.path.exists(picklepath): raise IOError(picklepath)
elif not os.path.exists(picklepath): raise IOError(picklepath)
pickleage = datetime.fromtimestamp(os.path.getmtime(picklepath))
# determine age of pickle file and compare against source age
else:
bc_method = False
pickleage = srcage
# parse export options
expargs = expargs.copy() # first copy, then modify...
lm3 = expargs.pop('lm3') # convert kg/m^2/s to m^3/m^2/s (water flux)
expformat = expargs.pop('format') # needed to get FileFormat object
exp_list= expargs.pop('exp_list') # this handled outside of export
compute_list = expargs.pop('compute_list', []) # variables to be (re-)computed - by default all
# initialize FileFormat class instance
fileFormat = getFileFormat(expformat, bc_method=bc_method, **expargs)
# get folder for target dataset and do some checks
expname = '{:s}_d{:02d}'.format(dataset_name,domain) if domain else dataset_name
expfolder = fileFormat.defineDataset(dataset=dataset, mode=mode, dataargs=dataargs, lwrite=True, ldebug=ldebug)
# prepare destination for new dataset
lskip = fileFormat.prepareDestination(srcage=max(srcage,pickleage), loverwrite=loverwrite)
# depending on last modification time of file or overwrite setting, start computation, or skip
if lskip:
# print message
skipmsg = "\n{:s} >>> Skipping: Format '{:s} for dataset '{:s}' already exists and is newer than source file.".format(pidstr,expformat,dataset_name)
skipmsg += "\n{:s} >>> ('{:s}')\n".format(pidstr,expfolder)
logger.info(skipmsg)
else:
## actually load datasets
source = loadfct() # load source data
# check period
if 'period' in source.atts and dataargs.periodstr != source.atts.period: # a NetCDF attribute
raise DateError, "Specifed period is inconsistent with netcdf records: '{:s}' != '{:s}'".format(periodstr,source.atts.period)
# load BiasCorrection object from pickle
if bc_method:
op = gzip.open if lgzip else open
with op(picklepath, 'r') as filehandle:
BC = pickle.load(filehandle)
# assemble logger entry
bcmsgstr = "(performing bias-correction using {:s} from {:s} towards {:s})".format(BC.long_name,bc_reference,bc_obs)
# print message
if mode == 'climatology': opmsgstr = 'Exporting Climatology ({:s}) to {:s} Format'.format(periodstr, expformat)
elif mode == 'time-series': opmsgstr = 'Exporting Time-series to {:s} Format'.format(expformat)
elif mode[-5:] == '-mean': opmsgstr = 'Exporting {:s}-Mean ({:s}) to {:s} Format'.format(mode[:-5], periodstr, expformat)
else: raise NotImplementedError, "Unrecognized Mode: '{:s}'".format(mode)
# print feedback to logger
logmsg = '\n{0:s} *** {1:^65s} *** \n{0:s} *** {2:^65s} *** \n'.format(pidstr,datamsgstr,opmsgstr)
if bc_method:
logmsg += "{0:s} *** {1:^65s} *** \n".format(pidstr,bcmsgstr)
logger.info(logmsg)
if not lparallel and ldebug: logger.info('\n'+str(source)+'\n')
# create GDAL-enabled target dataset
sink = Dataset(axes=(source.xlon,source.ylat), name=expname, title=source.title, atts=source.atts.copy())
addGDALtoDataset(dataset=sink, griddef=source.griddef)
assert sink.gdal, sink
# apply bias-correction
if bc_method:
source = BC.correct(source, asNC=False, varlist=bc_varlist, varmap=bc_varmap) # load bias-corrected variables into memory
# N.B.: for variables that are not bias-corrected, data are not loaded immediately but on demand; this way
# I/O and computing can be further disentangled and not all variables are always needed
# compute intermediate variables, if necessary
for varname in exp_list:
variables = None # variable list
var = None
# (re-)compute variable, if desired...
if varname in compute_list:
if varname == 'precip': var = newvars.computeTotalPrecip(source)
elif varname == 'waterflx': var = newvars.computeWaterFlux(source)
elif varname == 'liqwatflx': var = newvars.computeLiquidWaterFlux(source)
elif varname == 'netrad': var = newvars.computeNetRadiation(source, asVar=True)
elif varname == 'netrad_bb': var = newvars.computeNetRadiation(source, asVar=True, lrad=False, name='netrad_bb')
elif varname == 'netrad_bb0': var = newvars.computeNetRadiation(source, asVar=True, lrad=False, lA=False, name='netrad_bb0')
elif varname == 'vapdef': var = newvars.computeVaporDeficit(source)
elif varname in ('pet','pet_pm','petrad','petwnd') and 'pet' not in sink:
if 'petrad' in exp_list or 'petwnd' in exp_list:
variables = newvars.computePotEvapPM(source, lterms=True) # default; returns mutliple PET terms
else: var = newvars.computePotEvapPM(source, lterms=False) # returns only PET
elif varname == 'pet_th': var = None # skip for now
#var = computePotEvapTh(source) # simplified formula (less prerequisites)
# ... otherwise load from source file
if var is None and variables is None and varname in source:
var = source[varname].load() # load data (may not have to load all)
#else: raise VariableError, "Unsupported Variable '{:s}'.".format(varname)
# for now, skip variables that are None
if var or variables:
# handle lists as well
if var and variables: raise VariableError, (var,variables)
elif var: variables = (var,)
for var in variables:
addGDALtoVar(var=var, griddef=sink.griddef)
if not var.gdal and isinstance(fileFormat,ASCII_raster):
raise GDALError, "Exporting to ASCII_raster format requires GDAL-enabled variables."
# add to new dataset
sink += var
# convert units
if lm3:
for var in sink:
if var.units == 'kg/m^2/s':
var /= 1000. # divide to get m^3/m^2/s
var.units = 'm^3/m^2/s' # update units
# compute seasonal mean if we are in mean-mode
if mode[-5:] == '-mean':
sink = sink.seasonalMean(season=mode[:-5], lclim=True)
# N.B.: to remain consistent with other output modes,
# we need to prevent renaming of the time axis
sink = concatDatasets([sink,sink], axis='time', lensembleAxis=True)
sink.squeeze() # we need the year-axis until now to distinguish constant fields; now remove
# print dataset
if not lparallel and ldebug:
logger.info('\n'+str(sink)+'\n')
# export new dataset to selected format
fileFormat.exportDataset(sink)
# write results to file
writemsg = "\n{:s} >>> Export of Dataset '{:s}' to Format '{:s}' complete.".format(pidstr,expname, expformat)
writemsg += "\n{:s} >>> ('{:s}')\n".format(pidstr,expfolder)
logger.info(writemsg)
# clean up and return
source.unload(); #del source
return 0 # "exit code"
def loadHGS_StnTS(station=None, varlist=None, varatts=None, folder=None, name=None, title=None,
start_date=None, end_date=None, run_period=15, period=None, lskipNaN=False, lcheckComplete=True,
basin=None, WSC_station=None, basin_list=None, filename=None, prefix=None,
scalefactors=None, **kwargs):
''' Get a properly formatted WRF dataset with monthly time-series at station locations; as in
the hgsrun module, the capitalized kwargs can be used to construct folders and/or names '''
if folder is None or ( filename is None and station is None ): raise ArgumentError
# try to find meta data for gage station from WSC
HGS_station = station
if basin is not None and basin_list is not None:
station_name = station
station = getGageStation(basin=basin, station=station if WSC_station is None else WSC_station,
basin_list=basin_list) # only works with registered basins
if station_name is None: station_name = station.name # backup, in case we don't have a HGS station name
metadata = station.getMetaData() # load station meta data
if metadata is None: raise GageStationError(name)
else:
metadata = dict(); station = None; station_name = None
# prepare name expansion arguments (all capitalized)
expargs = dict(ROOT_FOLDER=root_folder, STATION=HGS_station, NAME=name, TITLE=title,
PREFIX=prefix, BASIN=basin, WSC_STATION=WSC_station)
for key,value in metadata.items():
if isinstance(value,basestring):
expargs['WSC_'+key.upper()] = value # in particular, this includes WSC_ID
if 'WSC_ID' in expargs:
if expargs['WSC_ID'][0] == '0': expargs['WSC_ID0'] = expargs['WSC_ID'][1:]
else: raise DatasetError('Expected leading zero in WSC station ID: {}'.format(expargs['WSC_ID']))
# exparg preset keys will get overwritten if capitalized versions are defined
for key,value in kwargs.items():
KEY = key.upper() # we only use capitalized keywords, and non-capitalized keywords are only used/converted
if KEY == key or KEY not in kwargs: expargs[KEY] = value # if no capitalized version is defined
# read folder and infer prefix, if necessary
folder = folder.format(**expargs)
if not os.path.exists(folder): raise IOError(folder)
if expargs['PREFIX'] is None:
with open('{}/{}'.format(folder,prefix_file), 'r') as pfx:
expargs['PREFIX'] = prefix = ''.join(pfx.readlines()).strip()
# now assemble file name for station timeseries
filename = filename.format(**expargs)
filepath = '{}/{}'.format(folder,filename)
if not os.path.exists(filepath): IOError(filepath)
if station_name is None:
station_name = filename[filename.index('hydrograph.')+1:-4] if station is None else station
# set meta data (and allow keyword expansion of name and title)
metadata['problem'] = prefix
metadata['station_name'] = metadata.get('long_name', station_name)
if name is not None: name = name.format(**expargs) # name expansion with capitalized keyword arguments
else: name = 'HGS_{:s}'.format(station_name)
metadata['name'] = name; expargs['Name'] = name.title() # name in title format
if title is None: title = '{{Name:s}} (HGS, {problem:s})'.format(**metadata)
title = title.format(**expargs) # name expansion with capitalized keyword arguments
metadata['long_name'] = metadata['title'] = title
# now determine start data for date_parser
if end_date is None:
if start_date and run_period: end_date = start_date + run_period
elif period: end_date = period[1]
else: raise ArgumentError("Need to specify either 'start_date' & 'run_period' or 'period' to infer 'end_date'.")
end_year,end_month,end_day = convertDate(end_date)
if start_date is None:
if end_date and run_period: start_date = end_date - run_period
elif period: start_date = period[0]
else: raise ArgumentError("Need to specify either 'end_date' & 'run_period' or 'period' to infer 'start_date'.")
start_year,start_month,start_day = convertDate(start_date)
if start_day != 1 or end_day != 1:
raise NotImplementedError('Currently only monthly data is supported.')
# import functools
# date_parser = functools.partial(date_parser, year=start_year, month=start_month, day=start_day)
# # now load data using pandas ascii reader
# data_frame = pd.read_table(filepath, sep='\s+', header=2, dtype=np.float64, index_col=['time'],
# date_parser=date_parser, names=ascii_varlist)
# # resample to monthly data
# data_frame = data_frame.resample(resampling).agg(np.mean)
# data = data_frame[flowvar].values
# parse header
if varlist is None: varlist = variable_list[:] # default list
with open(filepath, 'r') as f:
line = f.readline(); lline = line.lower() # 1st line
if not "hydrograph" in lline: raise GageStationError(line,filepath)
# parse variables and determine columns
line = f.readline(); lline = line.lower() # 2nd line
if not "variables" in lline: raise GageStationError(line)
variable_order = [v.strip('"').lower() for v in line[line.find('"'):].strip().split(',')]
# figure out varlist and data columns
if variable_order[0] == 'time': del variable_order[0] # only keep variables
else: raise GageStationError(variable_order)
variable_order = [hgs_variables[v] for v in variable_order] # replace HGS names with GeoPy names
vardict = {v:i+1 for i,v in enumerate(variable_order)} # column mapping; +1 because time was removed
variable_order = [v for v in variable_order if v in varlist or flow_to_flux[v] in varlist]
usecols = tuple(vardict[v] for v in variable_order) # variable columns that need to loaded (except time, which is col 0)
assert 0 not in usecols, usecols
# load data as tab separated values
data = np.genfromtxt(filepath, dtype=np.float64, delimiter=None, skip_header=3, usecols = (0,)+usecols)
assert data.shape[1] == len(usecols)+1, data.shape
if lskipNaN:
data = data[np.isnan(data).sum(axis=1)==0,:]
elif np.any( np.isnan(data) ):
raise DataError("Missing values (NaN) encountered in hydrograph file; use 'lskipNaN' to ignore.\n('{:s}')".format(filepath))
time_series = data[:,0]; flow_data = data[:,1:]
assert flow_data.shape == (len(time_series),len(usecols)), flow_data.shape
# original time deltas in seconds
time_diff = time_series.copy(); time_diff[1:] = np.diff(time_series) # time period between time steps
assert np.all( time_diff > 0 ), filepath
time_diff = time_diff.reshape((len(time_diff),1)) # reshape to make sure broadcasting works
# integrate flow over time steps before resampling
flow_data[1:,:] -= np.diff(flow_data, axis=0)/2. # get average flow between time steps
flow_data *= time_diff # integrate flow in time interval by multiplying average flow with time period
flow_data = np.cumsum(flow_data, axis=0) # integrate by summing up total flow per time interval
# generate regular monthly time steps
start_datetime = np.datetime64(dt.datetime(year=start_year, month=start_month, day=start_day), 'M')
end_datetime = np.datetime64(dt.datetime(year=end_year, month=end_month, day=end_day), 'M')
time_monthly = np.arange(start_datetime, end_datetime+np.timedelta64(1, 'M'), dtype='datetime64[M]')
assert time_monthly[0] == start_datetime, time_monthly[0]
assert time_monthly[-1] == end_datetime, time_monthly[-1]
# convert monthly time series to regular array of seconds since start date
time_monthly = ( time_monthly.astype('datetime64[s]') - start_datetime.astype('datetime64[s]') ) / np.timedelta64(1,'s')
assert time_monthly[0] == 0, time_monthly[0]
# interpolate integrated flow to new time axis
#flow_data = np.interp(time_monthly, xp=time_series[:,0], fp=flow_data[:,0],).reshape((len(time_monthly),1))
time_series = np.concatenate(([0],time_series), axis=0) # integrated flow at time zero must be zero...
flow_data = np.concatenate(([[0,]*len(usecols)],flow_data), axis=0) # ... this is probably better than interpolation
# N.B.: we are adding zeros here so we don't have to extrapolate to the left; on the right we just fill in NaN's
if ( time_monthly[-1] - time_series[-1] ) > 3*86400. and lcheckComplete:
warn("Data record ends more than 3 days befor end of period: {} days".format((time_monthly[-1]-time_series[-1])/86400.))
elif (time_monthly[-1]-time_series[-1]) > 5*86400.:
if lcheckComplete:
raise DataError("Data record ends more than 5 days befor end of period: {} days".format((time_monthly[-1]-time_series[-1])/86400.))
else:
warn("Data record ends more than 5 days befor end of period: {} days".format((time_monthly[-1]-time_series[-1])/86400.))
flow_interp = si.interp1d(x=time_series, y=flow_data, kind='linear', axis=0, copy=False,
bounds_error=False, fill_value=np.NaN, assume_sorted=True)
flow_data = flow_interp(time_monthly) # evaluate with call
# compute monthly flow rate from interpolated integrated flow
flow_data = np.diff(flow_data, axis=0) / np.diff(time_monthly, axis=0).reshape((len(time_monthly)-1,1))
flow_data *= 1000 # convert from m^3/s to kg/s
# construct time axis
start_time = 12*(start_year - 1979) + start_month -1
end_time = 12*(end_year - 1979) + end_month -1
time = Axis(name='time', units='month', atts=dict(long_name='Month since 1979-01'),
coord=np.arange(start_time, end_time)) # not including the last, e.g. 1979-01 to 1980-01 is 12 month
assert len(time_monthly) == end_time-start_time+1
assert flow_data.shape == (len(time),len(variable_order)), (flow_data.shape,len(time),len(variable_order))
# construct dataset
dataset = Dataset(atts=metadata)
dataset.station = station # add gage station object, if available (else None)
for i,flowvar in enumerate(variable_order):
data = flow_data[:,i]
fluxvar = flow_to_flux[flowvar]
if flowvar in varlist:
flowatts = variable_attributes[flowvar]
# convert variables and put into dataset (monthly time series)
if flowatts['units'] != 'kg/s':
raise VariableError("Hydrograph data is read as kg/s; flow variable does not match.\n{}".format(flowatts))
dataset += Variable(data=data, axes=(time,), **flowatts)
if fluxvar in varlist and 'shp_area' in metadata:
# compute surface flux variable based on drainage area
fluxatts = variable_attributes[fluxvar]
if fluxatts['units'] == 'kg/s' and fluxatts['units'] != 'kg/m^2/s': raise VariableError(fluxatts)
data = data / metadata['shp_area'] # need to make a copy
dataset += Variable(data=data, axes=(time,), **fluxatts)
# apply analysis period
if period is not None:
dataset = dataset(years=period)
# adjust scalefactors, if necessary
if scalefactors:
if isinstance(scalefactors,dict):
dataset = updateScalefactor(dataset, varlist=scalefactors, scalefactor=None)
elif isNumber(scalefactors):
scalelist = ('discharge','seepage','flow')
dataset = updateScalefactor(dataset, varlist=scalelist, scalefactor=scalefactors)
else:
raise TypeError(scalefactors)
# return completed dataset
return dataset
def loadGageStation(basin=None, station=None, varlist=None, varatts=None, mode='climatology',
aggregation=None, filetype='monthly', folder=None, name=None, period=None,
basin_list=None, lcheck=True, lexpand=True, lfill=True, lflatten=True,
lkgs=True, scalefactors=None, title=None):
''' function to load hydrograph climatologies and timeseries for a given basin '''
## resolve input
if mode == 'timeseries' and aggregation:
raise ArgumentError('Timeseries does not support aggregation.')
# get GageStation instance
station = getGageStation(basin=basin, station=station, name=name, folder=folder,
river=None, basin_list=basin_list, lcheck=True)
# variable attributes
if varlist is None: varlist = variable_list
elif not isinstance(varlist,(list,tuple)): raise TypeError
varlist = list(varlist) # make copy of varlist to avoid interference
if varatts is None:
if aggregation is None: varatts = variable_attributes_kgs if lkgs else variable_attributes_mms
else: varatts = agg_varatts_kgs if lkgs else agg_varatts_mms
elif not isinstance(varatts,dict): raise TypeError
## read csv data
# time series data and time coordinates
lexpand = True; lfill = True
if mode == 'climatology': lexpand = False; lfill = False; lflatten = False
data, time = station.getTimeseriesData(units='kg/s' if lkgs else 'm^3/s', lcheck=True, lexpand=lexpand,
lfill=lfill, period=period, lflatten=lflatten)
# station meta data
metadata = station.getMetaData(lcheck=True)
den = metadata['shp_area'] if lkgs else ( metadata['shp_area'] / 1000. )
## create dataset for station
dataset = Dataset(name='WSC', title=title or metadata['Station Name'], varlist=[], atts=metadata,)
if mode.lower() in ('timeseries','time-series'):
time = time.flatten(); data = data.flatten() # just to make sure...
# make time axis based on time coordinate from csv file
timeAxis = Axis(name='time', units='month', coord=time, # time series centered at 1979-01
atts=dict(long_name='Month since 1979-01'))
dataset += timeAxis
# load mean discharge
dataset += Variable(axes=[timeAxis], data=data, atts=varatts['discharge'])
# load mean runoff
doa = data / den
dataset += Variable(axes=[timeAxis], data=doa, atts=varatts['runoff'])
elif mode == 'climatology':
# N.B.: this is primarily for backwards compatibility; it should not be used anymore...
# make common time axis for climatology
te = 12 # length of time axis: 12 month
climAxis = Axis(name='time', units='month', length=12, coord=np.arange(1,te+1,1)) # monthly climatology
dataset.addAxis(climAxis, copy=False)
# extract variables (min/max/mean are separate variables)
# N.B.: this is mainly for backwards compatibility
doa = data / den
if aggregation is None or aggregation.lower() == 'mean':
# load mean discharge
tmpdata = nf.nanmean(data, axis=0)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['discharge'])
dataset.addVariable(tmpvar, copy=False)
# load mean runoff
tmpdata = nf.nanmean(doa, axis=0)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['runoff'])
dataset.addVariable(tmpvar, copy=False)
if aggregation is None or aggregation.lower() == 'std':
# load discharge standard deviation
tmpdata = nf.nanstd(data, axis=0, ddof=1) # very few values means large uncertainty!
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['discstd'])
dataset.addVariable(tmpvar, copy=False)
# load runoff standard deviation
tmpdata = nf.nanstd(doa, axis=0, ddof=1)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['roff_std'])
dataset.addVariable(tmpvar, copy=False)
if aggregation is None or aggregation.lower() == 'sem':
# load discharge standard deviation
tmpdata = nf.nansem(data, axis=0, ddof=1) # very few values means large uncertainty!
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['discsem'])
dataset.addVariable(tmpvar, copy=False)
# load runoff standard deviation
tmpdata = nf.nansem(doa, axis=0, ddof=1)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['roff_sem'])
dataset.addVariable(tmpvar, copy=False)
if aggregation is None or aggregation.lower() == 'max':
# load maximum discharge
tmpdata = nf.nanmax(data, axis=0)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['discmax'])
dataset.addVariable(tmpvar, copy=False)
# load maximum runoff
tmpdata = nf.nanmax(doa, axis=0)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['roff_max'])
dataset.addVariable(tmpvar, copy=False)
if aggregation is None or aggregation.lower() == 'min':
# load minimum discharge
tmpdata = nf.nanmin(data, axis=0)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['discmin'])
dataset.addVariable(tmpvar, copy=False)
# load minimum runoff
tmpdata = nf.nanmin(doa, axis=0)
tmpvar = Variable(axes=[climAxis], data=tmpdata, atts=varatts['roff_min'])
dataset.addVariable(tmpvar, copy=False)
else:
raise NotImplementedError, "Time axis mode '{}' is not supported.".format(mode)
# adjust scalefactors, if necessary
if scalefactors:
if isinstance(scalefactors,dict):
dataset = updateScalefactor(dataset, varlist=scalefactors, scalefactor=None)
elif isNumber(scalefactors):
scalelist = ('discharge','StdDisc','SEMDisc','MaxDisc','MinDisc',)
dataset = updateScalefactor(dataset, varlist=scalelist, scalefactor=scalefactors)
else:
raise TypeError(scalefactors)
# return station dataset
return dataset