You are viewing the latest unreleased documentation v3.1.dev0. You may prefer a stable version.

Source code for iris.coord_categorisation

# Copyright Iris contributors
#
# This file is part of Iris and is released under the LGPL license.
# See COPYING and COPYING.LESSER 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. Args: * cube (:class:`iris.cube.Cube`): the cube containing 'from_coord'. The new coord will be added into it. * name (string): name of the created coordinate * from_coord (:class:`iris.coords.Coord` or string): coordinate in 'cube', or the name of one * category_function (callable): function(coordinate, value), returning a category value for a coordinate point-value Kwargs: * units: 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 date of a time-coordinate point. Args: * coord (Coord): coordinate (must be Time-type) * time (float): value of a coordinate point Returns: datetime.date """ # NOTE: All of the currently defined categorisation functions are # calendar operations on Time coordinates. # - All these currently depend on Unit::num2date, which is deprecated (!!) # - We will want to do better, when we sort out our own Calendars. # - For now, just make sure these all call through this one function. return coord.units.num2date(time) # -------------------------------------------- # 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): """ Returns 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): """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. """ month_year_adjusts = [None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for season in seasons: months = _months_in_season(season) for month in months: if month > months[-1]: month_year_adjusts[month] = 1 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. 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". * seasons (:class:`list` of strings): 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, values 0..N-1 where N is the number of 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_number". * seasons (:class:`list` of strings): 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") ): """ 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): 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) # Define the adjustments to be made to the year. month_year_adjusts = _month_year_adjusts(seasons) # 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. 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. * season (string): Season defined by month abbreviations. Kwargs: * name (string): 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)