convectionsites.py

# coding=utf-8
import numpy as np
from autosubmit.config.log import Log
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.utils import Utils, TempFile
from earthdiagnostics.constants import Models


class ConvectionSites(Diagnostic):
    """
    Compute the intensity of convection in the four main convection sites

    :original author: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>

    :created: October 2013
    :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 nemo_version: NEMO's version
    :type nemo_version: str
    """

    def __init__(self, data_manager, startdate, member, chunk, nemo_version):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.nemo_version = nemo_version
        self.required_vars = ['vsftbarot']
        self.generated_vars = ['gyres']

    def __str__(self):
        return 'Convection sites Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member, self.chunk)

    @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: None
        :type options: list[str]
        :return:
        """
        if len(options) > 1:
            raise Exception('The convection sites diagnostic has no options')
        job_list = list()
        for startdate, member, chunk in diags.exp_manager.get_chunk_list():
            job_list.append(ConvectionSites(diags.data_manager, startdate, member, chunk, diags.nemo_version))
        return job_list

    def compute(self):
        """
        Runs the diagnostic
        """
        if self.nemo_version in [Models.ECEARTH_2_3_O1L42, Models.ECEARTH_3_0_O1L46,
                                 Models.NEMO_3_2_O1L42, Models.NEMO_3_3_O1L46,
                                 Models.NEMOVAR_O1L42]:
            labrador = [225, 245, 215, 255]
            irminger = [245, 290, 215, 245]
            gin = [260, 310, 245, 291]
            wedell = [225, 280, 1, 50]

        elif self.nemo_version in [Models.ECEARTH_3_0_O25L46, Models.ECEARTH_3_0_O25L75,
                                   Models.GLORYS2_V1_O25L75]:
            raise Exception("Option convection not available yet for {0}".format(self.nemo_version))
        else:
            raise Exception("Input grid {0} not recognized".format(self.nemo_version))

        mlotst_file = self.data_manager.get_file('ocean', 'mlotst', self.startdate, self.member, self.chunk)
        output = TempFile.get()

        self.mlotst_handler = Utils.openCdf(mlotst_file)
        handler = Utils.openCdf(output, 'w')
        handler.createDimension('time', self.mlotst_handler.variables['time'].shape[0])
        handler.createDimension('region', 4)
        Utils.copy_variable(self.mlotst_handler, handler, 'time')
        var_region = handler.createVariable('region', str, 'region')
        var_gyre = handler.createVariable('site', 'f', ('time', 'region'))
        var_gyre.short_name = 'site'
        var_gyre.long_name = 'convection sites'
        var_gyre.units = 'm^3/s'

        var_region[0] = 'labrador'
        var_gyre[:, 0] = self._convection_site(labrador)

        var_region[1] = 'irminger'
        var_gyre[:, 1] = self._convection_site(irminger)

        var_region[2] = 'gin'
        var_gyre[:, 2] = self._convection_site(gin)

        var_region[3] = 'wedell'
        var_gyre[:, 3] = self._convection_site(wedell)

        self.mlotst_handler.close()
        handler.close()

        self.data_manager.send_file(output, 'ocean', 'site', self.startdate, self.member, self.chunk)
        Log.info('Finished convection sites for startdate {0}, member {1}, chunk {2}',
                 self.startdate, self.member, self.chunk)

    def _convection_site(self, site):
        return np.max(self.mlotst_handler.variables['mlotst'][:, site[2]-1:site[3]-1, site[0]-1:site[1]-1], (1, 2))