period.py 4.2 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
12
13
        self.constraint = None
        self.input_coord = 'time'
        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
22
23
24
25
26
27
28
29

    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):
30
        super().__init__(['year', 'month_number'], 'mon')
31
32

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


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

59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    def __init__(self, seasons=("djf", "mam", "jja", "son")):
        super().__init__(("season_year", "season"), "sem")
        if isinstance(seasons, str):
            season = seasons.lower()
            self.first_month_number = self.YEAR.find(season) + 1
            self.length = len(season)
            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
            complement_season = Season.season_complement(season)
            self.seasons = (season, complement_season)
85
        else:
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
            self.seasons = seasons

    @staticmethod
    def season_complement(season):
        season = season.lower()
        length = len(season)
        index = Season.YEAR.find(season)
        if index < 0:
            # Can't match the season, raise an error.
            raise ValueError("unrecognised season: {!s}".format(season))
        complement_length = 12 - length
        complement_start = index + length
        complement_end = complement_start + complement_length
        complement_season = Season.YEAR[complement_start:complement_end]
        return complement_season
101

102
103
104
    def long_label(self):
        first_month = self.MONTHS[self.first_month_number-1]
        last_month = self.MONTHS[self.last_month_number-1]
105
        long_label = f'{first_month}-{last_month}'
106
107
108
        return long_label

    def add_coord_categorisation(self, cube):
109
110
111
112
113
114
115
116
117
118
119
120
        iris.coord_categorisation.add_season_year(
            cube,
            self.input_coord,
            name="season_year",
            seasons=self.seasons,
        )
        iris.coord_categorisation.add_season(
            cube,
            self.input_coord,
            name="season",
            seasons=self.seasons,
        )
121
122
        return self.output_coord

123
124
125
126

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

PERIODS = {
127
    'annual': Annual,
128
    'seasonal': Season,
129
    'monthly': Monthly,
130
131
132
133
134
135
136
}


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