cutsection.py

# coding=utf-8
import numpy as np
from autosubmit.config.log import Log

from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.box import Box
from earthdiagnostics.utils import Utils


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: str
    :param zonal: specifies if section is zonal or meridional
    :type zonal: bool
    :param value: value of the section's coordinate
    :type value: int
    """

    def __init__(self, data_manager, startdate, member, chunk, variable, domain, 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

    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:
        """
        num_options = len(options) - 1
        if num_options < 3:
            raise Exception('You must specify the variable, coordinate and coordinate value')
        if num_options > 4:
            raise Exception('You must specify between 3 and 4 parameters for the interpolation diagnostic')
        variable = options[1]
        zonal = options[2].lower() == 'true'
        value = int(options[3])
        if num_options >= 4:
            domain = options[4]
        else:
            domain = 'ocean'

        job_list = list()
        for startdate, member, chunk in diags.exp_manager.get_chunk_list():
            job_list.append(CutSection(diags.data_manager, startdate, member, chunk, variable, domain, zonal, value))
        return job_list

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

        box = Box()
        if self.zonal:
            box.max_lon = self.value
            box.min_lon = self.value
        else:
            box.max_lat = self.value
            box.min_lat = self.value

        self.data_manager.send_file(temp, self.domain, self.variable, self.startdate, self.member, self.chunk,
                                    box=box)
        Log.info('Finished cut section for startdate {0}, member {1}, chunk {2}',
                 self.startdate, self.member, self.chunk)