siasiesiv.py

# coding=utf-8
import netCDF4
import os
from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.utils import Utils, TempFile
import cdftoolspython
import numpy as np


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

    """
    e1t = None
    e2t = None
    gphit = None

    def __init__(self, data_manager, basin, startdate, member, chunk, 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.required_vars = ['sit', 'sic']
        self.generated_vars = ['siextents', 'sivols', 'siareas', 'siextentn', 'sivoln', 'siarean']

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

    @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:
        """
        if len(options) != 2:
            raise Exception('You must specify the basin for the siasiesiv diagnostic (and nothing else)')
        basin = Basins.parse(options[1])

        if basin != Basins.Global:
            mask_handler = Utils.openCdf('mask_regions.nc')
            mask = mask_handler.variables[basin.fullname][:, 0, :]
            mask_handler.close()
        else:
            mask_handler = Utils.openCdf('mask.nc')
            mask = np.asfortranarray(mask_handler.variables['tmask'][0, 0, :])
            mask_handler.close()

        job_list = list()
        for startdate, member, chunk in diags.exp_manager.get_chunk_list():
            job_list.append(Siasiesiv(diags.data_manager, basin, startdate, member, chunk, 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 compute(self):
        """
        Runs the diagnostic
        """
        sit_file = self.data_manager.get_file('seaIce', 'sit', self.startdate, self.member, self.chunk)
        sit_handler = Utils.openCdf(sit_file)
        sit = np.asfortranarray(sit_handler.variables['sit'][:])
        timesteps = sit_handler.dimensions['time'].size
        sit_handler.close()

        sic_file = self.data_manager.get_file('seaIce', 'sic', self.startdate, self.member, self.chunk)
        sic_handler = Utils.openCdf(sic_file)
        sic = np.asfortranarray(sic_handler.variables['sic'][:])
        sic_handler.close()

        result = np.empty((8, timesteps))
        for t in range(0, timesteps):
            try:

                result[:, t] = cdftoolspython.icediag.icediags(Siasiesiv.e1t, Siasiesiv.e2t, self.mask,
                                                               Siasiesiv.gphit, sit[t, :], sic[t, :])
            except Exception as ex:
                print ex

        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, result[4, :], 'sivols',
                                                                      "10^3 km3", "10^9 m3"),
                                    'seaIce', 'sivols', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)
        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, result[5, :], 'siareas',
                                                                      "10^6 km2", "10^9 m2"),
                                    'seaIce', 'siareas', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)
        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, result[7, :], 'siextents',
                                                                      "10^6 km2", "10^9 m2"),
                                    'seaIce', 'siextents', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)

        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, result[0, :], 'sivoln',
                                                                      "10^3 km3", "10^9 m3"),
                                    'seaIce', 'sivoln', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)
        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, result[1, :], 'siarean',
                                                                      "10^6 km2", "10^9 m2"),
                                    'seaIce', 'siarean', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)
        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, result[3, :], 'siextentn',
                                                                      "10^6 km2", "10^9 m2"),
                                    'seaIce', 'siextentn', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)

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

        # Create dimensions
        handler.createDimension('time')
        handler.createDimension('bnds', 2)

        # Copy time variable

        Utils.copy_variable(reference_handler, handler, 'time')
        Utils.copy_variable(reference_handler, handler, 'time_bnds')
        Utils.copy_variable(reference_handler, handler, 'leadtime')
        reference_handler.close()

        new_var = handler.createVariable(cmor_name, float, 'time', fill_value=0.0)
        factor = self._get_conversion_factor(target_units, output_units)
        values *= factor
        new_var[:] = values
        new_var.units = output_units
        new_var.short_name = cmor_name
        new_var.valid_min = 0.0
        new_var.valid_max = np.max(values)
        handler.close()
        return temp

    def _get_conversion_factor(self, input_units, output_units):
        units = input_units.split()
        if len(units) == 1:
            scale_unit = 1
            unit = units[0]
        else:
            if '^' in units[0]:
                values = units[0].split('^')
                scale_unit = pow(int(values[0]), int(values[1]))
            else:
                scale_unit = float(units[0])
            unit = units[1]

        units = output_units.split()
        if len(units) == 1:
            scale_new_unit = 1
            new_unit = units[0]
        else:
            if '^' in units[0]:
                values = units[0].split('^')
                scale_new_unit = pow(int(values[0]), int(values[1]))
            else:
                scale_new_unit = float(units[0])
            new_unit = units[1]

        factor = self._get_factor(new_unit, unit)
        invert = False
        if factor is None:
            factor = self._get_factor(unit, new_unit)
            invert = True

        if factor is None:
            raise Exception("Conversion from {0} to {1} not supported".format(input_units, output_units))

        if invert:
            factor = scale_unit / float(scale_new_unit * factor)
        else:
            factor = (factor * scale_unit) / float(scale_new_unit)
        return factor

    @staticmethod
    def _get_factor(new_unit, unit):
        # Add  only the conversions with a factor greater than 1
        if unit == new_unit:
            return 1
        if unit == 'km3':
            if new_unit == 'm3':
                return pow(1000, 3)
        elif unit == 'km2':
            if new_unit == 'm2':
                return pow(1000, 2)
        return None