Allow add_season_year() to optionally send spans backwards (#5573)

* New preceding year functionality for add_season_year().

* Modernise test_coord_categorisation.

* New test for backwards behaviour.

* What's New entry.

* Rename to use_year_at_season_start.

* Use np.zeros.

docs/src/whatsnew/latest.rst

@@ -38,6 +38,13 @@ This document explains the changes made to Iris for this release
    and :class:`~iris.coord_systems.RotatedMercator` coordinate systems,
    complete with NetCDF loading and saving. (:pull:`5548`)
 
+#. `@trexfeathers`_ added the ``use_year_at_season_start`` parameter to
+   :func:`iris.coord_categorisation.add_season_year`. When
+   ``use_year_at_season_start==True``: seasons spanning the year boundary (e.g.
+   Winter - December to February) will be assigned to the preceding year (e.g.
+   the year of December) instead of the following year (the default behaviour).
+   (:pull:`5573`)
+
 
 🐛 Bugs Fixed
 =============

lib/iris/coord_categorisation.py

@@ -271,20 +271,33 @@ def _validate_seasons(seasons):
     return
 
 
-def _month_year_adjusts(seasons):
-    """Compute the year adjustments required for each month.
+def _month_year_adjusts(seasons, use_year_at_season_start=False):
+    """
+    Compute the year adjustments required for each month.
 
-    These determine whether the month belongs to a season in the same
-    year or is in the start of a season that counts towards the next
-    year.
+    These adjustments ensure that no season spans two years by assigning months
+    to the **next** year (use_year_at_season_start is False) or the
+    **previous** year (use_year_at_season_start is True). E.g. Winter - djf:
+    either assign Dec to the next year, or Jan and Feb to the previous year.
 
     """
-    month_year_adjusts = [None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+    # 1 'slot' for each month, with an extra leading 'slot' because months
+    #  are 1-indexed - January is 1, therefore corresponding to the 2nd
+    #  array index.
+    month_year_adjusts = np.zeros(13, dtype=int)
+
     for season in seasons:
-        months = _months_in_season(season)
-        for month in months:
-            if month > months[-1]:
-                month_year_adjusts[month] = 1
+        months = np.array(_months_in_season(season))
+        if use_year_at_season_start:
+            months_to_shift = months < months[0]
+            year_shift = -1
+        else:
+            # Sending forwards.
+            months_to_shift = months > months[-1]
+            year_shift = 1
+        indices_to_shift = months[np.flatnonzero(months_to_shift)]
+        month_year_adjusts[indices_to_shift] = year_shift
+
     return month_year_adjusts
 
 
@@ -383,34 +396,40 @@ def _season_number(coord, value):
 
 
 def add_season_year(
-    cube, coord, name="season_year", seasons=("djf", "mam", "jja", "son")
+    cube,
+    coord,
+    name="season_year",
+    seasons=("djf", "mam", "jja", "son"),
+    use_year_at_season_start=False,
 ):
     """
-    Add a categorical year-of-season coordinate, with user specified
-    seasons.
-
-    Args:
-
-    * cube (:class:`iris.cube.Cube`):
-        The cube containing 'coord'. The new coord will be added into
-        it.
-    * coord (:class:`iris.coords.Coord` or string):
-        Coordinate in 'cube', or its name, representing time.
-
-    Kwargs:
-
-    * name (string):
-        Name of the created coordinate. Defaults to "season_year".
-    * seasons (:class:`list` of strings):
+    Add a categorical year-of-season coordinate, with user specified seasons.
+
+    Parameters
+    ----------
+    cube : :class:`iris.cube.Cube`
+        The cube containing `coord`. The new coord will be added into it.
+    coord : :class:`iris.coords.Coord` or str
+        Coordinate in `cube`, or its name, representing time.
+    name : str, default="season_year"
+        Name of the created coordinate.
+    seasons : tuple of str, default=("djf", "mam", "jja", "son")
         List of seasons defined by month abbreviations. Each month must
         appear once and only once. Defaults to standard meteorological
-        seasons ('djf', 'mam', 'jja', 'son').
+        seasons (``djf``, ``mam``, ``jja``, ``son``).
+    use_year_at_season_start: bool, default=False
+        Seasons spanning the year boundary (e.g. Winter ``djf``) will belong
+        fully to the following year by default (e.g. the year of Jan and Feb).
+        Set to ``True`` for spanning seasons to belong to the preceding
+        year (e.g. the year of Dec) instead.
 
     """
     # Check that the seasons are valid.
     _validate_seasons(seasons)
     # Define the adjustments to be made to the year.
-    month_year_adjusts = _month_year_adjusts(seasons)
+    month_year_adjusts = _month_year_adjusts(
+        seasons, use_year_at_season_start=use_year_at_season_start
+    )
 
     # Define a categorisation function.
     def _season_year(coord, value):