siasiesiv.py

# coding=utf-8
import netCDF4
import os

from bscearth.utils.log import Log

from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinOption
from earthdiagnostics.utils import Utils, TempFile
# noinspection PyUnresolvedReferences
import earthdiagnostics.cdftoolspython as cdftoolspython
import numpy as np

from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.constants import Basins


class Siasiesiv(Diagnostic):
    """
    Compute the sea ice extent , area and volume  in both hemispheres or a specified region.


    :original author: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Neven Fuckar <neven.fuckar@bsc.es>
    :contributor: Ruben Cruz <ruben.cruzgarcia@bsc.es>
    :contributor: Javier Vegas-Regidor <javier.vegas@bsc.es>

    :created: April 2012
    :last modified: June 2016

    """
    alias = 'siasiesiv'
    "Diagnostic alias for the configuration file"

    e1t = None
    e2t = None
    gphit = None

    def __init__(self, data_manager, startdate, member, chunk, basin, mask):
        """
        :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 mask: mask to use
        :type mask: numpy.array
        """
        Diagnostic.__init__(self, data_manager)
        self.basin = basin
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.mask = mask
        self.generated = {}

    def __str__(self):
        return 'Siasiesiv Startdate: {0} Member: {1} Chunk: {2} Basin: {3}'.format(self.startdate, self.member,
                                                                                   self.chunk, self.basin)

    @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: basin
        :type options: list[str]
        :return:
        """
        options_available = (DiagnosticBasinOption('basin', Basins().Global), )
        options = cls.process_options(options, options_available)

        if options['basin'] is None:
            Log.error('Basin not recognized')
            return ()

        mask = np.asfortranarray(Utils.get_mask(options['basin']))

        job_list = list()
        for startdate, member, chunk in diags.config.experiment.get_chunk_list():
            job_list.append(Siasiesiv(diags.data_manager, startdate, member, chunk, options['basin'], mask))
        mesh_handler = Utils.openCdf('mesh_hgr.nc')
        Siasiesiv.e1t = np.asfortranarray(mesh_handler.variables['e1t'][0, :])
        Siasiesiv.e2t = np.asfortranarray(mesh_handler.variables['e2t'][0, :])
        Siasiesiv.gphit = np.asfortranarray(mesh_handler.variables['gphit'][0, :])
        mesh_handler.close()

        return job_list

    def request_data(self):
        self.sit = self.request_chunk(ModelingRealms.seaIce, 'sit', self.startdate, self.member, self.chunk)
        self.sic = self.request_chunk(ModelingRealms.seaIce, 'sic', self.startdate, self.member, self.chunk)

    def declare_data_generated(self):
        self._declare_var('sivols')
        self._declare_var('siareas')
        self._declare_var('siextents')

        self._declare_var('sivoln')
        self._declare_var('siarean')
        self._declare_var('siextentn')

    def _declare_var(self, var_name):
        self.generated[var_name] = self.declare_chunk(ModelingRealms.seaIce, var_name, self.startdate, self.member,
                                                      self.chunk, region=self.basin.fullname)

    def compute(self):
        """
        Runs the diagnostic
        """
        sit_handler = Utils.openCdf(self.sit.local_file)
        sit = np.asfortranarray(sit_handler.variables['sit'][:])
        timesteps = sit_handler.dimensions['time'].size
        sit_handler.close()

        sic_handler = Utils.openCdf(self.sic.local_file)
        Utils.convert_units(sic_handler.variables['sic'], '1.0')
        sic = np.asfortranarray(sic_handler.variables['sic'][:])
        sic_handler.close()

        result = np.empty((8, timesteps))
        for t in range(0, timesteps):
            result[:, t] = cdftoolspython.icediag.icediags(Siasiesiv.e1t, Siasiesiv.e2t, self.mask,
                                                           Siasiesiv.gphit, sit[t, :], sic[t, :])

        self._extract_variable_and_rename(result[4, :], 'sivols', '10^9 m3')
        self._extract_variable_and_rename(result[5, :], 'siareas', '10^9 m2')
        self._extract_variable_and_rename(result[7, :], 'siextents', '10^9 m2')

        self._extract_variable_and_rename(result[0, :], 'sivoln', '10^9 m3')
        self._extract_variable_and_rename(result[1, :], 'siarean', '10^9 m2')
        self._extract_variable_and_rename(result[3, :], 'siextentn', '10^9 m2')

    def _extract_variable_and_rename(self, values, cmor_name, units):
        temp = TempFile.get()
        reference_handler = Utils.openCdf(self.sit.local_file)
        os.remove(temp)
        handler = netCDF4.Dataset(temp, 'w')

        Utils.copy_variable(reference_handler, handler, 'time', add_dimensions=True)
        Utils.copy_variable(reference_handler, handler, 'time_bnds', add_dimensions=True, must_exist=False)
        Utils.copy_variable(reference_handler, handler, 'leadtime', add_dimensions=True, must_exist=False)
        reference_handler.close()

        new_var = handler.createVariable(cmor_name, float, 'time', fill_value=1.0e20)
        new_var.units = units
        new_var.short_name = cmor_name
        new_var.valid_min = 0.0
        new_var[:] = values
        new_var.valid_max = np.max(values)
        handler.close()
        self.generated[cmor_name].set_local_file(temp)