siasiesiv.py

import netCDF4
import os
from earthdiagnostics import cdftools
from earthdiagnostics.basins 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

    """
    def __init__(self, data_manager, basin, startdate, member, chunk):
        Diagnostic.__init__(self, data_manager)
        self.basin = basin
        if basin == Basins.Global:
            self._region = None
        else:
            self._region = basin.shortname
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        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])
        job_list = list()
        for startdate in diags.startdates:
            for member in diags.members:
                for chunk in range(1, diags.chunks + 1):
                    job_list.append(Siasiesiv(diags.data_manager, basin, startdate, member, chunk))
        return job_list

    def compute(self):
        nco = Utils.nco

        mesh_handler = Utils.openCdf('mesh_hgr.nc')
        e1t = mesh_handler.variables['e1t'][0, :]
        e2t = mesh_handler.variables['e2t'][0, :]
        gphit = mesh_handler.variables['gphit'][0, :]
        mesh_handler.close()
        if self.basin != Basins.Global:
            mask_handler = Utils.openCdf('mask_regions.nc')
            mask = mask_handler.variables[self.basin.fullname][:, 0, :]
            mask_handler.close()
        else:
            mask_handler = Utils.openCdf('mask.nc')
            mask = mask_handler.variables['tmask'][0, 0, :]
            mask_handler.close()

        sit_file = self.data_manager.get_file('seaIce', 'sit', self.startdate, self.member, self.chunk)
        sic_file = self.data_manager.get_file('seaIce', 'sic', self.startdate, self.member, self.chunk)
        sit_handler = Utils.openCdf(sit_file)
        sic_handler = Utils.openCdf(sic_file)
        results = np.zeros()
        for t in range(0, sit_handler.dimensions['time'].size):
            try:
                sit = sit_handler.variables['sit'][t, :]
                sic = sic_handler.variables['sic'][t, :]
                results.append(cdftoolspython.icediag.icediags(e1t, e2t, mask, gphit, sit, sic))
            except Exception as ex:
                print ex

        sit_handler.close()
        sic_handler.close()

        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, results[0], 'sivols',
                                                                      "10^3 km3"),
                                    'seaIce', 'sivols', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)
        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, results[0], 'siareas',
                                                                      "10^6 km2"),
                                    'seaIce', 'siareas', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)
        self.data_manager.send_file(self._extract_variable_and_rename(sit_file, results[0], 'siextents',
                                                                      "10^6 km2"),
                                    'seaIce', 'siextents', self.startdate, self.member, self.chunk,
                                    region=self.basin.fullname)

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

        os.remove(temp)

    def _extract_variable_and_rename(self, reference_file, values, cmor_name, output_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('10^9 m3', 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(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