diff --git a/bluemath_tk/core/plotting/base_plotting.py b/bluemath_tk/core/plotting/base_plotting.py index 4dc786e..ceed5d6 100644 --- a/bluemath_tk/core/plotting/base_plotting.py +++ b/bluemath_tk/core/plotting/base_plotting.py @@ -1,15 +1,16 @@ from abc import ABC, abstractmethod -from typing import Tuple +from typing import Dict, List, Tuple, Union import cartopy.crs as ccrs import matplotlib.pyplot as plt +import numpy as np import plotly.graph_objects as go import xarray as xr from ...config.paths import PATHS -from .colors import hex_colors_land, hex_colors_water +from .colors import hex_colors_land, hex_colors_water, hex_colors_water_transition from .satellite import get_satellite_image -from .utils import join_colormaps +from .utils import format_ticks, join_colormaps, nice_ticks, single_colormap class BasePlotting(ABC): @@ -117,6 +118,11 @@ def plot_bathymetry( ax: plt.Axes, source: str, area: Tuple[float, float, float, float], + var_name: str = "elevation", + transition_depth: float = None, + isodepths: Union[float, List[float]] = None, + isodepths_labels: bool = True, + isodepths_kwargs: Dict = None, **kwargs, ) -> None: """ @@ -130,34 +136,164 @@ def plot_bathymetry( The source of the bathymetry data. Must be a key in the PATHS dictionary. area: Tuple[float, float, float, float] The area of the bathymetry data in the format (lon_min, lon_max, lat_min, lat_max). + transition_depth: float + Depth (negative, in data units) at which the ocean colormap turns from + blue to beige. Only used with the "albita_ocean" colormap. Default is + None, which spreads the ocean colours evenly over the depth range. + isodepths: Union[float, List[float]] + Depth / elevation values (in data units, so depths are negative) at which + to draw contour lines on top of the map. Default is None (no contours). + isodepths_labels: bool + Whether to annotate the contour lines with their value. Default is True. + isodepths_kwargs: Dict + Additional keyword arguments passed to the contour call (e.g. colors, + linewidths, linestyles), and, under the "labels_kwargs" key, a dictionary + of arguments passed to ax.clabel(). **kwargs Additional keyword arguments passed to the xr.Dataset.plot() function. """ - if source not in PATHS: - raise ValueError(f"Source {source} not found in PATHS") - else: - bathymetry_ds = ( - xr.open_dataset(PATHS[source]) - .sel(lon=slice(area[0], area[1]), lat=slice(area[2], area[3])) - .elevation - ) + if isinstance(source, str): + if source not in PATHS: + raise ValueError(f"Source '{source}' not found in PATHS.") + + else: + bathymetry_ds = ( + xr.open_dataset(PATHS[source]) + .sel(lon=slice(area[0], area[1]), lat=slice(area[2], area[3]))[var_name] + ) + + elif isinstance(source, xr.Dataset): + bathymetry_ds = source[var_name].sel(lon=slice(area[0], area[1]), lat=slice(area[2], area[3])) + + elif isinstance(source, xr.DataArray): + bathymetry_ds = source cmap = kwargs.pop("cmap", self.bathymetry_defaults.get("cmap")) if cmap == "albita_ocean": - cmap, norm = join_colormaps( - cmap1=hex_colors_water, - cmap2=hex_colors_land, - value_range1=(bathymetry_ds.min(), 0.0), - value_range2=(0.0, bathymetry_ds.max()), + vmin, vmax = float(bathymetry_ds.min()), float(bathymetry_ds.max()) + anchor = ( + None + if transition_depth is None + else (hex_colors_water_transition, transition_depth) ) + + # Only keep the halves of the colormap the data actually covers, so that + # e.g. an all-ocean map does not carry an unused land ramp in its colorbar + if vmax <= 0.0: + cmap, norm = single_colormap( + cmap=hex_colors_water, + value_range=(vmin, vmax), + anchor=anchor, + ) + elif vmin >= 0.0: + cmap, norm = single_colormap( + cmap=hex_colors_land, + value_range=(vmin, vmax), + ) + else: + cmap, norm = join_colormaps( + cmap1=hex_colors_water, + cmap2=hex_colors_land, + value_range1=(vmin, 0.0), + value_range2=(0.0, vmax), + anchor1=anchor, + ) + + # The norm is boundary-based, so ask for a colorbar that stays linear in + # data space and is labelled with round values instead of raw boundaries + ticks = None + if kwargs.get("add_colorbar", True): + cbar_kwargs = dict(kwargs.pop("cbar_kwargs", None) or {}) + cbar_kwargs.setdefault("spacing", "proportional") + if "ticks" not in cbar_kwargs: + ticks = nice_ticks((vmin, vmax), include=0.0) + cbar_kwargs["ticks"] = ticks + kwargs["cbar_kwargs"] = cbar_kwargs + p = bathymetry_ds.plot(ax=ax, cmap=cmap, norm=norm, **kwargs) - # Hide minor ticks on colorbar if hasattr(p, "colorbar") and p.colorbar is not None: + # Hide minor ticks on colorbar p.colorbar.minorticks_off() + if ticks is not None: + # End ticks sit at the exact data limits, so round their labels + p.colorbar.set_ticklabels(format_ticks(ticks)) else: bathymetry_ds.plot(ax=ax, cmap=cmap, **kwargs) + if isodepths is not None: + self.plot_isodepths( + ax=ax, + bathymetry=bathymetry_ds, + isodepths=isodepths, + labels=isodepths_labels, + transform=kwargs.get("transform"), + **(isodepths_kwargs or {}), + ) + + @staticmethod + def plot_isodepths( + ax: plt.Axes, + bathymetry: xr.DataArray, + isodepths: Union[float, List[float]], + labels: bool = True, + transform=None, + **kwargs, + ): + """ + Draw isodepth contour lines on top of a bathymetry map. + + Parameters + ---------- + ax: plt.Axes + The axes on which to plot the contours. + bathymetry: xr.DataArray + The bathymetry data to contour. + isodepths: Union[float, List[float]] + Depth / elevation values (in data units, so depths are negative). + labels: bool + Whether to annotate the contour lines with their value. Default is True. + transform + Cartopy transform of the data, if any. + **kwargs + Additional keyword arguments passed to the contour call. The + "labels_kwargs" key holds a dictionary of arguments passed to ax.clabel(). + + Returns + ------- + matplotlib.contour.QuadContourSet + The contour set, so that it can be further customized. + """ + + levels = sorted(np.atleast_1d(isodepths).astype(float)) + labels_kwargs = kwargs.pop("labels_kwargs", {}) + if transform is not None: + kwargs.setdefault("transform", transform) + + contours = bathymetry.plot.contour( + ax=ax, + levels=levels, + colors=kwargs.pop("colors", "black"), + linewidths=kwargs.pop("linewidths", 0.5), + # Solid by default, as matplotlib dashes negative levels (i.e. all depths) + linestyles=kwargs.pop("linestyles", "solid"), + add_colorbar=False, + add_labels=False, + **kwargs, + ) + if labels: + ax.clabel( + contours, + **{ + "fmt": "%g", + "inline": True, + "fontsize": 7, + **labels_kwargs, + }, + ) + + return contours + def plot_satellite( self, ax: plt.Axes, diff --git a/bluemath_tk/core/plotting/colors.py b/bluemath_tk/core/plotting/colors.py index b6c457a..d082584 100644 --- a/bluemath_tk/core/plotting/colors.py +++ b/bluemath_tk/core/plotting/colors.py @@ -35,6 +35,10 @@ "#f1e2c6", ] +# Fractional position within `hex_colors_water` where the deep blues give way to the +# shallow beige/green tones (index of "#c8ebd8" over the number of colour intervals). +hex_colors_water_transition = 7 / (len(hex_colors_water) - 1) + hex_colors_land = [ "#cfe2bd", "#aece91", diff --git a/bluemath_tk/core/plotting/utils.py b/bluemath_tk/core/plotting/utils.py index ff690b4..f1f65fc 100644 --- a/bluemath_tk/core/plotting/utils.py +++ b/bluemath_tk/core/plotting/utils.py @@ -4,6 +4,7 @@ import matplotlib.pyplot as plt import numpy as np from matplotlib.colors import BoundaryNorm, Colormap, ListedColormap +from matplotlib.ticker import MaxNLocator def get_list_of_colors_for_colormap( @@ -55,6 +56,192 @@ def create_cmap_from_colors( return colors.LinearSegmentedColormap.from_list(name, rgb_colors, N=256) +def _as_colormap(cmap: Union[str, List[str], Colormap], name: str = "cmap") -> Colormap: + """ + Convert a colormap name or a list of hex colors into a Colormap. + """ + + if isinstance(cmap, str): + return plt.get_cmap(cmap) + elif isinstance(cmap, list): + return colors.LinearSegmentedColormap.from_list(name, cmap) + + return cmap + + +def _bounds_with_anchor( + value_range: Tuple[float, float], + num: int, + anchor: Tuple[float, float] = None, +) -> np.ndarray: + """ + Build `num` boundaries spanning `value_range`, optionally pinning a colormap + fraction to a given data value. + + Parameters + ---------- + value_range : Tuple[float, float] + Data range covered by the colormap. + num : int + Number of boundaries to generate. + anchor : Tuple[float, float], optional + (colormap_fraction, data_value) pair. The colour sitting at + `colormap_fraction` of the colormap is placed at `data_value`, and each + side of it is stretched linearly. If None, boundaries are evenly spaced. + + Returns + ------- + np.ndarray + Monotonically increasing boundaries of length `num`. + """ + + vmin, vmax = float(value_range[0]), float(value_range[1]) + if anchor is None: + return np.linspace(vmin, vmax, num) + + fraction, value = float(anchor[0]), float(anchor[1]) + if not 0.0 < fraction < 1.0: + raise ValueError(f"Anchor fraction must be in (0, 1), got {fraction}.") + if not vmin < value < vmax: + raise ValueError( + f"Anchor value {value} must lie strictly within the data range ({vmin}, {vmax})." + ) + + num_low = int(round(fraction * (num - 1))) + 1 + + return np.concatenate( + [ + np.linspace(vmin, value, num_low), + np.linspace(value, vmax, num - num_low + 1)[1:], + ] + ) + + +def single_colormap( + cmap: Union[str, List[str], Colormap], + value_range: Tuple[float, float], + name: str = "cmap", + cmap_range: Tuple[float, float] = (0.0, 1.0), + anchor: Tuple[float, float] = None, + num_colors: int = 256, +) -> Tuple[ListedColormap, BoundaryNorm]: + """ + Build a single colormap and its matching norm over a data range. + + Same interface as `join_colormaps`, for the cases where only one of the two + colormaps is actually needed (e.g. bathymetry without any emerged land). + + Parameters + ---------- + cmap : Union[str, List[str], Colormap] + Input colormap (name, list of hex codes, or Colormap object). + value_range : Tuple[float, float] + Value range in the data domain covered by the colormap. + name : str + Name of the output colormap. + cmap_range : Tuple[float, float] + Portion of the colormap to use (from 0 to 1). + anchor : Tuple[float, float] + Optional (colormap_fraction, data_value) pair pinning a colour to a data + value, stretching both sides of it linearly. + num_colors : int + Number of discrete colors to sample. + + Returns + ------- + ListedColormap + Colormap object. + BoundaryNorm + Normalization for mapping data to colors. + """ + + cmap = _as_colormap(cmap, name=name) + newcolors = cmap(np.linspace(cmap_range[0], cmap_range[1], num_colors)) + bounds = _bounds_with_anchor(value_range, num_colors + 1, anchor) + norm = BoundaryNorm(boundaries=bounds, ncolors=num_colors) + + return colors.ListedColormap(newcolors, name=name), norm + + +def nice_ticks( + value_range: Tuple[float, float], + num_ticks: int = 7, + include: float = None, + ends: bool = True, +) -> np.ndarray: + """ + Compute round tick values covering a data range. + + Parameters + ---------- + value_range : Tuple[float, float] + Data range to cover. + num_ticks : int + Approximate number of ticks. Default is 7. + include : float + Optional value to force into the ticks (e.g. the shoreline at 0), if it + falls inside the range. + ends : bool + Whether to always tick both ends of the range. Round ticks landing too + close to an end are dropped in their favour. Default is True. + + Returns + ------- + np.ndarray + Sorted tick values, all within the data range. + """ + + vmin, vmax = float(value_range[0]), float(value_range[1]) + ticks = MaxNLocator(nbins=num_ticks, steps=[1, 2, 2.5, 5, 10]).tick_values( + vmin, vmax + ) + step = ticks[1] - ticks[0] if len(ticks) > 1 else (vmax - vmin) + ticks = ticks[(ticks >= vmin) & (ticks <= vmax)] + + if ends: + # Leave room for the end labels, then place them at the exact data limits + ticks = ticks[(ticks > vmin + 0.3 * step) & (ticks < vmax - 0.3 * step)] + if include is not None and vmin <= include <= vmax: + ticks = np.union1d(ticks, [float(include)]) + if ends: + ticks = np.union1d(ticks, [vmin, vmax]) + + return ticks + + +def format_ticks(ticks: Union[List[float], np.ndarray], max_decimals: int = 3) -> List[str]: + """ + Format tick values as short labels, using the fewest decimals that still + keeps every label distinct. + + Useful for ticks placed at exact data limits, so that e.g. a shallowest + value of -0.3 m reads as "0" rather than "-0.3". + + Parameters + ---------- + ticks : Union[List[float], np.ndarray] + Tick values to label. + max_decimals : int + Maximum number of decimals to fall back on. Default is 3. + + Returns + ------- + List[str] + One label per tick. + """ + + values = np.asarray(ticks, dtype=float) + + for decimals in range(max_decimals + 1): + rounded = np.round(values, decimals) + rounded[rounded == 0.0] = 0.0 # Avoid "-0" labels + labels = [f"{value:.{decimals}f}" for value in rounded] + if len(set(labels)) == len(labels): + break + + return labels + + def join_colormaps( cmap1: Union[str, List[str], Colormap], cmap2: Union[str, List[str], Colormap], @@ -63,6 +250,8 @@ def join_colormaps( range2: Tuple[float, float] = (0.0, 1.0), value_range1: Tuple[float, float] = None, value_range2: Tuple[float, float] = None, + anchor1: Tuple[float, float] = None, + anchor2: Tuple[float, float] = None, ) -> Tuple[ListedColormap, BoundaryNorm]: """ Join two colormaps into one, with value ranges specified for each. @@ -77,6 +266,9 @@ def join_colormaps( Portion of each colormap to use (from 0 to 1). value_range1, value_range2 : Tuple[float, float] Value ranges in the data domain corresponding to each colormap. + anchor1, anchor2 : Tuple[float, float] + Optional (colormap_fraction, data_value) pairs pinning a colour of each + colormap to a data value, stretching both sides of it linearly. Returns ------- @@ -86,16 +278,9 @@ def join_colormaps( Normalization for mapping data to colors. """ - # Convert cmap1 to a Colormap if needed - if isinstance(cmap1, str): - cmap1 = plt.get_cmap(cmap1) - elif isinstance(cmap1, list): - cmap1 = colors.LinearSegmentedColormap.from_list("cmap1", cmap1) - - if isinstance(cmap2, str): - cmap2 = plt.get_cmap(cmap2) - elif isinstance(cmap2, list): - cmap2 = colors.LinearSegmentedColormap.from_list("cmap2", cmap2) + # Convert each input to a Colormap if needed + cmap1 = _as_colormap(cmap1, name="cmap1") + cmap2 = _as_colormap(cmap2, name="cmap2") # Get colors from each colormap colors1 = cmap1(np.linspace(range1[0], range1[1], 128)) @@ -104,8 +289,9 @@ def join_colormaps( # Create corresponding boundaries in data space if value_range1 is not None and value_range2 is not None: - bounds1 = np.linspace(value_range1[0], value_range1[1], 129) - bounds2 = np.linspace(value_range2[0], value_range2[1], 129) + # Values are cast to float, as scalar xarray objects break np.linspace wrapping + bounds1 = _bounds_with_anchor(value_range1, 129, anchor1) + bounds2 = _bounds_with_anchor(value_range2, 129, anchor2) all_bounds = np.sort(np.concatenate([bounds1[:-1], bounds2])) norm = BoundaryNorm(boundaries=all_bounds, ncolors=len(newcolors))