timemean.py

# coding=utf-8
"""Time mean diagnostics"""
import iris
import iris.coord_categorisation
import iris.analysis
import iris.exceptions

import numpy as np
from bscearth.utils.log import Log

from earthdiagnostics.diagnostic import (
    Diagnostic,
    DiagnosticOption,
    DiagnosticDomainOption,
    DiagnosticFrequencyOption,
    DiagnosticVariableOption,
)
from earthdiagnostics.frequency import Frequencies
from earthdiagnostics.utils import TempFile, Utils


class TimeMean(Diagnostic):
    """
    Base class for all time mean diagnostics

    :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's name
    :type variable: str
    :param domain: variable's domain
    :type domain: ModelingRealm
    :param frequency: original frequency
    :type frequency: str
    :param grid: original data grid
    :type grid: str
    """

    def __init__(
        self,
        data_manager,
        startdate,
        member,
        chunk,
        domain,
        variable,
        frequency,
        grid,
    ):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.variable = variable
        self.domain = domain
        self.frequency = frequency
        self.grid = grid
        self._target_frequency = None

        self.variable_file = None
        self.mean_file = None

    def __str__(self):
        return (
            "Calculate {0._target_frequency} mean "
            "Startdate: {0.startdate} Member: {0.member} Chunk: {0.chunk} "
            "Variable: {0.domain}:{0.variable} "
            "Original frequency: {0.frequency} Grid: {0.grid}".format(self)
        )

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

    def __eq__(self, other):
        if self._different_type(other):
            return False
        return (
            self.startdate == other.startdate
            and self.member == other.member
            and self.chunk == other.chunk
            and self.domain == other.domain
            and self.variable == other.variable
            and self.frequency == other.frequency
            and self.grid == other.grid
            and self._target_frequency == other._target_frequency
        )

    @classmethod
    def _process_options(cls, diags, options):
        options_available = (
            DiagnosticDomainOption(),
            DiagnosticVariableOption(diags.data_manager.config.var_manager),
            DiagnosticFrequencyOption(),
            DiagnosticOption("grid", ""),
        )
        options = cls.process_options(options, options_available)
        return options

    @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, domain, frequency=day, grid=''
        :type options: list[str]
        :return:
        """
        options = cls._process_options(diags, options)
        job_list = list()
        chunk_list = diags.config.experiment.get_chunk_list()
        for startdate, member, chunk in chunk_list:
            job_list.append(
                cls(
                    diags.data_manager,
                    startdate,
                    member,
                    chunk,
                    options["domain"],
                    options["variable"],
                    options["frequency"],
                    options["grid"],
                )
            )
        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,
            frequency=self.frequency,
            grid=self.grid,
        )

    def compute_mean(self, cube):
        """
        Compute the time mean

        Parameters
        ----------
        cube: iris.cube.Cube

        Returns
        -------
        iris.cube.Cube
        """
        raise NotImplementedError()

    def compute(self):
        """Run the diagnostic"""
        temp = TempFile.get()
        Log.info('Load data')
        cube = iris.load_cube(self.variable_file.local_file)
        time_centered = [
            coord
            for coord in cube.coords()
            if coord.var_name == "time_centered"
        ]
        if time_centered:
            cube.remove_coord(time_centered[0])
        iris.coord_categorisation.add_day_of_month(cube, "time")
        iris.coord_categorisation.add_month_number(cube, "time")
        iris.coord_categorisation.add_year(cube, "time")
        Log.info('Compute mean')
        cube = self.compute_mean(cube)
        cube.remove_coord("day_of_month")
        cube.remove_coord("month_number")
        cube.remove_coord("year")
        try:
            region_coord = cube.coord("region")
            cube.remove_coord(region_coord)
        except iris.exceptions.CoordinateNotFoundError:
            region_coord = None
        cube.data
        iris.save(cube, temp)
        if region_coord:
            handler = Utils.open_cdf(temp)
            region = handler.createVariable("region", str, ("dim0",))
            region.standard_name = region_coord.standard_name
            region[...] = region_coord.points.astype(np.dtype(str))

            handler.variables[self.variable].coordinates += " region"
            handler.close()

            Utils.rename_variable(temp, "dim0", "region", False)

        self.mean_file.set_local_file(temp)


class DailyMean(TimeMean):
    """
    Calculates daily mean for a given variable

    :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's name
    :type variable: str
    :param domain: variable's domain
    :type domain: ModelingRealm
    :param frequency: original frequency
    :type frequency: str
    :param grid: original data grid
    :type grid: str
    """

    alias = "daymean"
    "Diagnostic alias for the configuration file"

    def __init__(
        self,
        data_manager,
        startdate,
        member,
        chunk,
        domain,
        variable,
        frequency,
        grid,
    ):
        TimeMean.__init__(
            self,
            data_manager,
            startdate,
            member,
            chunk,
            domain,
            variable,
            frequency,
            grid,
        )
        self._target_frequency = "daily"

    def compute_mean(self, cube):
        """
        Compute the time mean

        Parameters
        ----------
        cube: iris.cube.Cube

        Returns
        -------
        iris.cube.Cube
        """
        return cube.aggregated_by(
            ["day_of_month", "month_number", "year"], iris.analysis.MEAN
        )

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        self.mean_file = self.declare_chunk(
            self.domain,
            self.variable,
            self.startdate,
            self.member,
            self.chunk,
            frequency=Frequencies.daily,
            grid=self.grid,
        )


class MonthlyMean(TimeMean):
    """
    Calculates monthly mean for a given variable

    :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's name
    :type variable: str
    :param domain: variable's domain
    :type domain: ModelingRealm
    :param frequency: original frequency
    :type frequency: str
    :param grid: original data grid
    :type grid: str
    """

    alias = "monmean"
    "Diagnostic alias for the configuration file"

    def __init__(
        self,
        data_manager,
        startdate,
        member,
        chunk,
        domain,
        variable,
        frequency,
        grid,
    ):
        TimeMean.__init__(
            self,
            data_manager,
            startdate,
            member,
            chunk,
            domain,
            variable,
            frequency,
            grid,
        )
        self._target_frequency = "monthly"

    def compute_mean(self, cube):
        """
        Compute the time mean

        Parameters
        ----------
        cube: iris.cube.Cube

        Returns
        -------
        iris.cube.Cube
        """
        return cube.aggregated_by(["month_number", "year"], iris.analysis.MEAN)

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        self.mean_file = self.declare_chunk(
            self.domain,
            self.variable,
            self.startdate,
            self.member,
            self.chunk,
            frequency=Frequencies.monthly,
            grid=self.grid,
        )


class YearlyMean(TimeMean):
    """
    Calculates monthly mean for a given variable

    :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's name
    :type variable: str
    :param domain: variable's domain
    :type domain: ModelingRealm
    :param frequency: original frequency
    :type frequency: str
    :param grid: original data grid
    :type grid: str
    """

    alias = "yearmean"
    "Diagnostic alias for the configuration file"

    def __init__(
        self,
        data_manager,
        startdate,
        member,
        chunk,
        domain,
        variable,
        frequency,
        grid,
    ):
        TimeMean.__init__(
            self,
            data_manager,
            startdate,
            member,
            chunk,
            domain,
            variable,
            frequency,
            grid,
        )
        self._target_frequency = "yearly"

    def compute_mean(self, cube):
        """
        Compute the time mean

        Parameters
        ----------
        cube: iris.cube.Cube

        Returns
        -------
        iris.cube.Cube
        """
        return cube.aggregated_by(["year"], iris.analysis.MEAN)

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        self.mean_file = self.declare_chunk(
            self.domain,
            self.variable,
            self.startdate,
            self.member,
            self.chunk,
            frequency=Frequencies.yearly,
            grid=self.grid,
        )