# coding=utf-8 """Diagnostic to calculate a region total""" import os import iris import iris.util import iris.coords import iris.analysis import iris.exceptions import numpy as np import netCDF4 from earthdiagnostics import cdftools from earthdiagnostics.box import Box 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.regsum as regsum from diagonals.mesh_helpers.nemo import Nemo class RegionSum(Diagnostic): """ Computes the sum 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 :created: March 2017 :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 = 'regsum' "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, domain, variable, grid_point, box, save3d, basins, grid): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.domain = domain self.variable = variable self.grid_point = grid_point self.box = box self.save3d = save3d self.basins = basins self.grid = grid self.declared = {} self.lat_name = 'lat' self.lon_name = 'lon' def __eq__(self, other): if self._different_type(other): return False 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 and \ self.grid_point == other.grid_point and self.grid == other.grid \ and self.basin == other.basin def __str__(self): return 'Region sum Startdate: {0.startdate} Member: {0.member}' \ 'Chunk: {0.chunk} Variable: {0.variable} ' \ 'Grid point: {0.grid_point} Box: {0.box} Save 3D: {0.save3d}' \ 'Original grid: {0.grid} Basin: {0.basins}'.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: """ options_available = (DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager), DiagnosticOption('grid_point', 'T'), DiagnosticBasinListOption('basins', Basins().Global), DiagnosticIntOption('min_depth', -1), DiagnosticIntOption('max_depth', -1), DiagnosticIntOption('min_lat', -1), DiagnosticIntOption('max_lat', -1), DiagnosticIntOption('min_lon', -1), DiagnosticIntOption('max_lon', -1), DiagnosticBoolOption('save3D', True), 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_list.append(RegionSum(diags.data_manager, startdate, member, chunk, options['domain'], options['variable'], options['grid_point'].lower(), box, options['save3D'], options['basins'], options['grid'])) return job_list def request_data(self): """Request data required by the diagnostic""" self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, grid=self.grid) def declare_data_generated(self): """Declare data to be generated by the diagnostic""" if self.box.min_depth == -1 and self.box.max_depth == -1: box_save = None else: box_save = self.box self._declare_var('sum', False, box_save) self._declare_var('sum', True, box_save) def compute(self): """Run the diagnostic""" 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 -1 and self.box.max_lat is not -1 and self.box.min_lon is not -1 and \ self.box.max_lat is not -1: 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) if has_levels: self._sum_3d_variable(data, mesh, masks) else: self._sum_2d_var(data, mesh, masks) def _sum_2d_var(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) varsum = regsum.compute_regsum_2D(data.data, masks, areacello) self._save_result_2D('sum', varsum, data) def _sum_3d_variable(self, data, mesh, masks): areacello = mesh.get_areacello(cell_point=self.grid_point) tmask = iris.load_cube('mesh_hgr.nc', '{0}mask'.format(self.grid_point)) e3 = self._try_load_cube(3) e3 = self._rename_depth(e3) e3.coord('depth').bounds = data.coord('depth').bounds if self.box.min_depth is not -1 and self.box.max_depth is not -1: depth_constraint = iris.Constraint(depth=lambda c: self.box.min_depth <= c <= self.box.max_depth) e3 = e3.extract(depth_constraint) data = data.extract(depth_constraint) tmask = tmask.extract(depth_constraint) volcello = areacello*e3.data.astype(np.float32) varsum = regsum.compute_regsum_3D(data.data, masks, volcello, tmask.data) self._save_result_2D('sum', varsum, data) if self.save3d: varsum3d = regsum.compute_regsum_levels(data.data, masks, volcello, tmask) self._save_result_3D('sum', varsum3d, data) def _try_load_cube(self, number): try: cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}'.format(number, self.grid_point)) except iris.exceptions.ConstraintMismatchError: cube = iris.load_cube('mesh_hgr.nc', 'e{0}{1}_0'.format(number, self.grid_point)) cube = iris.util.squeeze(cube) dims = len(cube.shape) try: cube.coord('i') except iris.exceptions.CoordinateNotFoundError: cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 1]), var_name='i'), dims - 1) try: cube.coord('j') except iris.exceptions.CoordinateNotFoundError: cube.add_dim_coord(iris.coords.DimCoord(np.arange(cube.shape[dims - 2]), var_name='j'), dims - 2) return cube def _load_data(self): 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 return data 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) else: final_name = '{1}{0}'.format(var, self.variable) self.declared[final_name] = self.declare_chunk( ModelingRealms.ocean, final_name, self.startdate, self.member, self.chunk, box=box_save, region=self.basins ) def _send_var(self, threed, mean_file): var = 'sum' if threed: if not self.save3d: return False original_name = '{0}_{1}'.format(var, self.variable) final_name = '{1}3d{0}'.format(var, self.variable) levels = ',lev' else: original_name = '{0}_3D{1}'.format(var, self.variable) final_name = '{1}{0}'.format(var, self.variable) levels = '' temp2 = TempFile.get() Utils.nco().ncks( input=mean_file, output=temp2, options=('-v {0},{2.lat_name},{2.lon_name}{1}'.format(original_name, levels, self),) ) handler = Utils.open_cdf(temp2) var_handler = handler.variables[original_name] if hasattr(var_handler, 'valid_min'): del var_handler.valid_min if hasattr(var_handler, 'valid_max'): del var_handler.valid_max handler.close() self.declared[final_name].set_local_file( temp2, diagnostic=self, rename_var=original_name, region=self.basins ) 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) else: final_name = '{1}{0}'.format(var, self.variable) self.declared[final_name] = self.declare_chunk( ModelingRealms.ocean, final_name, self.startdate, self.member, self.chunk, box=box_save, region=self.basins, grid=self.grid ) 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}3d{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) handler_temp.createDimension('lev', data.shape[1]) var_level = handler_temp.createVariable('lev', float, 'lev') var_level[...] = data.coord('depth').points var_level.units = 'm' var_level.axis = 'Z' var_level.positive = 'down' var_level.long_name = 'ocean depth coordinate' var_level.standard_name = 'depth' var_region = handler_temp.createVariable('region', 'S1', ('region', 'region_length')) var = handler_temp.createVariable('{1}3d{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)