cutsection.py

# coding=utf-8
import numpy as np
from bscearth.utils.log import Log

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


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):
        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)

    @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, 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

    def request_data(self):
        self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk)

    def declare_data_generated(self):
        self.section = self.declare_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk,
                                          box=self.box)

    def compute(self):
        """
        Runs the diagnostic
        """
        nco = Utils.nco

        handler = Utils.openCdf('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.openCdf('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.openCdf(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.openCdf(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)