period.py

# -*- coding: utf-8 -*-

from collections import namedtuple

import iris
import iris.coord_categorisation


class Period:
    def __init__(self, output_coord, label):
        self.constraint = None
        self.input_coord = "time"
        self.output_coord = output_coord
        self.label = label


class Annual(Period):
    def __init__(self):
        super().__init__("year", "yr")

    def add_coord_categorisation(self, cube):
        iris.coord_categorisation.add_year(
            cube, self.input_coord, name=self.output_coord
        )
        return self.output_coord


class Monthly(Period):
    def __init__(self):
        super().__init__(["year", "month_number"], "mon")

    def add_coord_categorisation(self, cube):
        iris.coord_categorisation.add_year(cube, self.input_coord, name="year")
        iris.coord_categorisation.add_month_number(
            cube, self.input_coord, name="month_number"
        )
        return self.output_coord


class Season(Period):
    YEAR = "JFMAMJJASOND"
    MONTHS = [
        "January",
        "February",
        "March",
        "April",
        "May",
        "June",
        "July",
        "August",
        "September",
        "October",
        "November",
        "December",
    ]

    def __init__(self, specifier):
        super().__init__("season_year", "sem")
        self.specifier = specifier.upper()
        self.first_month_number = (self.YEAR * 2).find(self.specifier) + 1
        self.length = len(self.specifier)
        last_month_number = self.first_month_number + self.length - 1
        if last_month_number > 12:
            last_month_number %= 12

            def selector(cell):
                m = cell.point.month
                return (self.first_month_number <= m <= 12) | (
                    1 <= m <= last_month_number
                )

        else:

            def selector(cell):
                m = cell.point.month
                return self.first_month_number <= m <= last_month_number

        self.constraint = iris.Constraint(time=selector)
        self.last_month_number = last_month_number

    def long_label(self):
        first_month = self.MONTHS[self.first_month_number - 1]
        last_month = self.MONTHS[self.last_month_number - 1]
        long_label = f"{first_month}-{last_month}"
        return long_label

    def add_coord_categorisation(self, cube):
        iris.coord_categorisation.add_season_year(
            cube, self.input_coord, name=self.output_coord
        )
        return self.output_coord


PeriodSpecification = namedtuple("PeriodSpecification", "type parameters")

PERIODS = {
    "annual": Annual,
    "seasonal": Season,
    "monthly": Monthly,
}


def build_period(period_spec):
    try:
        Period = PERIODS[period_spec.type]
    except KeyError:
        raise ValueError("Unknown period specification <{period_spec}>")
    if period_spec.parameters is None:
        period = Period()
    else:
        period = Period(period_spec.parameters)
    return period