siasiesiv.py

# coding=utf-8
"""Compute the sea ice extent, area and volume"""
import numpy as np

import netCDF4

import iris
import iris.analysis
import iris.coords
import iris.util
from bscearth.utils.log import Log

from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import (
    Diagnostic,
    DiagnosticBasinListOption,
    DiagnosticBoolOption,
)
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile

import diagonals.siasie as siasie
from diagonals.mesh_helpers.nemo import Nemo


# 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: list of Basin
    mask: numpy.array
    omit_vol: bool
    """

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

    def __init__(
        self,
        data_manager,
        startdate,
        member,
        chunk,
        masks,
        var_manager,
        data_convention,
        omit_vol,
    ):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.masks = masks
        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
        self.data_convention = data_convention

        self.results = {}
        for var in ("siarean", "siareas", "siextentn", "siextents"):
            self.results[var] = {}

    def __str__(self):
        return (
            "Siasiesiv Startdate: {0.startdate} Member: {0.member} "
            "Chunk: {0.chunk} Basins: {1} Omit volume: {0.omit_volume}".format(
                self, ",".join(str(basin) for basin in self.masks.keys())
            )
        )

    def __hash__(self):
        return hash(str(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 = (
            DiagnosticBasinListOption("basins", [Basins().Global]),
            DiagnosticBoolOption("omit_volume", True),
        )
        options = cls.process_options(options, options_available)

        basins = options["basins"]
        if not basins:
            Log.error("Basins not recognized")
            return ()

        masks = {}
        basins.sort()
        for basin in basins:
            masks[basin] = Utils.get_mask(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,
                    masks,
                    diags.config.var_manager,
                    diags.config.data_convention,
                    options["omit_volume"],
                )
            )

        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,
        )

    def compute(self):
        """Run the diagnostic"""

        self._fix_coordinates_attribute(self.sic.local_file, self.sic_varname)
        sic = iris.load_cube(self.sic.local_file)
        if sic.units.origin != "%":
            sic.convert_units("%")

        sic_slices = []
        for sic_data in sic.slices_over("time"):
            sic_data.data = np.ma.filled(sic_data.data, 0.0).astype(np.float32)
            sic_data.data[np.isinf(sic_data.data)] = 0.0
            sic_slices.append(sic_data)
        mesh = Nemo("mesh_hgr.nc", "mask_regions.nc")
        areacello = mesh.get_areacello(cell_point="T")
        gphit = mesh.get_grid_latitude(cell_point="T")

        if not self.omit_volume:
            sit = iris.load_cube(self.sit.local_file)
            sit_slices = []
            for sit_data in sit.slices_over("time"):
                sit_data.data = np.ma.filled(sit_data.data, 0.0).astype(
                    np.float32
                )
                sit_data.data[np.isinf(sit_data.data)] = 0.0
                sit_slices.append(sit_data)
            results = siasie.compute(
                gphit, areacello, sic_slices, self.masks, sit_slices
            )
            self.results["siextentn"] = results[0]
            self.results["siextents"] = results[1]
            self.results["siarean"] = results[2]
            self.results["siareas"] = results[3]
            self.results["sivoln"] = results[4]
            self.results["sivols"] = results[5]
        else:
            results = siasie.compute(
                gphit, areacello, sic_slices, self.masks, None
            )
            self.results["siextentn"] = results[0]
            self.results["siextents"] = results[1]
            self.results["siarean"] = results[2]
            self.results["siareas"] = results[3]
        self.save()

    def _fix_coordinates_attribute(self, filepath, var_name):
        add_coordinates = {
            "time",
            "leadtime",
            "time_centered",
            self.data_convention.lon_name,
            self.data_convention.lat_name,
        }
        handler = Utils.open_cdf(filepath)
        coordinates = handler.variables[var_name].coordinates.split()
        handler.variables[var_name].coordinates = " ".join(
            set(coordinates) | add_coordinates
        )
        handler.close()

    def save(self):
        for var in self.results.keys():
            res = self.results[var]
            temp = TempFile.get()
            handler_source = Utils.open_cdf(self.sic.local_file)
            handler_temp = Utils.open_cdf(temp, "w")
            Utils.copy_variable(
                handler_source, handler_temp, "time", True, True
            )
            handler_temp.createDimension("region", len(self.masks))
            handler_temp.createDimension("region_length", 50)
            var_region = handler_temp.createVariable(
                "region", "S1", ("region", "region_length")
            )
            var_res = handler_temp.createVariable(
                "{0}".format(var), float, ("time", "region",)
            )
            if var in ("sivoln", "sivols"):
                var_res.units = "m^3"
            else:
                var_res.units = "m^2"
            for i, basin in enumerate(self.masks):
                var_region[i, ...] = netCDF4.stringtoarr(str(basin), 50)
                var_res[..., i] = res[i, ...]
            handler_temp.close()
            self.generated[var].set_local_file(temp, diagnostic=self)