Source code for utopya.eval.plots._utils

"""Utility functions that assist plot functions"""

import itertools
import logging
from typing import Dict, Tuple

import numpy as np
import xarray as xr

log = logging.getLogger(__name__)

# -----------------------------------------------------------------------------

[docs] def calc_pxmap_rectangles( *, x_coords: np.ndarray, y_coords: np.ndarray, x_scale: str = "lin", y_scale: str = "lin", default_pos: Tuple[float, float] = (0.0, 0.0), default_distance: float = 1.0, size_factor: float = 1.0, extend: bool = False, ) -> Tuple[xr.Dataset, Dict[str, tuple]]: """Calculates the positions and sizes of rectangles centered at the given coordinates in such a way that they fully cover the 2D space spanned by the coordinates. The exact values of the coordinates are arbitrary, but they should be ordered. The resulting rectangles will not necessarily be centered around the coordinate, but the distance between rectangle edges is set such that it lies halfway to the next coordinate in that dimension; the "halfway" is evaluated according to a certain scale. Args: x_coords (numpy.ndarray): The x coordinates y_coords (numpy.ndarray): The y coordinates x_scale (str, optional): The x-axis scale, used to determine rectangle sizes. y_scale (str, optional): The y-axis scale, used to determine rectangle sizes. default_pos (Tuple[float, float], optional): If any of the coordinates is not given, this information will be used for creating the rectangle specification. default_distance (float, optional): If any of the coordinates is not given or of size 1, this distance to the next data point is assumed in that dimension. In such a case, the scale does not have an effect on the resulting rectangle. size_factor (float, optional): Scaling factor for the rectangle sizes extend (bool, optional): Whether to extend the rectangles of the points at the border of the domain such that the coordinate is _not_ at the edge of the rectangle but in the bulk. Returns: Tuple[xarray.Dataset, Dict[str, tuple]]: The first tuple element is the dataset of rectangle specifications, each available as a data variable: - ``pos_x``, ``pos_y``: position of the lower-value coordinate of the rectangle. Together, this specifies the bottom left- hand corner of the rectangle (in a right-hand coordinate system) - ``len_x``, ``len_y``: lengths of the rectangle sides in the specified dimensions. Adding both these to the position leads to the top-right hand corner of the rectangle - ``rect_spec``: A matplotlib-compatible rectangle specification, i.e. (position, width, height) The second tuple element are the limits in x and y direction, given as dict with keys x and y. """ # Define some allowed linear and logarithmic scales LIN_SCALES = ["lin", "linear"] LOG_SCALES = ["log", "symlog"] CATEGORIAL_SCALES = ["cat", "categorical", "cat_at_value"] # Helper functions ........................................................ def determine_diffs_and_scale( coords: np.ndarray, *, scale: str, axis_name: str ) -> Tuple[np.ndarray, str]: """Given the coordinates and the scale, determine the distances between coordinates in one dimension. This is also used to check the scale. Args: coords (array-like): Sequence of coordinates scale (str): The scale to use axis_name (str): The name of the axis; used for error message only Returns: Tuple[numpy.ndarray, str]: Coordinate distances and the new scale Raises: ValueError: On invalid scale argument """ if scale in LIN_SCALES: diffs = np.diff(coords) elif scale in LOG_SCALES: diffs = np.diff(np.log(coords)) elif scale in CATEGORIAL_SCALES: # For categorial scale at value, use just some arbitrary integeres # for the coordinates; the scale can then be assumed linear. if scale != "cat_at_value": coords = list(range(len(coords))) diffs = np.diff(coords) scale = "lin" else: raise ValueError( "Invalid scale argument for {} axis: '{}'! " "Expected one of: {}" "".format( scale, axis_name, ", ".join(LIN_SCALES + LOG_SCALES + CATEGORIAL_SCALES), ) ) return diffs, scale def expand_diffs(diffs, *, extend: bool, default: float) -> np.ndarray: """Adds the boundary values to the given diffs""" if diffs.size < 1: # Need to use the default value diffs = np.array([default, default]) elif extend: # Add the first item to the beginning and the last to the end of # the diffs sequences to account for the borders diffs = np.insert(diffs, 0, diffs[0]) diffs = np.append(diffs, diffs[-1]) else: # No extension of borders desired --> add zeros to diffs diffs = np.insert(diffs, 0, 0.0) diffs = np.append(diffs, 0.0) return diffs def determine_limits(*, coords, sizes, scale: str) -> Tuple[float, float]: """Determine the limits of the space that is covered by the rectangles at the given coordinates and sizes. """ if scale in LIN_SCALES: return (float(coords[0]) - sizes[0], float(coords[-1]) + sizes[-1]) elif scale in LOG_SCALES: return ( np.exp(np.log(float(coords[0])) - sizes[0]), np.exp(np.log(float(coords[-1])) + sizes[-1]), ) # NOTE This point is not reached; scale is checked elsewhere # The functions below create a rectangle specifier given the x and y # positions and the desired sizes. # Have 4 explicit functions (one for each scale specification) because it's # very tedious to create the rectangle specifier in a general way ... Grml. def calc_linlin_rect( x: float, y: float, x_sizes: tuple, y_sizes: tuple ) -> tuple: """ Create the rectangle specification for a single rectangle embedded in a lin-lin plot. Args: x (float): The x-coordinate y (float): The y-coordinate x_sizes (tuple): The distance to the rectangle edges in x direction y_sizes (tuple): The distance to the rectangle edges in y direction Returns: tuple: (x pos. bottom left-hand corner, y pos. bottom left-hand corner, total x length, total y length) """ return ( x - x_sizes[0], y - y_sizes[0], x_sizes[0] + x_sizes[1], y_sizes[0] + y_sizes[1], ) def calc_linlog_rect(x, y, x_sizes, y_sizes) -> tuple: """See calc_linlin_rect""" return ( x - x_sizes[0], np.exp(np.log(y) - y_sizes[0]), x_sizes[0] + x_sizes[1], np.exp(np.log(y) + y_sizes[1]) - np.exp(np.log(y) - y_sizes[0]), ) def calc_loglin_rect(x, y, x_sizes, y_sizes) -> tuple: """See calc_linlin_rect""" return ( np.exp(np.log(x) - x_sizes[0]), y - y_sizes[0], np.exp(np.log(x) + x_sizes[1]) - np.exp(np.log(x) - x_sizes[0]), y_sizes[0] + y_sizes[1], ) def calc_loglog_rect(x, y, x_sizes, y_sizes) -> tuple: """See calc_linlin_rect""" return ( np.exp(np.log(x) - x_sizes[0]), np.exp(np.log(y) - y_sizes[0]), np.exp(np.log(x) + x_sizes[1]) - np.exp(np.log(x) - x_sizes[0]), np.exp(np.log(y) + y_sizes[1]) - np.exp(np.log(y) - y_sizes[0]), ) # Create a map to select the callable via `in LIN_SCALES` boolean key pair FUNC_MAP = { # x lin? y lin? (True, True): calc_linlin_rect, (True, False): calc_linlog_rect, (False, True): calc_loglin_rect, (False, False): calc_loglog_rect, } # ......................................................................... log.debug("Calculating rectangle positions and sizes ...") # Allow None, in which case the rectangle is set to a default position on # that axis x_coords = x_coords if x_coords is not None else [default_pos[0]] y_coords = y_coords if y_coords is not None else [default_pos[1]] # Check the lengths if len(x_coords) < 1 or len(y_coords) < 1: raise ValueError( "Coordinate sequences need to be at least of length " "one or None (to use defaults), but were: {} and {}" "".format(x_coords, y_coords) ) # Calculate distances, distinguishing between linear and logarithmic scale x_diffs, x_scale = determine_diffs_and_scale( x_coords, scale=x_scale, axis_name="x" ) y_diffs, y_scale = determine_diffs_and_scale( y_coords, scale=y_scale, axis_name="y" ) # Determine rectangle calculation method calc_rect = FUNC_MAP[(x_scale in LIN_SCALES, y_scale in LIN_SCALES)] # Need two additional borders: the lower border of the first rectangle and # the upper border of the last rectangle. x_diffs = expand_diffs(x_diffs, extend=extend, default=default_distance) y_diffs = expand_diffs(y_diffs, extend=extend, default=default_distance) # Divide by two to have the _sizes_ of the rectangle, not the distance # between the points, and then apply the size factor. x_sizes = x_diffs / 2.0 * size_factor y_sizes = y_diffs / 2.0 * size_factor # Create rects Dataset, which is populated below. It contains the positions # and sizes as separate data variables. shape = (len(x_coords), len(y_coords)) rect_spec = np.dtype( [("xy", (float, (2,))), ("width", float), ("height", float)] ) rects = xr.Dataset( data_vars=dict( pos_x=(("x", "y"), np.full(shape, np.nan)), pos_y=(("x", "y"), np.full(shape, np.nan)), len_x=(("x", "y"), np.full(shape, np.nan)), len_y=(("x", "y"), np.full(shape, np.nan)), rect_spec=(("x", "y"), np.full(shape, np.nan, dtype=rect_spec)), ), coords=dict(x=(("x",), x_coords), y=(("y",), y_coords)), ) # Now, iterate over the coordinate combinations, calculate the rectangle # specification, and store it in the data array. it = itertools.product( enumerate(rects.coords["x"]), enumerate(rects.coords["y"]) ) for (x_idx, x_coord), (y_idx, y_coord) in it: # Get the relevant sizes calculated from the diffs and calculate the # rectangle specification from that pos_x, pos_y, len_x, len_y = calc_rect( x_coord, y_coord, (x_sizes[x_idx], x_sizes[x_idx + 1]), # left # right (y_sizes[y_idx], y_sizes[y_idx + 1]), # bottom ) # top # Store in the corresponding data variable in the xr.Dataset rects["pos_x"][x_idx, y_idx] = pos_x rects["pos_y"][x_idx, y_idx] = pos_y rects["len_x"][x_idx, y_idx] = len_x rects["len_y"][x_idx, y_idx] = len_y rects["rect_spec"][x_idx, y_idx] = (pos_x, pos_y), len_x, len_y # Calculate limits of the space that is covered by the rectangles x_lims = determine_limits(coords=x_coords, sizes=x_sizes, scale=x_scale) y_lims = determine_limits(coords=y_coords, sizes=y_sizes, scale=y_scale) # Done. return rects, dict(x=x_lims, y=y_lims)