Source code for earthdiagnostics.ocean.cutsection

# coding=utf-8
"""Cut meridional or zonal sections"""
import numpy as np
from bscearth.utils.log import Log

from earthdiagnostics.box import Box
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBoolOption, DiagnosticIntOption, \
    DiagnosticDomainOption, DiagnosticVariableOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils


[docs]class CutSection(Diagnostic): """ Cuts a meridional or zonal section :original author: Virginie Guemas <virginie.guemas@bsc.es> :contributor: Javier Vegas-Regidor<javier.vegas@bsc.es> :created: September 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's name :type variable: str :param domain: variable's domain :type domain: Domain :param zonal: specifies if section is zonal or meridional :type zonal: bool :param value: value of the section's coordinate :type value: int """ alias = 'cutsection' "Diagnostic alias for the configuration file" def __init__(self, data_manager, startdate, member, chunk, domain, variable, zonal, value): Diagnostic.__init__(self, data_manager) self.startdate = startdate self.member = member self.chunk = chunk self.variable = variable self.domain = domain self.zonal = zonal self.value = value self.box = Box() if self.zonal: self.box.max_lon = self.value self.box.min_lon = self.value else: self.box.max_lat = self.value self.box.min_lat = self.value 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.variable == other.variable and self.zonal == other.zonal and \ self.value == other.value def __str__(self): return 'Cut section Startdate: {0} Member: {1} Chunk: {2} Variable: {3}:{4} ' \ 'Zonal: {5} Value: {6}'.format(self.startdate, self.member, self.chunk, self.domain, self.variable, self.zonal, self.value)
[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, zonal, value, domain=ocean :type options: list[str] :return: """ options_available = (DiagnosticVariableOption(diags.data_manager.config.var_manager), DiagnosticBoolOption('zonal'), DiagnosticIntOption('value'), DiagnosticDomainOption(default_value=ModelingRealms.ocean)) options = cls.process_options(options, options_available) job_list = list() for startdate, member, chunk in diags.config.experiment.get_chunk_list(): job_list.append(CutSection(diags.data_manager, startdate, member, chunk, options['domain'], options['variable'], options['zonal'], options['value'])) return job_list
[docs] 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)
[docs] def declare_data_generated(self): """Declare data to be generated by the diagnostic""" self.section = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk, box=self.box)
[docs] def compute(self): """Run the diagnostic""" nco = Utils.nco handler = Utils.open_cdf('mesh_hgr.nc') dimi = handler.dimensions['i'].size dimj = handler.dimensions['j'].size dimlev = handler.dimensions['lev'].size lon = handler.variables['lon'][:] lat = handler.variables['lat'][:] handler.close() handler = Utils.open_cdf('mask.nc') mask_lev = handler.variables['tmask'][:] mask_lev = mask_lev.astype(float) # noinspection PyTypeChecker np.place(mask_lev, mask_lev == 0, [1e20]) handler.close() # Latitude / longitude of the zonal / meridional section exactpos = self.value if not self.zonal: while exactpos < np.min(lon): exactpos += 360 while exactpos > np.max(lon): exactpos -= 360 size = dimj else: size = dimi # Collect the indexes defining the section listi = np.empty(size, dtype=int) listj = np.empty(size, dtype=int) for jpt in range(0, size): if not self.zonal: vector = lon[jpt, :] else: vector = lat[:, jpt] distance = abs(vector - exactpos) pos = np.where(distance == min(distance)) if not self.zonal: listi[jpt] = pos[0][0] listj[jpt] = jpt else: listi[jpt] = jpt listj[jpt] = pos[0][0] temp = self.data_manager.get_file(self.domain, self.variable, self.startdate, self.member, self.chunk) handler = Utils.open_cdf(temp) dimtime = handler.dimensions['time'].size var_array = handler.variables[self.variable][:] handler.close() var = np.empty([dimtime, dimlev, size], dtype=handler.variables[self.variable].dtype) new_coord = np.empty(size, dtype=float) if self.zonal: old_coord = lon else: old_coord = lat for jpt in range(0, size): var[:, :, jpt] = np.maximum(var_array[:, :, listj[jpt], listi[jpt]], mask_lev[:, :, listj[jpt], listi[jpt]]) new_coord[jpt] = old_coord[listj[jpt], listi[jpt]] nco.ncks(input=temp, output=temp, options=('-O -v lev,time',)) handler = Utils.open_cdf(temp) if not self.zonal: handler.createDimension('lat', size) coord_var = handler.createVariable('lat', float, 'lat') file_var = handler.createVariable(self.variable, float, ('time', 'lev', 'lat')) else: handler.createDimension('lon', size) coord_var = handler.createVariable('lon', float, 'lon') file_var = handler.createVariable(self.variable, float, ('time', 'lev', 'lon')) coord_var[:] = new_coord[:] file_var[:] = var[:] file_var.missing_value = 1e20 handler.close() self.section.set_local_file(temp) Log.info('Finished cut section for startdate {0}, member {1}, chunk {2}', self.startdate, self.member, self.chunk)