regionmean.py

# coding=utf-8
import os

from earthdiagnostics import cdftools
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


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, grid_point, box, save3d, basin,
                 variance, grid):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.domain = domain
        self.variable = variable
        self.grid_point = grid_point.upper()
        self.box = box
        self.save3d = save3d
        self.basin = basin
        self.variance = variance
        self.grid = grid
        self.declared = {}

    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} ' \
               'Grid point: {0.grid_point} Box: {0.box} Save 3D: {0.save3d} Save variance: {0.variance} ' \
               'Original grid: {0.grid}'.format(self)

    @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: 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']

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

    def request_data(self):
        self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk,
                                                grid=self.grid)

    def declare_data_generated(self):
        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):
        """
        Runs the diagnostic
        """
        mean_file = TempFile.get()

        variable_file = self.variable_file.local_file

        handler = Utils.openCdf(variable_file)
        self.save3d &= 'lev' in handler.dimensions
        handler.close()

        cdfmean_options = [self.variable, self.grid_point, 0, 0, 0, 0, self.box.min_depth, self.box.max_depth]
        if self.variance:
            cdfmean_options += ['-var']
        if self.basin != Basins().Global:
            cdfmean_options.append('-M')
            cdfmean_options.append('mask_regions.3d.nc')
            cdfmean_options.append(self.basin.name)

        cdftools.run('cdfmean', input=variable_file, output=mean_file, options=cdfmean_options)
        Utils.rename_variables(mean_file, {'gdept': 'lev', 'gdepw': 'lev'}, must_exist=False, rename_dimension=True)

        self.send_var('mean', False, mean_file)
        self.send_var('mean', True, mean_file)

        if self.variance:
            self.send_var('var', False, mean_file)
            self.send_var('var', True, mean_file)

        os.remove(mean_file)

    def send_var(self, var, threed, mean_file):
        if threed:
            if not self.save3d:
                return False
            original_name = '{0}_{1}'.format(var, self.variable)
            final_name = '{1}3d{0}'.format(var, self.variable)
            levels = ',lev'
        else:
            original_name = '{0}_3D{1}'.format(var, self.variable)
            final_name = '{1}{0}'.format(var, self.variable)
            levels = ''

        temp2 = TempFile.get()
        Utils.nco.ncks(input=mean_file, output=temp2, options=('-v {0},lat,lon{1}'.format(original_name, levels),))
        self.declared[final_name].set_local_file(temp2, rename_var=original_name)

    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, grid=self.grid)