Newer
Older
"""Diagnostic to compute regional averages"""
import iris.coords
import iris.exceptions
from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticIntOption, DiagnosticDomainOption, \
DiagnosticBoolOption, DiagnosticBasinListOption, DiagnosticVariableOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile
import diagonals.regmean as regmean
from diagonals.mesh_helpers.nemo import Nemo
class RegionMean(Diagnostic):
"""
Computes the mean value of the field (3D, weighted).
For 3D fields, a horizontal mean for each level is also given. If a spatial window
is specified, the mean value is computed only in this window.
:original author: Javier Vegas-Regidor <javier.vegas@bsc.es>
:param data_manager: data management object
:type data_manager: DataManager
:param startdate: startdate
:type startdate: str
:param member: member number
:type member: int
:param chunk: chunk's number
:type chunk: int
:param variable: variable to average
:type variable: str
:param box: box used to restrict the vertical mean
:type box: Box
"""
alias = 'regmean'
"Diagnostic alias for the configuration file"
def __init__(self, data_manager, startdate, member, chunk, domain, variable, box, save3d,
Diagnostic.__init__(self, data_manager)
self.startdate = startdate
self.member = member
self.chunk = chunk
self.domain = domain
self.variable = variable
self.box = box
self.lat_name = 'lat'
self.lon_name = 'lon'
return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \
self.box == other.box and self.variable == other.variable
def __str__(self):
return 'Region mean Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} Variable: {0.variable} ' \
'Box: {0.box} Save 3D: {0.save3d} Save variance: {0.variance} Grid point: {0.grid_point}'.format(self)
def __hash__(self):
return hash(str(self))
@classmethod
def generate_jobs(cls, diags, options):
"""
Create a job for each chunk to compute the diagnostic
:param diags: Diagnostics manager class
:type diags: Diags
:param options: variable, minimum depth (level), maximum depth (level)
:type options: list[str]
:return:
"""
DiagnosticVariableOption(diags.data_manager.config.var_manager),
DiagnosticIntOption('min_depth', -1),
DiagnosticIntOption('max_depth', -1),
DiagnosticIntOption('min_lat', None),
DiagnosticIntOption('max_lat', None),
DiagnosticIntOption('min_lon', None),
DiagnosticIntOption('max_lon', None),
DiagnosticBoolOption('variance', False),
DiagnosticOption('grid', ''))
options = cls.process_options(options, options_available)
box = Box()
box.min_depth = options['min_depth']
box.max_depth = options['max_depth']
box.min_lat = options['min_lat']
box.max_lat = options['max_lat']
box.min_lon = options['min_lon']
box.max_lon = options['max_lon']
basins = options['basins']
if not basins:
Log.error('Basins not recognized')
return()
job_list = list()
for startdate, member, chunk in diags.config.experiment.get_chunk_list():
job = RegionMean(diags.data_manager, startdate, member, chunk,
options['domain'], options['variable'], box,
options['save3D'], options['variance'], options['basins'],
job_list.append(job)
self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk)
if self.box.min_depth == -1 and self.box.max_depth == -1:
box_save = None
else:
box_save = self.box
self._declare_var('mean', False, box_save)
self._declare_var('mean', True, box_save)
has_levels = self._fix_file_metadata()
data = self._load_data()
masks = {}
self.basins.sort()
for basin in self.basins:
masks[basin] = Utils.get_mask(basin)
mesh = Nemo('mesh_hgr.nc', 'mask_regions.nc')
if self.box.min_lat is not None and self.box.max_lat is not None and self.box.min_lon is not None and self.box.max_lat is not None:
name = '{0}_{1}'.format(Box.get_lat_str(self.box), Box.get_lon_str(self.box))
masks[name] = mesh.get_region_mask(self.box.min_lat,
self.box.max_lat,
self.box.min_lon,
self.box.max_lon)
self._mean_2d_var(data, mesh, masks)
def _mean_2d_var(self, data, mesh, masks):
areacello = mesh.get_areacello()
mean = regmean.compute_regmean_2D(data.data, masks, areacello)
self._save_result_2D('mean', mean, data)
def _meand_3d_variable(self, data, mesh, masks):
areacello = mesh.get_areacello()
e3 = self._rename_depth(e3)
depth_constraint = iris.Constraint(depth=lambda c: self.box.min_depth <= c <= self.box.max_depth)
volcello = areacello*e3.extract(depth_constraint).data
mean = regmean.compute_regmean_3D(data.data, masks, volcello)
self._save_result_2D('mean', mean, data)
if self.save3d:
mean3d = regmean.compute_regmean_levels(data.data, masks, volcello)
self._save_result_3D('mean3d', mean3d, data)
coords = []
handler = Utils.open_cdf(self.variable_file.local_file)
for variable in handler.variables:
if variable in ('time', 'lev', 'lat', 'lon', 'latitude', 'longitude', 'leadtime', 'time_centered'):
coords.append(variable)
if variable == 'time_centered':
handler.variables[variable].standard_name = ''
handler.variables[self.variable].coordinates = ' '.join(coords)
handler.close()
data = iris.load_cube(self.variable_file.local_file)
return self._rename_depth(data)
def _rename_depth(self, data):
for coord_name in ('model_level_number', 'Vertical T levels', 'lev'):
if data.coords(coord_name):
coord = data.coord(coord_name)
coord.standard_name = 'depth'
coord.long_name = 'depth'
break
def _fix_file_metadata(self):
handler = Utils.open_cdf(self.variable_file.local_file)
var = handler.variables[self.variable]
coordinates = ''
has_levels = False
for dimension in handler.variables.keys():
if dimension in ['time', 'lev', 'lat', 'latitude', 'lon', 'longitude', 'i', 'j']:
coordinates += ' {0}'.format(dimension)
if dimension == 'lev':
has_levels = True
var.coordinates = coordinates
handler.close()
return has_levels
def _declare_var(self, var, threed, box_save):
if threed:
if not self.save3d:
return False
final_name = '{1}3d{0}'.format(var, self.variable)
final_name = '{1}{0}'.format(var, self.variable)
self.declared[final_name] = self.declare_chunk(ModelingRealms.ocean, final_name, self.startdate, self.member,
def _send_var(self, var, threed, cube_list):
if not self.save3d and threed:
return False
final_name = '{1}3d{0}'.format(var, self.variable)
final_name = '{1}{0}'.format(var, self.variable)
cube = cube_list.merge_cube()
cube.remove_coord('latitude')
cube.remove_coord('longitude')
try:
cube.remove_coord('depth')
except iris.exceptions.CoordinateNotFoundError:
pass
try:
cube.remove_coord('lev')
except iris.exceptions.CoordinateNotFoundError:
pass
temp = TempFile.get()
iris.save(cube, temp)
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
self.declared[final_name].set_local_file(temp, diagnostic=self, rename_var='result', region=self.basins)
def _save_result_2D(self, var, result, data):
final_name = '{1}{0}'.format(var, self.variable)
temp = TempFile.get()
handler_source = Utils.open_cdf(self.variable_file.local_file)
handler_temp = Utils.open_cdf(temp, 'w')
Utils.copy_variable(handler_source, handler_temp, 'time', True, True)
handler_temp.createDimension('region', len(result))
handler_temp.createDimension('region_length', 50)
var_region = handler_temp.createVariable('region', 'S1',
('region', 'region_length'))
var = handler_temp.createVariable('{1}{0}'.format(var, self.variable),
float, ('time', 'region',),)
var.units = '{0}'.format(data.units)
for i, basin in enumerate(result):
var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50)
var[..., i] = result[basin]
handler_temp.close()
self.declared[final_name].set_local_file(temp, diagnostic=self)
def _save_result_3D(self, var, result, data):
final_name = '{1}{0}'.format(var, self.variable)
temp = TempFile.get()
handler_source = Utils.open_cdf(self.variable_file.local_file)
handler_temp = Utils.open_cdf(temp, 'w')
Utils.copy_variable(handler_source, handler_temp, 'time', True, True)
Utils.copy_variable(handler_source, handler_temp, 'lev', True, True)
handler_temp.createDimension('region', len(result))
handler_temp.createDimension('region_length', 50)
var_region = handler_temp.createVariable('region', 'S1',
('region', 'region_length'))
var = handler_temp.createVariable('{1}{0}'.format(var, self.variable),
float, ('time', 'lev', 'region',),)
var.units = '{0}'.format(data.units)
for i, basin in enumerate(result):
var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50)
var[..., i] = result[basin]
handler_temp.close()
self.declared[final_name].set_local_file(temp, diagnostic=self)