# coding=utf-8 """Extract levels from variable""" from earthdiagnostics.box import Box from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticDomainOption, \ DiagnosticVariableListOption, DiagnosticIntOption from earthdiagnostics.utils import Utils, TempFile class SelectLevels(Diagnostic): """ Extract levels from file Parameters ---------- data_manager: DataManager startdate: str member: int chunk: init domain: ModellingRealm variable: str grid: str or None first_level: int last_level: int """ alias = 'selev' "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, domain, variable, grid, first_level, last_level): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.variable = variable self.domain = domain self.grid = grid self.box = Box(False) self.box.min_depth = first_level self.box.max_depth = last_level def __str__(self): return 'Select levels Startdate: {0} Member: {1} Chunk: {2} ' \ 'Variable: {3}:{4} Levels: {6}-{7} ' \ 'Grid: {5}'.format(self.startdate, self.member, self.chunk, self.domain, self.variable, self.grid, self.box.min_depth, self.box.max_depth) def __eq__(self, other): return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \ self.domain == other.domain and self.variable == other.variable and self.grid == self.grid @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: domain,variables,grid :type options: list[str] :return: """ options_available = (DiagnosticDomainOption(), DiagnosticVariableListOption(diags.data_manager.config.var_manager, 'variables'), DiagnosticIntOption('first_level'), DiagnosticIntOption('last_level'), DiagnosticOption('grid', '')) options = cls.process_options(options, options_available) job_list = list() variables = options['variables'] for var in variables: for startdate, member, chunk in diags.config.experiment.get_chunk_list(): job_list.append(SelectLevels(diags.data_manager, startdate, member, chunk, options['domain'], var, options['grid'], options['first_level'], options['last_level'])) 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, to_modify=True) def declare_data_generated(self): """Request data required by the diagnostic""" self.result = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, grid=self.grid) def compute(self): """Run the diagnostic""" temp = TempFile.get() Utils.nco.ncks(input=self.variable_file, output=temp, options=('-O -d lev,{0.min_depth},{0.max_depth}'.format(self.box),)) self.result.set_local_file(temp) @staticmethod def _create_var(var_name, var_values, source, destiny): old_var = source.variables[var_name] new_var = destiny.createVariable(var_name, old_var.dtype, dimensions=(var_name, )) new_var[:] = var_values Utils.copy_attributes(new_var, old_var) vertices_name = '{0}_vertices'.format(var_name) if vertices_name in source.variables: var_vertices = source.variables[vertices_name] if var_name == 'lon': vertices_values = var_vertices[0:1, ...] else: vertices_values = var_vertices[:, 0:1, :] new_lat_vertices = destiny.createVariable(vertices_name, var_vertices.dtype, dimensions=(var_name, 'vertices')) new_lat_vertices[:] = vertices_values Utils.copy_attributes(new_lat_vertices, var_vertices)