def run(self):
if self.netcdf3: io.netcdf3()
else: io.netcdf4(level=self.netcdf4level)
cdms2.setAutoBounds(0)
self.notice('Masking %s to %s', self.inpfile, self.outfile)
inpfile = cdms2.open(self.inpfile)
if os.path.isfile(self.outfile):
if os.path.samefile(self.inpfile, self.outfile):
raise Exception('Cannot use same input and output file')
if self.overwrite:
os.remove(self.outfile)
else:
raise Exception('Output file already exists and overwriting is not requested')
outfile = cdms2.open(self.outfile, 'w')
# Copy global attributes
for a,v in inpfile.attributes.items():
setattr(outfile, a, v)
# Keep grid masks in memory to improve performances (TODO: use a file cache as for basemaps ?)
mask_cache = dict()
stats = True # self.is_verbose()
# Iterate over input variables
for varid,filevar in inpfile.variables.items():
# Process only gridded/specified variables
grid = filevar.getGrid()
maskit = grid is not None
if maskit and self.included_variables and varid not in self.included_variables: maskit = False
if maskit and self.excluded_variables and varid in self.excluded_variables: maskit = False
if not maskit and self.masked_only:
self.verbose('Ignoring %s: variable has no grid and masked_only was specified')
continue
self.logger.info(bases.psinfo())
self.notice('Processing variable: %s, grid: %s', varid, grid.shape if grid else None)
# NOTE: With scalar variables, memvar could be numpy.<type> instead of cdms2.tvariable.TransientVariable
memvar = filevar()
self.info(bases.describe(memvar, stats=stats))
self.logger.info(bases.psinfo())
if maskit:
# Build the mask, check if it is already cached ?
cache_id = id(grid)
self.info('Get mask for grid %s', cache_id)
mask = mask_cache.get(cache_id, None)
if mask is None:
self.notice('Loading mask: %s', dict(resolution=self.resolution, mode=self.mode, thresholds=self.thresholds, reverse=self.reverse))
mask = masking.polygon_mask(grid, self.resolution, mode=self.mode, thresholds=self.thresholds)
if self.reverse: mask = ~mask
mask_cache[cache_id] = mask
self.info(bases.describe(mask, stats=stats))
self.logger.info(bases.psinfo())
# Mask variable
# TODO: check/handle dimensions count and order
self.notice('Masking variable: %s, mask: %s', memvar.shape, mask.shape)
mask = MV2.resize(mask, filevar.shape)
memvar[:] = MV2.masked_where(mask, memvar)
self.info(bases.describe(memvar, stats=stats))
# Special scalar case which could fail if directly written (because fill_value is None)
if not filevar.shape:
fill_value = filevar.getMissing()
if fill_value is None:
fill_value = -memvar
memvar = cdms2.createVariable(
memvar, id=filevar.id, shape=(), typecode=filevar.typecode(),
fill_value=fill_value, attributes=filevar.attributes)
# Write masked variable to output file
self.notice('Writing variable to file')
outfile.write(memvar)
self.logger.info(bases.psinfo())
outfile.close()
inpfile.close()
def dump(self, restart_file=None):
"""Dump the current instance to a netcdf file"""
# File
netcdf4()
if restart_file is None:
restart_file = self.restart_file
if restart_file is None:
restart_file = self.default_restart_file
self.restart_file = restart_file
if os.path.exists(restart_file):
os.remove(restart_file)
f = cdms2.open(restart_file, 'w')
# Config
# - what was initially asked and some more
for sname in self.all_stats + ('sum', 'sqr', 'prod', 'stats'):
for st in 'st':
value = getattr(self, st+sname)
if value is None: continue
if isinstance(value, bool): value = int(value)
setattr(f, st+sname, value)
# - current status
f.iterindex = self.iterindex
f.nitems = int(self.nitems)
f.withtime = -1 if self.withtime is None else int(self.withtime)
if self.withtime and self.lasttime:
f.lasttime = str(self.lasttime)
if self.bins is None:
f.bin_edges = 0
else:
f.bin_edges = self.bins
if f.nitems==0: # Still no data
f.close()
return
# - already had some data
f.dual = int(self.dual)
f.ns = self.ns
f.nt = self.nt
f.nts = self._nts
f.tstats = int(self.tstats)
f.sstats = int(self.sstats)
# Spatial statistics
if self.sstats:
# Time axes
if self.withtime:
taxes = ()
for i, tt in enumerate(self._stimes):
taxis = MV2_axisConcatenate(tt)
taxis.stataccum_oldid = taxis.id
taxis.id = 't'+str(i)
taxes += taxis,
else:
taxes = [cdms2.createAxis(N.arange(self.nt), id='t')]
# Count
self._dump_array_(f, var=self._scount, id='scount', axis=taxes[0])
# Other stats
for key, stats in self._sstats.items():
if isinstance(stats, tuple):
for i, stat in enumerate(stats):
self._dump_array_(f, var=stat, id='s%s%s'%(key, str(i)),
axis=taxes[i])
else:
self._dump_array_(f, var=stats, id='s'+key, axis=taxes[0])
# Temporal statistics
if self.tstats:
# Main axis
axis = cdms2.createAxis(N.arange(self.ns), id='s')
# Count
self._dump_array_(f, var=self._tcount, id='tcount', axis=axis)
# Other stats
for key in self._dual_accums+self._single_accums:
if not getattr(self, 't'+key): continue
value = getattr(self, '_t'+key)
if key in self._dual_accums:
self._dump_array_(f, var=value, id='t'+key, axis=axis)
else:
for i, stat in enumerate(value):
self._dump_array_(f, var=stat, id='t%s%s'%(key, str(i)), axis=axis)
# Templates
ttemplates = [self._ttemplates]
# if self.thist:
# ttemplates.append(self._thtemplates)
for ttpls in ttemplates:
for i, ttpl in enumerate(ttpls):
# ttpl = ttpl.clone(copyData=0)
# for ia, axis in enumerate(ttpl.getAxisList()): #FIXME:grid cloning
# ttpl.setAxis(ia, axis.clone(copyData=0))
_add_id_prefix_(ttpl, 'var%i_'%i, exc=self._baxis)
f.write(ttpl)
_rm_id_prefix_(ttpl, 'var%i_'%i, exc=self._baxis)
# Attributes
for ivar, atts in enumerate(self._atts):
var = MV2.zeros(1)
set_atts(var, atts)
var.id = 'var%i_atts'%ivar
var.stataccum_id = atts['id']
var.getAxis(0).id = 'var_att_axis'
f.write(var)
# self._ttemplates, self._thtemplates, self._atts, self._tbase
f.close()
return f.id
def dump(self, restart_file=None):
"""Dump the current instance to a netcdf file"""
# File
netcdf4()
if restart_file is None:
restart_file = self.restart_file
if restart_file is None:
restart_file = self.default_restart_file
self.restart_file = restart_file
if os.path.exists(restart_file):
os.remove(restart_file)
f = cdms2.open(restart_file, 'w')
# Config
# - what was initially asked and some more
for sname in self.all_stats + ('sum', 'sqr', 'prod', 'stats'):
for st in 'st':
value = getattr(self, st+sname)
if value is None: continue
if isinstance(value, bool): value = int(value)
setattr(f, st+sname, value)
# - current status
f.iterindex = self.iterindex
f.nitems = int(self.nitems)
f.withtime = -1 if self.withtime is None else int(self.withtime)
if self.withtime and self.lasttime:
f.lasttime = str(self.lasttime)
if self.bins is None:
f.bin_edges = 0
else:
f.bin_edges = self.bins
if f.nitems==0: # Still no data
f.close()
return
# - already had some data
f.dual = int(self.dual)
f.ns = self.ns
f.nt = self.nt
f.nts = self._nts
f.tstats = int(self.tstats)
f.sstats = int(self.sstats)
# Spatial statistics
if self.sstats:
# Time axes
if self.withtime:
taxes = ()
for i, tt in enumerate(self._stimes):
taxis = MV2_axisConcatenate(tt)
taxis.stataccum_oldid = taxis.id
taxis.id = 't'+str(i)
taxes += taxis,
else:
taxes = [cdms2.createAxis(N.arange(self.nt), id='t')]
# Count
self._dump_array_(f, var=self._scount, id='scount', axis=taxes[0])
# Other stats
for key, stats in self._sstats.items():
if isinstance(stats, tuple):
for i, stat in enumerate(stats):
self._dump_array_(f, var=stat, id='s%s%s'%(key, str(i)),
axis=taxes[i])
else:
self._dump_array_(f, var=stats, id='s'+key, axis=taxes[0])
# Temporal statistics
if self.tstats:
# Main axis
axis = cdms2.createAxis(N.arange(self.ns), id='s')
# Count
self._dump_array_(f, var=self._tcount, id='tcount', axis=axis)
# Other stats
for key in self._dual_accums+self._single_accums:
if not getattr(self, 't'+key): continue
value = getattr(self, '_t'+key)
if key in self._dual_accums:
self._dump_array_(f, var=value, id='t'+key, axis=axis)
else:
for i, stat in enumerate(value):
self._dump_array_(f, var=stat, id='t%s%s'%(key, str(i)), axis=axis)
# Templates
for i, ttpl in enumerate(self._ttemplates):
ttpl = ttpl.clone()
_add_id_prefix_(ttpl, 'var%i_'%i, exc=self._baxis)
f.write(ttpl)
# Attributes
for ivar, atts in enumerate(self._atts):
var = MV2.zeros(1)
set_atts(var, atts)
var.id = 'var%i_atts'%ivar
var.stataccum_id = atts['id']
var.getAxis(0).id = 'var_att_axis'
f.write(var)
# self._ttemplates, self._thtemplates, self._atts, self._tbase
f.close()
return f.id
