Source code for eradiate.plot
from __future__ import annotations
import typing as t
import warnings
import matplotlib.pyplot as _plt
import numpy as np
import xarray.plot
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from xarray.plot import FacetGrid
[docs]def set_style(rc=None):
"""
Apply Eradiate Matplotlib style (the Seaborn ``ticks`` style with the
viridis colormap).
Parameters
----------
rc : dict, optional
A Matplotlib rc parameter dictionary to be applied in addition to the
Eradiate style.
"""
try:
import seaborn
except ModuleNotFoundError:
warnings.warn(
"To use Eradiate's Matplotlib style, you must install Seaborn.\n"
"See instructions on https://seaborn.pydata.org/installing.html."
)
raise
if rc is None:
rc = {}
seaborn.set_theme(style="ticks", rc={"image.cmap": "viridis", **rc})
[docs]def detect_axes(from_=None):
"""
Try and extract a :class:`~matplotlib.axes.Axes` list from a data structure.
Parameters
----------
from_ : matplotlib figure or axes or list of axes or :class:`xarray.plot.FacetGrid`, optional
Data structure to get an :class:`~matplotlib.axes.Axes` list from.
If unset, :func:`matplotlib.pyplot.gca()` is used.
Returns
-------
list of matplotlib axes
Extracted list of :class:`~matplotlib.axes.Axes`.
Raises
------
TypeError
If ``from_`` is of unsupported type.
"""
if from_ is None:
from_ = _plt.gca()
if isinstance(from_, Figure):
return from_.axes
if isinstance(from_, Axes):
return [from_]
if isinstance(from_, FacetGrid):
return list(from_.axes.flatten())
if isinstance(from_, list):
if all([isinstance(x, Axes) for x in from_]):
return from_
raise TypeError("unsupported type")
[docs]def get_axes_from_facet_grid(
facet_grid: xarray.plot.FacetGrid, exclude: str = None
) -> list[Axes]:
"""
Extract a flat list of :class:`~matplotlib.axes.Axes` from a
:class:`~xarray.plot.FacetGrid`.
Parameters
----------
facet_grid : :class:`xarray.plot.FacetGrid`
Object to extract a list of :class:`~matplotlib.axes.Axes`.
exclude : str, optional
Exclude selected :class:`~matplotlib.axes.Axes` from
the returned list. Supported values:
* ``"lower_left"``: exclude lower-left corner.
Returns
-------
list of matplotlib axes
List of extracted :class:`~matplotlib.axes.Axes` objects.
"""
index = np.full((facet_grid.axes.size,), True, dtype=bool)
n_rows, n_cols = facet_grid.axes.shape
# This filtering system can be extended in the future
if exclude == "lower_left":
index[(n_rows - 1) * n_cols] = False
return list(facet_grid.axes.flatten()[index])
[docs]def remove_xylabels(from_=None) -> None:
"""
Remove x- and y-axis labels from ``from_``
(processed by :func:`detect_axes`).
Parameters
----------
from_ : matplotlib figure or axes or list of axes or :class:`xarray.plot.FacetGrid`, optional
Data structure to get an :class:`~matplotlib.axes.Axes` list from.
If unset, :func:`matplotlib.pyplot.gca()` is used.
See Also
--------
:func:`detect_axes`
"""
for ax in detect_axes(from_):
ax.set_xlabel("")
ax.set_ylabel("")
[docs]def remove_xyticks(from_=None) -> None:
"""
Remove x and y axis tick labels from ``from_``
(processed by :func:`detect_axes`).
Parameters
----------
from_ : matplotlib figure or axes or list of axes or :class:`xarray.plot.FacetGrid`, optional
Data structure to get an :class:`~matplotlib.axes.Axes` list from.
If unset, :func:`matplotlib.pyplot.gca()` is used.
See Also
--------
:func:`detect_axes`
"""
for ax in detect_axes(from_):
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
[docs]def make_ticks(num_ticks: int, limits: t.Sequence[float]):
"""
Generate ticks and their respective tickmarks.
Parameters
----------
num_ticks : int
Number of ticks to generate, including the limits
of the given range
limits : tuple[float, float]
List of two values, limiting the ticks inclusive
Returns
------
steps : list of float
Values for the ticks.
marks : list of str
Tick values converted to degrees as string tickmarks.
"""
step_width = float(limits[1] - limits[0]) / (num_ticks - 1)
steps = [limits[0] + step_width * i for i in range(num_ticks)]
marks = [f"{i / np.pi * 180}°" for i in steps]
return steps, marks
