from earthdiagnostics import cdftools from earthdiagnostics.box import Box from earthdiagnostics.diagnostic import Diagnostic from earthdiagnostics.utils import Utils, TempFile class VerticalMean(Diagnostic): """ Chooses vertical level in ocean, or vertically averages between 2 or more ocean levels :original author: Virginie Guemas :contributor: Eleftheria Exarchou :contributor: Javier Vegas-Regidor :created: February 2012 :last modified: June 2016 """ 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 self.required_vars = [variable] self.generated_vars = [variable + 'vmean'] def __str__(self): return 'Vertical mean Startdate: {0} Member: {1} Chunk: {2} Variable: {3} ' \ 'Box: {4}'.format(self.startdate, self.member, self.chunk, self.variable, self.box) @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: """ num_options = len(options) - 1 if num_options < 1: raise Exception('You must specify the variable to average vertically') if num_options > 3: raise Exception('You must specify between one and three parameters for the vertical mean') variable = options[1] box = Box(True) if num_options >= 2: box.min_depth = float(options[2]) if num_options >= 2: box.max_depth = float(options[3]) job_list = list() for startdate in diags.startdates: for member in diags.members: for chunk in range(1, diags.chunks + 1): job_list.append(VerticalMean(diags.data_manager, startdate, member, chunk, variable, box)) return job_list def compute(self): temp = TempFile.get() variable_file = self.data_manager.get_file('ocean', self.variable, self.startdate, self.member, self.chunk) handler = Utils.openCdf(variable_file) if self.box.min_depth is None: lev_min = handler.variables['lev'][0] else: lev_min = self.box.min_depth if self.box.max_depth is None: lev_max = handler.variables['lev'][-1] else: lev_max = self.box.max_depth handler.close() cdftools.run('cdfvertmean', input=variable_file, output=temp, options=[self.variable, 'T', lev_min, lev_max, '-debug']) Utils.setminmax(temp, '{0}_vert_mean'.format(self.variable)) self.data_manager.send_file(temp, 'ocean', self.variable + 'vmean', self.startdate, self.member, self.chunk, box=self.box, rename_var='{0}_vert_mean'.format(self.variable))