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.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticBasinListOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile
from bscearth.utils.log import Log

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

import iris
import iris.coords

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 'Heat content layer Startdate: {0} Member: {1} Chunk: {2} Box: {3}'.format(self.startdate, self.member,
                                                                                          self.chunk, 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()
        results1D = 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_sl1D = 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.e20)
        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()

        results1D = TempFile.get()
        handler_results1D = Utils.open_cdf(results1D, 'w')
        Utils.copy_variable(handler, handler_results1D, 'time', True, True)
        handler_results1D.createDimension('region', len(self.basins))
        handler_results1D.createDimension('region_length', 50)
        var_region = handler_results1D.createVariable(
            'region', 'S1', ('region', 'region_length'))
        var_ohc1D = handler_results1D.createVariable(
            'heatcsum', float, ('time', 'region',),)
        handler_results1D.sync()
        for i, basin in enumerate(self.basins):
            var_region[i, ...] = netCDF4.stringtoarr(basin.name, 50)
            var_ohc1D[..., i] = heatc_sl1D[i]
        handler_results1D.close()

        Utils.setminmax(results, 'heatc')
        self.heatc.set_local_file(results)
        self.heatcsum.set_local_file(results1D)