Source code for iris.coord_categorisation

# Copyright Iris contributors
#
# This file is part of Iris and is released under the BSD license.
# See LICENSE in the root of the repository for full licensing details.
"""Cube functions for coordinate categorisation.

All the functions provided here add a new coordinate to a cube.

* The function :func:`add_categorised_coord` performs a generic
  coordinate categorisation.
* The other functions all implement specific common cases
  (e.g. :func:`add_day_of_month`).
  Currently, these are all calendar functions, so they only apply to
  "Time coordinates".

"""

import calendar
import collections

import numpy as np

import iris.coords


[docs] def add_categorised_coord(cube, name, from_coord, category_function, units="1"): """Add a new coordinate to a cube, by categorising an existing one. Make a new :class:`iris.coords.AuxCoord` from mapped values, and add it to the cube. Parameters ---------- cube : :class:`iris.cube.Cube` The cube containing 'from_coord'. The new coord will be added into it. name : str Name of the created coordinate. from_coord : :class:`iris.coords.Coord` or str Coordinate in 'cube', or the name of one. category_function : callable Function(coordinate, value), returning a category value for a coordinate point-value. units : str, default="1" Units of the category value, typically 'no_unit' or '1'. """ # Interpret coord, if given as a name if isinstance(from_coord, str): from_coord = cube.coord(from_coord) if len(cube.coords(name)) > 0: msg = 'A coordinate "%s" already exists in the cube.' % name raise ValueError(msg) # Construct new coordinate by mapping values, using numpy.vectorize to # support multi-dimensional coords. # Test whether the result contains strings. If it does we must manually # force the dtype because of a numpy bug (see numpy #3270 on GitHub). result = category_function(from_coord, from_coord.points.ravel()[0]) if isinstance(result, str): str_vectorised_fn = np.vectorize(category_function, otypes=[object]) def vectorised_fn(*args): # Use a common type for string arrays (N.B. limited to 64 chars). return str_vectorised_fn(*args).astype("|U64") else: vectorised_fn = np.vectorize(category_function) new_coord = iris.coords.AuxCoord( vectorised_fn(from_coord, from_coord.points), units=units, attributes=from_coord.attributes.copy(), ) new_coord.rename(name) # Add into the cube cube.add_aux_coord(new_coord, cube.coord_dims(from_coord))
# ====================================== # Specific functions for particular purposes # # NOTE: all the existing ones are calendar operations, so are for 'Time' # coordinates only # # Private "helper" function def _pt_date(coord, time): """Return the datetime of a time-coordinate point. Parameters ---------- coord : Coord Coordinate (must be Time-type). time : float Value of a coordinate point. Returns ------- cftime.datetime """ # NOTE: All of the currently defined categorisation functions are # calendar operations on Time coordinates. return coord.units.num2date(time, only_use_cftime_datetimes=True) # -------------------------------------------- # Time categorisations : calendar date components
[docs] def add_year(cube, coord, name="year"): """Add a categorical calendar-year coordinate.""" add_categorised_coord(cube, name, coord, lambda coord, x: _pt_date(coord, x).year)
[docs] def add_month_number(cube, coord, name="month_number"): """Add a categorical month coordinate, values 1..12.""" add_categorised_coord(cube, name, coord, lambda coord, x: _pt_date(coord, x).month)
[docs] def add_month_fullname(cube, coord, name="month_fullname"): """Add a categorical month coordinate, values 'January'..'December'.""" add_categorised_coord( cube, name, coord, lambda coord, x: calendar.month_name[_pt_date(coord, x).month], units="no_unit", )
[docs] def add_month(cube, coord, name="month"): """Add a categorical month coordinate, values 'Jan'..'Dec'.""" add_categorised_coord( cube, name, coord, lambda coord, x: calendar.month_abbr[_pt_date(coord, x).month], units="no_unit", )
[docs] def add_day_of_month(cube, coord, name="day_of_month"): """Add a categorical day-of-month coordinate, values 1..31.""" add_categorised_coord(cube, name, coord, lambda coord, x: _pt_date(coord, x).day)
[docs] def add_day_of_year(cube, coord, name="day_of_year"): """Add a categorical day-of-year coordinate, values 1..365 (1..366 in leap years).""" # Note: cftime.datetime objects return a normal tuple from timetuple(), # unlike datetime.datetime objects that return a namedtuple. # Index the time tuple (element 7 is day of year) instead of using named # element tm_yday. add_categorised_coord( cube, name, coord, lambda coord, x: _pt_date(coord, x).timetuple()[7] )
# -------------------------------------------- # Time categorisations : days of the week
[docs] def add_weekday_number(cube, coord, name="weekday_number"): """Add a categorical weekday coordinate, values 0..6 [0=Monday].""" add_categorised_coord( cube, name, coord, lambda coord, x: _pt_date(coord, x).dayofwk )
[docs] def add_weekday_fullname(cube, coord, name="weekday_fullname"): """Add a categorical weekday coordinate, values 'Monday'..'Sunday'.""" add_categorised_coord( cube, name, coord, lambda coord, x: calendar.day_name[_pt_date(coord, x).dayofwk], units="no_unit", )
[docs] def add_weekday(cube, coord, name="weekday"): """Add a categorical weekday coordinate, values 'Mon'..'Sun'.""" add_categorised_coord( cube, name, coord, lambda coord, x: calendar.day_abbr[_pt_date(coord, x).dayofwk], units="no_unit", )
# -------------------------------------------- # Time categorisations : hour of the day
[docs] def add_hour(cube, coord, name="hour"): """Add a categorical hour coordinate, values 0..23.""" add_categorised_coord(cube, name, coord, lambda coord, x: _pt_date(coord, x).hour)
# ---------------------------------------------- # Time categorisations : meteorological seasons def _months_in_season(season): """Return a list of month numbers corresponding to each month in the given season.""" cyclic_months = "jfmamjjasondjfmamjjasond" m0 = cyclic_months.find(season.lower()) if m0 < 0: # Can't match the season, raise an error. raise ValueError("unrecognised season: {!s}".format(season)) m1 = m0 + len(season) return [(month % 12) + 1 for month in range(m0, m1)] def _validate_seasons(seasons): """Check that a set of seasons is valid. Validity means that all months are included in a season, and no month is assigned to more than one season. Raises ValueError if either of the conditions is not met, returns None otherwise. """ c = collections.Counter() for season in seasons: c.update(_months_in_season(season)) # Make a list of months that are not present... not_present = [ calendar.month_abbr[month] for month in range(1, 13) if month not in c ] if not_present: raise ValueError( "some months do not appear in any season: {!s}".format( ", ".join(not_present) ) ) # Make a list of months that appear multiple times... multi_present = [ calendar.month_abbr[month] for month in range(1, 13) if c[month] > 1 ] if multi_present: raise ValueError( "some months appear in more than one season: {!s}".format( ", ".join(multi_present) ) ) return def _month_year_adjusts(seasons, use_year_at_season_start=False): """Compute the year adjustments required for each month. 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. """ # 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 = 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 def _month_season_numbers(seasons): """Compute a mapping between months and season number. Returns a list to be indexed by month number, where the value at each index is the number of the season that month belongs to. """ month_season_numbers = [None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for season_number, season in enumerate(seasons): for month in _months_in_season(season): month_season_numbers[month] = season_number return month_season_numbers
[docs] def add_season(cube, coord, name="season", seasons=("djf", "mam", "jja", "son")): """Add a categorical season-of-year 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" Name of the created coordinate. Defaults to "season". seasons : :class:`list` of str, optional List of seasons defined by month abbreviations. Each month must appear once and only once. Defaults to standard meteorological seasons ('djf', 'mam', 'jja', 'son'). """ # Check that the seasons are valid. _validate_seasons(seasons) # Get a list of the season number each month is is, using month numbers # as the indices. month_season_numbers = _month_season_numbers(seasons) # Define a categorisation function. def _season(coord, value): dt = _pt_date(coord, value) return seasons[month_season_numbers[dt.month]] # Apply the categorisation. add_categorised_coord(cube, name, coord, _season, units="no_unit")
[docs] def add_season_number( cube, coord, name="season_number", seasons=("djf", "mam", "jja", "son") ): """Add a categorical season-of-year coordinate. Add a categorical season-of-year coordinate, values 0..N-1 where N is the number of 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" Name of the created coordinate. Defaults to "season_number". seasons : :class:`list` of str, optional List of seasons defined by month abbreviations. Each month must appear once and only once. Defaults to standard meteorological seasons ('djf', 'mam', 'jja', 'son'). """ # Check that the seasons are valid. _validate_seasons(seasons) # Get a list of the season number each month is is, using month numbers # as the indices. month_season_numbers = _month_season_numbers(seasons) # Define a categorisation function. def _season_number(coord, value): dt = _pt_date(coord, value) return month_season_numbers[dt.month] # Apply the categorisation. add_categorised_coord(cube, name, coord, _season_number)
[docs] def add_season_year( 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. 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``). 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, use_year_at_season_start=use_year_at_season_start ) # Define a categorisation function. def _season_year(coord, value): dt = _pt_date(coord, value) year = dt.year year += month_year_adjusts[dt.month] return year # Apply the categorisation. add_categorised_coord(cube, name, coord, _season_year)
[docs] def add_season_membership(cube, coord, season, name="season_membership"): """Add a categorical season membership coordinate for a user specified season. The coordinate has the value True for every time that is within the given season, and the value False otherwise. 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. season : str Season defined by month abbreviations. name : str, default="season_membership" Name of the created coordinate. Defaults to "season_membership". """ months = _months_in_season(season) def _season_membership(coord, value): dt = _pt_date(coord, value) if dt.month in months: return True return False add_categorised_coord(cube, name, coord, _season_membership)