Fix seasonal period (backport of !168)

climix/period.py

@@ -40,7 +40,7 @@ class Monthly(Period):
 
 
 class Season(Period):
-    YEAR = 'JFMAMJJASOND'
+    YEAR = "jfmamjjasond" * 3
     MONTHS = [
         'January',
         'February',
@@ -56,29 +56,48 @@ class Season(Period):
         '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
-                )
-
+    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)
         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
+            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
 
     def long_label(self):
         first_month = self.MONTHS[self.first_month_number-1]
@@ -87,9 +106,18 @@ class Season(Period):
         return long_label
 
     def add_coord_categorisation(self, cube):
-        iris.coord_categorisation.add_season_year(cube,
-                                                  self.input_coord,
-                                                  name=self.output_coord)
+        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,
+        )
         return self.output_coord