# 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)
-3.10.0.dev18-gold?style=flat){kind=icon}