Source code for earthdiagnostics.general.module

# coding=utf-8
"""Compute module of two variables"""
import numpy as np

from earthdiagnostics.diagnostic import Diagnostic, DiagnosticVariableOption, DiagnosticDomainOption, DiagnosticOption
from earthdiagnostics.utils import Utils, TempFile


[docs]class Module(Diagnostic): """ Compute the module of the vector given by two scalar variables :original author: Javier Vegas-Regidor<javier.vegas@bsc.es> :created: July 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 domain: variable's domain :type domain: ModelingRealm """ alias = 'module' "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, domain, componentu, componentv, module_var, grid): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.domain = domain self.componentu = componentu self.componentv = componentv self.module = module_var self.grid = grid self.original_values = None def __str__(self): return 'Calculate module Startdate: {0} Member: {1} Chunk: {2} ' \ 'Variables: {3}:{4},{5},{6} ' \ 'Grid: {7}'.format(self.startdate, self.member, self.chunk, self.domain, self.componentu, self.componentv, self.module, self.grid) 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.domain == other.domain and self.componentu == other.componentu and \ self.componentv == other.componentv and self.module == other.module and self.grid == other.grid
[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, domain, grid :type options: list[str] :return: """ options_available = (DiagnosticDomainOption(), DiagnosticVariableOption(diags.data_manager.config.var_manager, 'componentu'), DiagnosticVariableOption(diags.data_manager.config.var_manager, 'componentv'), DiagnosticVariableOption(diags.data_manager.config.var_manager, 'module'), DiagnosticOption('grid', '')) options = cls.process_options(options, options_available) job_list = list() for startdate, member, chunk in diags.config.experiment.get_chunk_list(): job_list.append(Module(diags.data_manager, startdate, member, chunk, options['domain'], options['componentu'], options['componentv'], options['module'], options['grid'])) return job_list
[docs] def request_data(self): """Request data required by the diagnostic""" self.component_u_file = self.request_chunk(self.domain, self.componentu, self.startdate, self.member, self.chunk, grid=self.grid) self.component_v_file = self.request_chunk(self.domain, self.componentv, self.startdate, self.member, self.chunk, grid=self.grid)
[docs] def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.module_file = self.declare_chunk(self.domain, self.module, self.startdate, self.member, self.chunk, grid=self.grid)
[docs] def compute(self): """Run the diagnostic""" temp = TempFile.get() Utils.copy_file(self.component_u_file.local_file, temp) component_u = Utils.open_cdf(temp) component_v = Utils.open_cdf(self.component_v_file.local_file) variable_u = component_u.variables[self.componentu] variable_v = component_v.variables[self.componentv] variable_u[:] = np.sqrt(variable_u[:] ** 2 + variable_v[:] ** 2) if 'table' in variable_u.ncattrs(): del variable_u.table if 'code' in variable_u.ncattrs(): del variable_u.code component_u.close() component_v.close() self.module_file.set_local_file(temp, rename_var=self.componentu)