siasiesiv.py

# coding=utf-8
"""Compute the sea ice extent , area and volume  in both hemispheres or a specified region"""
import os

import netCDF4
import numpy as np
from bscearth.utils.log import Log

from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinOption, DiagnosticBoolOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile


# noinspection PyUnresolvedReferences


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

    Parameters
    ----------
    data_manager: DataManager
    startdate: str
    member: int
    chunk: init
    domain: ModellingRealm
    variable: str
    basin: Basin
    mask: numpy.array
    omit_vol: bool
    """

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

    e1t = None
    e2t = None
    gphit = None

    def __init__(self, data_manager, startdate, member, chunk, basin, mask, var_manager, omit_vol):
        Diagnostic.__init__(self, data_manager)
        self.basin = basin
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.mask = mask
        self.generated = {}
        self.var_manager = var_manager
        self.omit_volume = omit_vol
        self.sic_varname = self.var_manager.get_variable('sic').short_name
        self.sit_varname = self.var_manager.get_variable('sit').short_name

    def __str__(self):
        return 'Siasiesiv Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} ' \
               'Basin: {0.basin} Omit volume: {0.omit_volume}'.format(self)

    @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: basin
        :type options: list[str]
        :return:
        """
        options_available = (DiagnosticBasinOption('basin', Basins().Global),
                             DiagnosticBoolOption('omit_volume', False))
        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,
                                      diags.config.var_manager, options['omit_volume']))
        mesh_handler = Utils.open_cdf('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):
        """Request data required by the diagnostic"""
        if not self.omit_volume:
            self.sit = self.request_chunk(ModelingRealms.seaIce, self.sit_varname,
                                          self.startdate, self.member, self.chunk)
        self.sic = self.request_chunk(ModelingRealms.seaIce, self.sic_varname,
                                      self.startdate, self.member, self.chunk)

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        if not self.omit_volume:
            self._declare_var('sivols')
            self._declare_var('sivoln')

        self._declare_var('siareas')
        self._declare_var('siextents')

        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.name)

    def compute(self):
        """Run the diagnostic"""
        import earthdiagnostics.cdftoolspython as cdftoolspython

        sic_handler = Utils.open_cdf(self.sic.local_file)
        Utils.convert_units(sic_handler.variables[self.sic_varname], '1.0')
        sic = np.asfortranarray(sic_handler.variables[self.sic_varname][:])
        timesteps = sic_handler.dimensions['time'].size
        sic_handler.close()

        if self.omit_volume:
            sit = sic
        else:
            sit_handler = Utils.open_cdf(self.sit.local_file)
            sit = np.asfortranarray(sit_handler.variables[self.sit_varname][:])
            sit_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[5, :], 'siareas', '10^9 m2')
        self._extract_variable_and_rename(result[7, :], 'siextents', '10^9 m2')

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

        if not self.omit_volume:
            self._extract_variable_and_rename(result[4, :], 'sivols', '10^9 m3')
            self._extract_variable_and_rename(result[0, :], 'sivoln', '10^9 m3')

    def _extract_variable_and_rename(self, values, cmor_name, units):
        temp = TempFile.get()
        reference_handler = Utils.open_cdf(self.sic.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)