heatcontentlayer.py

# coding=utf-8
"""Point-wise Ocean Heat Content in a specified ocean thickness (J/m-2)"""
import numpy as np

from earthdiagnostics.box import Box
from earthdiagnostics.diagnostic import (
    Diagnostic,
    DiagnosticIntOption,
    DiagnosticBasinListOption,
)
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile
from bscearth.utils.log import Log

from diagonals import ohc
from diagonals.mesh_helpers.nemo import Nemo

import netCDF4


class HeatContentLayer(Diagnostic):
    """
    Point-wise Ocean Heat Content in a specified ocean thickness (J/m-2)

    :original author: Isabel Andreu Burillo
    :contributor: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Eleftheria Exarchou <eleftheria.exarchou@bsc.es>
    :contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>

    :created: June 2012
    :last modified: June 2016

    :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 box: box to use for the calculations
    :type box: Box
    """

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

    def __init__(
        self,
        data_manager,
        startdate,
        member,
        chunk,
        box,
        areas,
        weight,
        layers,
        basins,
        data_convention,
        min_level,
        max_level,
    ):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.box = box
        self.areas = areas
        self.weight = weight
        self.layers = layers
        self.basins = basins
        self.data_convention = data_convention
        self.min_level = min_level
        self.max_level = max_level

    def __str__(self):
        return (
            f"Heat content layer Startdate: {self.startdate} "
            f"Member: {self.member} Chunk: {self.chunk} "
            f"Box: {self.box}"
        )

    @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: minimum depth, maximum depth, basin=Global
        :type options: list[str]
        """
        options_available = (
            DiagnosticIntOption("min_depth"),
            DiagnosticIntOption("max_depth"),
            DiagnosticBasinListOption("basins", "global"),
        )
        options = cls.process_options(options, options_available)

        box = Box(True)
        box.min_depth = options["min_depth"]
        box.max_depth = options["max_depth"]
        basins = options["basins"]
        if not basins:
            Log.error("Basins not recognized")
            return ()

        mesh = Nemo("mesh_hgr.nc", "mask_regions.nc")

        layers = ((box.min_depth, box.max_depth),)

        masks = {}
        basins.sort()
        for basin in basins:
            masks[basin] = Utils.get_mask(basin)

        areacello = mesh.get_areacello()
        areas = ohc.get_basin_area(areacello, masks)

        job_list = list()

        e3t = mesh.get_k_length()
        mask = mesh.get_landsea_mask()
        depth = mesh.get_depth(cell_point="W")
        weight = ohc.get_weights(layers, mask, e3t, depth)
        max_level, min_level = cls._get_used_levels(weight)
        weight[0] = weight[0][:, min_level:max_level, :, :]

        del mask, depth, e3t

        for (
            startdate,
            member,
            chunk,
        ) in diags.config.experiment.get_chunk_list():
            job_list.append(
                HeatContentLayer(
                    diags.data_manager,
                    startdate,
                    member,
                    chunk,
                    box,
                    areas,
                    weight,
                    layers,
                    basins,
                    diags.config.data_convention,
                    min_level,
                    max_level,
                )
            )
        return job_list

    @classmethod
    def _get_mesh_info(cls):
        handler = Utils.open_cdf("mesh_zgr.nc")
        # mask = Utils.get_mask(options['basin'])
        mask = handler.variables["tmask"][:]
        if "e3t" in handler.variables:
            e3t = handler.variables["e3t"][:]
        elif "e3t_0" in handler.variables:
            e3t = handler.variables["e3t_0"][:]
        else:
            raise Exception("e3t variable can not be found")
        if "gdepw" in handler.variables:
            depth = handler.variables["gdepw"][:]
        elif "gdepw_0" in handler.variables:
            depth = handler.variables["gdepw_0"][:]
        else:
            raise Exception("gdepw variable can not be found")
        e3t_3d = e3t.shape != depth.shape
        if e3t_3d:
            mask = e3t_3d * mask
        else:
            e3t = e3t[0, :]
        while len(depth.shape) < 4:
            depth = np.expand_dims(depth, -1)
        handler.close()
        return depth, e3t, mask

    def request_data(self):
        """Request data required by the diagnostic"""
        self.thetao = self.request_chunk(
            ModelingRealms.ocean,
            "thetao",
            self.startdate,
            self.member,
            self.chunk,
        )

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        self.heatc = self.declare_chunk(
            ModelingRealms.ocean,
            "heatc",
            self.startdate,
            self.member,
            self.chunk,
            box=self.box,
        )
        self.heatcsum = self.declare_chunk(
            ModelingRealms.ocean,
            "heatcsum",
            self.startdate,
            self.member,
            self.chunk,
            box=self.box,
        )

    @classmethod
    def _get_used_levels(cls, weight):
        # Now we will reduce to the levels with any weight != 0 to avoid
        # loading too much data on memory
        levels = weight[0].shape[1]
        min_level = 0
        while min_level < levels and not weight[0][:, min_level, :].any():
            min_level += 1
        max_level = min_level
        while (
            max_level < (levels - 1) and weight[0][:, max_level + 1, :].any()
        ):
            max_level += 1
        return max_level, min_level

    def compute(self):
        """Run the diagnostic"""
        thetao_file = TempFile.get()
        results = TempFile.get()
        results_1d = TempFile.get()
        Utils.copy_file(self.thetao.local_file, thetao_file)

        handler = Utils.open_cdf(thetao_file)
        Utils.convert_units(handler.variables["thetao"], "K")
        heatc_sl, heatc_sl_1d = ohc.compute(
            self.layers,
            self.weight,
            handler.variables["thetao"][
                :, self.min_level:self.max_level, :, :
            ],
            self.areas,
        )
        handler.sync()
        handler.renameVariable("thetao", "heatc_sl")

        results = TempFile.get()
        handler_results = Utils.open_cdf(results, "w")
        lat_name = next(
            alias
            for alias in ("lat", "latitude")
            if alias in handler.variables.keys()
        )
        lon_name = next(
            alias
            for alias in ("lon", "longitude")
            if alias in handler.variables.keys()
        )
        Utils.copy_variable(handler, handler_results, "time", True, True)
        Utils.copy_variable(handler, handler_results, "i", False, True)
        Utils.copy_variable(handler, handler_results, "j", False, True)
        Utils.copy_variable(handler, handler_results, lat_name, True, True)
        Utils.copy_variable(handler, handler_results, lon_name, True, True)
        var = handler_results.createVariable(
            "heatc", float, ("time", "j", "i"), fill_value=1.0e20
        )
        var.units = "J m-2" ""
        var.coordinates = " ".join((lat_name, lon_name))
        handler_results.sync()
        # temporary fix, needs to loop over layers
        handler_results.variables["heatc"][:] = heatc_sl[0]
        handler_results.close()

        results_1d = TempFile.get()
        handler_results_1d = Utils.open_cdf(results_1d, "w")
        Utils.copy_variable(handler, handler_results_1d, "time", True, True)
        handler_results_1d.createDimension("region", len(self.basins))
        handler_results_1d.createDimension("region_length", 50)
        var_region = handler_results_1d.createVariable(
            "region", "S1", ("region", "region_length")
        )
        var_ohc_1d = handler_results_1d.createVariable(
            "heatcsum", float, ("time", "region",),
        )
        handler_results_1d.sync()
        for i, basin in enumerate(self.basins):
            var_region[i, ...] = netCDF4.stringtoarr(basin.name, 50)
            var_ohc_1d[..., i] = heatc_sl_1d[i]
        handler_results_1d.close()
        var_ohc_1d.units = "J" ""
        Utils.setminmax(results, "heatc")
        self.heatc.set_local_file(results)
        self.heatcsum.set_local_file(results_1d)