heatcontent.py

# coding=utf-8
import shutil

from earthdiagnostics import cdftools
from earthdiagnostics.constants import Basins
from earthdiagnostics.utils import Utils, TempFile
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticBasinOption, DiagnosticIntOption
from earthdiagnostics.box import Box
from earthdiagnostics.modelingrealm import ModelingRealms
import numpy as np


class HeatContent(Diagnostic):
    """
    Compute the total ocean heat content

    :original author: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>

    :created: May 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 mixed_layer: If 1, restricts calculation to the mixed layer, if -1 exclude it. If 0, no effect
    :type mixed_layer: int
    :param box: box to use for the average
    :type box: Box

    """

    alias = 'ohc'
    "Diagnostic alias for the configuration file"

    def __init__(self, data_manager, startdate, member, chunk, basin, mixed_layer, box, min_level, max_level):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.basin = basin
        self.mxloption = mixed_layer
        self.box = box
        self.min_level = min_level
        self.max_level = max_level
        self.required_vars = ['so', 'mlotst']
        self.generated_vars = ['scvertsum']

    def __eq__(self, other):
        return self.startdate == other.startdate and self.member == other.member and self.chunk == other.chunk and \
            self.box == other.box and self.basin == other.basin and self.mxloption == other.mxloption

    def __str__(self):
        return 'Heat content Startdate: {0} Member: {1} Chunk: {2} Mixed layer: {3} Box: {4} ' \
               'Basin: {5}'.format(self.startdate, self.member, self.chunk, self.mxloption, self.box,
                                   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, mixed layer option (1 to only compute at the mixed layer, -1 to exclude it, 0 to ignore),
                        minimum depth, maximum depth
        :type options: list[str]
        :return:
        """
        options_available = (DiagnosticBasinOption('basin'),
                             DiagnosticIntOption('mixed_layer', None, -1, 1),
                             DiagnosticIntOption('min_depth'),
                             DiagnosticIntOption('max_depth'))
        options = cls.process_options(options, options_available)
        box = Box(True)
        box.min_depth = options['min_depth']
        box.max_depth = options['max_depth']
        min_level = 0
        max_level = 0

        if box.min_depth > 0 or box.max_depth > 0:

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

            if 'gdepw_1d' in handler.variables:
                depth = handler.variables['gdepw_1d'][0, :]
            else:
                raise Exception('gdepw 1D variable can not be found')
            handler.close()

            def find_nearest(array, value):
                idx = (np.abs(array - value)).argmin()
                return idx, round(array[idx])

            if box.min_depth > 0:
                min_level, box.min_depth = find_nearest(depth, box.min_depth)
            else:
                min_level = 1

            if box.max_depth > 0:
                max_level, box.max_depth = find_nearest(depth, box.max_depth)
            else:
                max_level = len(depth)
                box.max_depth = round(depth[-1])

        job_list = list()
        for startdate, member, chunk in diags.config.experiment.get_chunk_list():
            job_list.append(HeatContent(diags.data_manager, startdate, member, chunk,
                                        options['basin'], options['mixed_layer'], box, min_level, max_level))
        return job_list

    def compute(self):
        """
        Runs the diagnostic
        """
        nco = Utils.nco
        temperature_file = self.data_manager.get_file(ModelingRealms.ocean, 'thetao', self.startdate,
                                                      self.member, self.chunk)
        if self.mxloption != 0:
            mlotst_file = self.data_manager.get_file(ModelingRealms.ocean, 'mlotst', self.startdate,
                                                     self.member, self.chunk)
            nco.ncks(input=mlotst_file, output=temperature_file, options='-A -v mlotst')

        para = list()
        if self.box.min_depth != 0:
            para.append('-zoom')
            para.append(0)
            para.append(0)
            para.append(0)
            para.append(0)
            para.append(self.min_level)
            para.append(self.max_level)
        if self.mxloption != 0:
            para.append('-mxloption')
            para.append(str(self.mxloption))
        if self.basin != Basins.Global:
            handler = Utils.openCdf('mask_regions.3d.nc')
            if self.basin.fullname not in handler.variables:
                raise Exception('Basin {0} is not defined on mask_regions.nc'.format(self.basin.fullname))

            handler.close()
            para.append('-M')
            para.append('mask_regions.3d.nc')
            para.append(self.basin.fullname)

        temp2 = TempFile.get()

        cdftools.run('cdfheatc', options=para, input=temperature_file, output=temp2, input_option='-f')

        results = Utils.openCdf(temp2)
        heatcsum_temp = TempFile.get()
        heatcvmean_temp = TempFile.get()
        nco.ncks(input=temperature_file, output=heatcsum_temp, options='-O -v time')
        shutil.copy(heatcsum_temp, heatcvmean_temp)

        heatcsum_handler = Utils.openCdf(heatcsum_temp)
        thc = heatcsum_handler.createVariable('heatcsum', float, 'time')
        thc.standard_name = "integral_of_sea_water_potential_temperature_expressed_as_heat_content"
        thc.long_name = "Total heat content"
        thc.units = "J"
        thc[:] = results.variables['heatc3d'][:, 0, 0]
        heatcsum_handler.close()

        heatcvmean_handler = Utils.openCdf(heatcvmean_temp)
        uhc = heatcvmean_handler.createVariable('heatcvmean', float, 'time')
        uhc.standard_name = "integral_of_sea_water_potential_temperature_expressed_as_heat_content"
        uhc.long_name = "Heat content per unit volume"
        uhc.units = "J*m^-3"
        uhc[:] = results.variables['heatc3dpervol'][:, 0, 0]
        heatcvmean_handler.close()

        results.close()

        if self.box.min_depth == 0:
            # For cdftools, this is all levels
            box_save = None
        else:
            box_save = self.box

        Utils.setminmax(heatcsum_temp, 'heatcsum')
        self.send_file(heatcsum_temp, ModelingRealms.ocean, 'heatcsum', self.startdate, self.member, self.chunk,
                       box=box_save, region=self.basin.fullname)
        Utils.setminmax(heatcvmean_temp, 'heatcvmean')
        self.send_file(heatcvmean_temp, ModelingRealms.ocean, 'heatcvmean', self.startdate, self.member, self.chunk,
                       box=box_save, region=self.basin.fullname)