# coding=utf-8 import os from earthdiagnostics import cdftools from earthdiagnostics.box import Box from earthdiagnostics.constants import Basins from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticIntOption, DiagnosticDomainOption, \ DiagnosticBoolOption, DiagnosticBasinOption, DiagnosticVariableOption from earthdiagnostics.modelingrealm import ModelingRealms from earthdiagnostics.utils import Utils, TempFile class RegionSum(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 :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, basin, 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.upper() self.box = box self.save3d = save3d self.basin = basin self.grid = grid self.declared = {} self.lat_name = 'lat' self.lon_name = 'lon' def __eq__(self, other): 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 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.basin}'.format(self) @classmethod def generate_jobs(cls, diags, options): """ Creates 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'), DiagnosticBasinOption('basin', Basins().Global), DiagnosticIntOption('min_depth', 0), DiagnosticIntOption('max_depth', 0), 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'] 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'], box, options['save3D'], options['basin'], options['grid'])) return job_list def request_data(self): self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, grid=self.grid) def declare_data_generated(self): if self.box.min_depth == 0: # To cdftools, this means all levels 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): """ Runs the diagnostic """ mean_file = TempFile.get() variable_file = self.variable_file.local_file handler = Utils.openCdf(variable_file) self.save3d &= 'lev' in handler.dimensions if "latitude" in handler.variables: self.lat_name = 'latitude' if "longitude" in handler.variables: self.lon_name = 'longitude' handler.close() cdfmean_options = ['-v', self.variable, '-p', self.grid_point, '-zoom', 0, 0, 0, 0, self.box.min_depth, self.box.max_depth] if self.basin != Basins().Global: cdfmean_options.append('-M') cdfmean_options.append('mask_regions.3d.nc') cdfmean_options.append(self.basin.name) cdftools.run('cdfsum', input_file=variable_file, input_option='-f', output_file=mean_file, options=cdfmean_options) Utils.rename_variables(mean_file, {'gdept': 'lev', 'gdepw': 'lev'}, must_exist=False, rename_dimension=True) self.send_var(False, mean_file) self.send_var(True, mean_file) os.remove(mean_file) 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.openCdf(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.basin) 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.basin, grid=self.grid)