period.py 2.94 KB
Newer Older
1
2
3
4
5
# -*- coding: utf-8 -*-

from collections import namedtuple

import iris
6
import iris.coord_categorisation
7
8


9
class Period:
10
    def __init__(self, output_coord, label):
11
        self.constraint = None
12
        self.input_coord = "time"
13
        self.output_coord = output_coord
14
        self.label = label
15
16
17
18


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

    def add_coord_categorisation(self, cube):
22
23
24
        iris.coord_categorisation.add_year(
            cube, self.input_coord, name=self.output_coord
        )
25
26
27
28
29
        return self.output_coord


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

    def add_coord_categorisation(self, cube):
33
34
35
36
        iris.coord_categorisation.add_year(cube, self.input_coord, name="year")
        iris.coord_categorisation.add_month_number(
            cube, self.input_coord, name="month_number"
        )
37
38
39
        return self.output_coord


40
class Season(Period):
41
    YEAR = "JFMAMJJASOND"
42
    MONTHS = [
43
44
45
46
47
48
49
50
51
52
53
54
        "January",
        "February",
        "March",
        "April",
        "May",
        "June",
        "July",
        "August",
        "September",
        "October",
        "November",
        "December",
55
56
57
    ]

    def __init__(self, specifier):
58
        super().__init__("season_year", "sem")
59
        self.specifier = specifier.upper()
60
        self.first_month_number = (self.YEAR * 2).find(self.specifier) + 1
61
        self.length = len(self.specifier)
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
        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
80

81
    def long_label(self):
82
83
84
        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}"
85
86
87
        return long_label

    def add_coord_categorisation(self, cube):
88
        iris.coord_categorisation.add_season_year(
89
90
            cube, self.input_coord, name=self.output_coord
        )
91
92
        return self.output_coord

93

94
PeriodSpecification = namedtuple("PeriodSpecification", "type parameters")
95
96

PERIODS = {
97
98
99
    "annual": Annual,
    "seasonal": Season,
    "monthly": Monthly,
100
101
102
103
104
105
106
}


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