regionmean.py

# coding=utf-8
"""Diagnostic to compute regional averages"""
import iris
import iris.exceptions

from earthdiagnostics.box import Box
from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticOption, DiagnosticIntOption, DiagnosticDomainOption, \
    DiagnosticBoolOption, DiagnosticBasinOption, DiagnosticVariableOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile
from bscearth.utils.log import Log


class RegionMean(Diagnostic):
    """
    Computes the mean value of the field (3D, weighted).

    For 3D fields, a horizontal mean for each level is also given. If a spatial window
    is specified, the mean value is computed only in this window.

    :original author: Javier Vegas-Regidor <javier.vegas@bsc.es>

    :created: March 2017

    :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 variable: variable to average
    :type variable: str
    :param box: box used to restrict the vertical mean
    :type box: Box
    """

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

    def __init__(self, data_manager, startdate, member, chunk, domain, variable, box, save3d, weights_file,
                 variance, basin):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.domain = domain
        self.variable = variable
        self.box = box
        self.save3d = save3d
        self.weights_file = weights_file
        self.variance = variance
        self.basin = basin

        self.declared = {}

        self.lat_name = 'lat'
        self.lon_name = 'lon'

    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.variable == other.variable

    def __str__(self):
        return 'Region mean Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} Variable: {0.variable} ' \
               'Box: {0.box} Save 3D: {0.save3d} Save variance: {0.variance}'.format(self)

    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: variable, minimum depth (level), maximum depth (level)
        :type options: list[str]
        :return:
        """
        options_available = (DiagnosticDomainOption(),
                             DiagnosticVariableOption(diags.data_manager.config.var_manager),
                             DiagnosticOption('grid_point', 'T'),
                             DiagnosticBasinOption('basin', Basins().Global),
                             DiagnosticIntOption('min_depth', 0),
                             DiagnosticIntOption('max_depth', 0),
                             DiagnosticBoolOption('save3D', True),
                             DiagnosticBoolOption('variance', False),
                             DiagnosticOption('grid', ''))
        options = cls.process_options(options, options_available)

        box = Box()
        box.min_depth = options['min_depth']
        box.max_depth = options['max_depth']

        weights_file = TempFile.get()
        weight_diagnostics = ComputeWeights(diags.data_manager, options['grid_point'],  options['basin'], weights_file)

        job_list = list()
        for startdate, member, chunk in diags.config.experiment.get_chunk_list():
            job = RegionMean(diags.data_manager, startdate, member, chunk,
                             options['domain'], options['variable'], box,
                             options['save3D'], weights_file, options['variance'], options['basin'])
            job.add_subjob(weight_diagnostics)
            job_list.append(job)

        return job_list

    def request_data(self):
        """Request data required by the diagnostic"""
        self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk)

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        if self.box.min_depth == 0:
            # To cdftools, this means all levels
            box_save = None
        else:
            box_save = self.box

        self._declare_var('mean', False, box_save)
        self._declare_var('mean', True, box_save)

        if self.variance:
            self._declare_var('var', False, box_save)
            self._declare_var('var', True, box_save)

    def compute(self):
        """Run the diagnostic"""
        data = iris.load_cube(self.variable_file, self.variable)
        weights = iris.load_cube(self.weights_file, 'weights')
        mean = data.collapsed(['latitude', 'longitude'], iris.analysis.MEAN, weights=weights)


    def _declare_var(self, var, threed, box_save):
        if threed:
            if not self.save3d:
                return False
            final_name = '{1}3d{0}'.format(var, self.variable)
        else:
            final_name = '{1}{0}'.format(var, self.variable)

        self.declared[final_name] = self.declare_chunk(ModelingRealms.ocean, final_name, self.startdate, self.member,
                                                       self.chunk, box=box_save, region=self.basin)


class ComputeWeights(Diagnostic):
    """
    Diagnostic used to compute regional mean and sum weights

    Parameters
    ----------
    data_manager: DataManager
    startdate: str
    member: int
    chunk: int
    weights_file: str
    """

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

    @classmethod
    def generate_jobs(cls, diags, options):
        pass

    def __init__(self, data_manager, grid_point, basin, weights_file):
        Diagnostic.__init__(self, data_manager)
        self.weights_file = weights_file
        self.basin = basin
        self.grid_point = grid_point.lower()

    def __str__(self):
        return 'Computing weights for region averaging: Point {0.grid_point} Basin: {0.basin}'.format(self)

    def __hash__(self):
        return hash(str(self))

    def compute(self):
        """Compute weights"""
        mask = Utils.get_mask(self.basin)
        e1 = iris.load_cube('mesh_hgr.nc', 'e1{0}'.format(self.grid_point))
        e2 = iris.load_cube('mesh_hgr.nc', 'e2{0}'.format(self.grid_point))
        try:
            e3 = iris.load_cube('mesh_hgr.nc', 'e3{0}'.format(self.grid_point))
        except iris.exceptions.ConstraintMismatchError:
            e3 = iris.load_cube('mesh_hgr.nc', 'e3{0}_0'.format(self.grid_point))
        weights = e3 * e1 * e2 * mask
        weights.var_name = 'weights'
        Log.info(str(weights))
        iris.save(weights, self.weights_file)

    def request_data(self):
        """Request data required by the diagnostic"""
        None

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        pass