Source code for earthdiagnostics.ocean.verticalgradient

# coding=utf-8
"""Calculate the gradient between 2 ocean levels"""
from earthdiagnostics.box import Box
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticVariableOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile


[docs]class VerticalGradient(Diagnostic): """ Calculate the gradient between 2 ocean levels :original author: Virginie Guemas <virginie.guemas@bsc.es> :contributor: Eleftheria Exarchou <eleftheria.exarchou@bsc.es> :contributor: Javier Vegas-Regidor <javier.vegas@bsc.es> :created: February 2012 :last modified: June 2016 :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 = 'vgrad' "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, variable, box): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.variable = variable self.box = box 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 def __str__(self): return 'Vertical gradient Startdate: {0} Member: {1} Chunk: {2} Variable: {3} ' \ 'Box: {4}'.format(self.startdate, self.member, self.chunk, self.variable, self.box)
[docs] @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 = (DiagnosticVariableOption(diags.data_manager.config.var_manager), DiagnosticIntOption('upper_level', 1), DiagnosticIntOption('low_level', 2)) options = cls.process_options(options, options_available) box = Box(False) if options['upper_level'] >= 0: box.min_depth = options['upper_level'] if options['low_level'] >= 0: box.max_depth = options['low_level'] job_list = list() for startdate, member, chunk in diags.config.experiment.get_chunk_list(): job_list.append(VerticalGradient(diags.data_manager, startdate, member, chunk, options['variable'], box)) return job_list
[docs] def request_data(self): """Request data required by the diagnostic""" self.variable_file = self.request_chunk(ModelingRealms.ocean, self.variable, self.startdate, self.member, self.chunk)
[docs] def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.gradient_file = self.declare_chunk(ModelingRealms.ocean, self.variable + 'vgrad', self.startdate, self.member, self.chunk, box=self.box)
[docs] def compute(self): """Run the diagnostic""" handler = Utils.open_cdf(self.variable_file) if 'lev' not in handler.dimensions: raise Exception('Variable {0} does not have a level dimension') var_handler = handler.variables[self.variable] upper_level = var_handler[:, self.box.min_depth - 1, ...] lower_level = var_handler[:, self.box.max_depth - 1, ...] gradient = upper_level - lower_level temp = TempFile.get() new_file = Utils.open_cdf(temp, 'w') for var in handler.variables.keys(): if var in (self.variable, 'lev', 'lev_bnds'): continue Utils.copy_variable(handler, new_file, var, add_dimensions=True) new_var = new_file.createVariable(self.variable + 'vgrad', var_handler.dtype, dimensions=('time', 'j', 'i'), zlib=True) Utils.copy_attributes(new_var, var_handler) new_var[...] = gradient[...] new_var.long_name += ' Vertical gradient' new_var.standard_name += '_vertical_gradient' self.gradient_file.set_local_file(temp)