daysoverpercentile.py 4.83 KB
Newer Older
# coding=utf-8
from bscearth.utils.date import parse_date, add_months

from earthdiagnostics.diagnostic import *
from earthdiagnostics.frequency import Frequencies
import iris
import iris.coord_categorisation
from iris.time import PartialDateTime
import iris.analysis


class DaysOverPercentile(Diagnostic):
    """
    Calculates the montlhy percentiles

    :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
    """

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

    def __init__(self, data_manager, startdate, member, chunk, domain, variable, leadtime, percentile):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.variable = variable
        self.domain = domain
        self.leadtime = leadtime
        self.percentile = percentile

    def __eq__(self, other):
        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.leadtime == other.leadtime and \
            self.percentile == other.percentile

    def __str__(self):
        return 'Days over percentile Startdate: {0} Member: {1} Chunk: {2} ' \
               'Variable: {3}:{4} Leadtime: {5}'.format(self.startdate, self.member, self.chunk, self.domain,
                                                        self.variable, self.leadtime)

    @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: domain, variable, percentil number, maximum depth (level)
        :type options: list[str]
        :return:
        """
        options_available = (DiagnosticDomainOption(),
                             DiagnosticOption('variable'),
                             DiagnosticListIntOption('leadtimes'),
                             DiagnosticFloatOption('percentile'))
        options = cls.process_options(options, options_available)

        job_list = list()
        for startdate, member, chunk in diags.config.experiment.get_chunk_list():
            for leadtime in options['leadtimes']:
                job_list.append(DaysOverPercentile(diags.data_manager, startdate, member, chunk,
                                                   options['domain'], options['variable'], leadtime,
                                                   options['percentile']))
        return job_list

    def request_data(self):
        var_name = self.variable + 'prct'
        self.percentiles_file = self.request_chunk(self.domain, var_name, None, None, None,
                                                   frequency=Frequencies.climatology)

        self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, self.member, self.chunk)

    def declare_data_generated(self):
        var_name = self.variable + '_daysover'.format(self.leadtime)
        self.percentiles_file = self.declare_chunk(self.domain, var_name, None, None, None,
                                                   frequency=Frequencies.climatology, vartype=VariableType.STATISTIC)

    def compute(self):
        """
        Runs the diagnostic
        """
        percentiles = iris.load_cube(self.percentiles_file.local_file)

        var = iris.load_cube(self.variable_file.local_file)

        date = parse_date(self.startdate)
        leadtimes = {0: PartialDateTime(date.year, date.month)}

        def assign_leadtime(coord, x):
            try:
                leadtime_month = 0
                partial_date = leadtimes[leadtime_month]
                while coord.units.num2date(x) != partial_date:
                    leadtime_month += 1
                    try:
                        partial_date = leadtimes[leadtime_month]
                    except KeyError:
                        new_date = add_months(date, leadtime_month, self.experiment_config.calendar)
                        partial_date = PartialDateTime(new_date.year, new_date.month)
                        leadtimes[leadtime_month] = partial_date
                return leadtime_month
            except Exception as ex:
                pass

        for leadtime_slice in var.slices_over('leadtime'):
            percentiles_leadtime = percentiles.extract(iris.Constraint(leadtime=leadtime_slice.coord('leadtime').points[0]))
            for percentile_slice in percentiles_leadtime.slices_over('percentile'):
                over = leadtime_slice > percentile_slice
                print(over)