heatcontentlayer.py

# coding=utf-8
import numpy as np

from earthdiagnostics.box import Box
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics import cdo
from earthdiagnostics.utils import Utils, TempFile


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: 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, weight):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.box = box
        self.weight = weight
        self.required_vars = ['so', 'mlotst']
        self.generated_vars = ['scvertsum']

    def __str__(self):
        return 'Heat content layer Startdate: {0} Member: {1} Chunk: {2}'.format(self.startdate, self.member,
                                                                                 self.chunk)

    @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: minimum depth, maximum depth
        :type options: list[str]
        :return:
        """
        num_options = len(options) - 1
        if num_options < 2:
            raise Exception('You must specify the minimum and maximum depth to use')
        if num_options > 2:
            raise Exception('You must specify 2 for the heat content layer diagnostic')
        box = Box(True)
        box.min_depth = int(options[1])
        box.max_depth = int(options[2])
        job_list = list()

        handler = Utils.openCdf('mesh_zgr.nc')

        tmask = handler.variables['tmask'][:]

        if 'e3t' in handler.variables:
            tmask = handler.variables['e3t'][:] * tmask
        elif 'e3t_0' in handler.variables:
            tmask = handler.variables['e3t_0'][:] * tmask
        else:
            raise Exception('e3t variable can not be found')

        if 'gdept' in handler.variables:
            depth = handler.variables['gdept'][:]
        if 'gdept_0' in handler.variables:
            depth = handler.variables['gdept_0'][:]
        else:
            raise Exception('gdept_0 variable can not be found')
        handler.close()

        def calculate_weight(a):
            level = 0
            previous_level = 0

            while a[level] <= box.min_depth:
                previous_level = a[level]
                a[level] = 0
                level += 1
                if level >= a.size:
                    return a

            if previous_level != box.min_depth:
                weight = (a[level] - box.min_depth) / (a[level] - previous_level)
                previous_level = a[level]
                a[level] = weight
                level += 1
                if level >= a.size:
                    return a

            while a[level] <= box.max_depth:
                previous_level = a[level]
                a[level] = 1
                level += 1
                if level >= a.size:
                    return a

            if previous_level != box.max_depth:
                weight = (box.max_depth - previous_level) / (a[level] - previous_level)
                a[level] = weight
                level += 1
                if level >= a.size:
                    return a

            a[level:] = 0
            return a

        weight = tmask * np.apply_along_axis(calculate_weight, 1, depth)

        for startdate, member, chunk in diags.exp_manager.get_chunk_list():
            job_list.append(HeatContentLayer(diags.data_manager, startdate, member, chunk, box, weight))
        return job_list

    def compute(self):
        """
        Runs the diagnostic
        """
        nco = Utils.nco
        results = TempFile.get()

        thetao_file = self.data_manager.get_file('ocean', 'thetao', self.startdate, self.member, self.chunk)

        handler = Utils.openCdf(thetao_file)
        heatc_sl = np.sum(handler.variables['thetao'][:] * 1020 * 4000 * self.weight, 1)
        handler.sync()
        handler.renameVariable('thetao', 'heatc_sl')
        handler.close()

        nco.ncks(input=thetao_file, output=results, options='-O -v lon,lat,time')
        Utils.rename_variables(results, {'x': 'i', 'y': 'j'}, False, True)
        handler_results = Utils.openCdf(results)
        handler_results.createVariable('ohc', float, ('time', 'j', 'i'))
        handler_results.close()

        handler_results = Utils.openCdf(results)
        handler_results.variables['ohc'][:] = heatc_sl
        handler_results.close()

        Utils.setminmax(results, 'ohc')
        self.data_manager.send_file(results, 'ocean', 'ohc', self.startdate, self.member, self.chunk, box=self.box)