Source code for qcodes.plotting.matplotlib_helpers

from __future__ import annotations

import copy
import logging
from typing import TYPE_CHECKING, Any, Literal, cast

if TYPE_CHECKING:
    import matplotlib
    import matplotlib.colorbar

import numpy as np

from .auto_range import DEFAULT_PERCENTILE, auto_range_iqr

DEFAULT_COLOR_OVER = "Magenta"
DEFAULT_COLOR_UNDER = "Cyan"

_LOG = logging.getLogger(__name__)

_EXTEND_TYPE = Literal["neither", "both", "min", "max"]


def _set_colorbar_extend(
    colorbar: matplotlib.colorbar.Colorbar,
    extend: _EXTEND_TYPE,
) -> None:
    """
    Workaround for a missing setter for the extend property of a matplotlib
    colorbar.

    The colorbar object in matplotlib has no setter method and setting the
    colorbar extend does not take any effect.
    Calling a subsequent update will cause a runtime
    error because of the internal implementation of the rendering of the
    colorbar. To circumvent this we need to manually specify the property
    `_inside`, which is a slice that describes which of the colorbar levels
    lie inside of the box and it is thereby dependent on the extend.

    Args:
        colorbar: the colorbar for which to set the extend
        extend: the desired extend ('neither', 'both', 'min' or 'max')

    """
    colorbar.extend = extend
    _slice_dict = {
        "neither": slice(0, None),
        "both": slice(1, -1),
        "min": slice(1, None),
        "max": slice(0, -1),
    }
    colorbar._inside = _slice_dict[extend]  # type: ignore[attr-defined]


[docs] def apply_color_scale_limits( colorbar: matplotlib.colorbar.Colorbar, new_lim: tuple[float | None, float | None], data_lim: tuple[float, float] | None = None, data_array: np.ndarray | None = None, color_over: Any = DEFAULT_COLOR_OVER, color_under: Any = DEFAULT_COLOR_UNDER, ) -> None: """ Applies limits to colorscale and updates extend. This function applies the limits `new_lim` to the heatmap plot associated with the provided `colorbar`, updates the colorbar limits, and also adds the colorbar clipping indicators in form of small triangles on the top and bottom of the colorbar, according to where the limits are exceeded. Args: colorbar: The actual colorbar to be updated. new_lim: 2-tuple of the desired minimum and maximum value of the color scale. If any is `None` it will be left unchanged. data_lim: 2-tuple of the actual minimum and maximum value of the data. If left out the minimum and maximum are deduced from the provided data, or the data associated with the colorbar. data_array: Numpy array containing the data to be considered for scaling. Must be left out if `data_lim` is provided. If neither is provided the data associated with the colorbar is used. color_over: Matplotlib color representing the datapoints clipped by the upper limit. color_under: Matplotlib color representing the datapoints clipped by the lower limit. Raise: RuntimeError: If not received mesh data. Or if you specified both `data_lim` and `data_array`. """ import matplotlib.collections # browse the input data and make sure that `data_lim` and `new_lim` are # available if not isinstance(colorbar.mappable, matplotlib.collections.QuadMesh): raise RuntimeError( "Can only scale mesh data, but received " f'"{type(colorbar.mappable)}" instead' ) if data_lim is None: if data_array is None: data_array = cast(np.ndarray, colorbar.mappable.get_array()) data_lim = np.nanmin(data_array), np.nanmax(data_array) elif data_array is not None: raise RuntimeError( "You may not specify `data_lim` and `data_array` " "at the same time. Please refer to the docstring of " "`apply_color_scale_limits for details:\n\n`" f"{apply_color_scale_limits.__doc__!s}" ) else: data_lim = cast(tuple[float, float], tuple(sorted(data_lim))) # if `None` is provided in the new limits don't change this limit vlim = [new or old for new, old in zip(new_lim, colorbar.mappable.get_clim())] # sort limits in case they were given in a wrong order vlim = sorted(vlim) # detect exceeding colorscale and apply new limits exceeds_min, exceeds_max = (data_lim[0] < vlim[0], data_lim[1] > vlim[1]) if exceeds_min and exceeds_max: extend: _EXTEND_TYPE = "both" elif exceeds_min: extend = "min" elif exceeds_max: extend = "max" else: extend = "neither" _set_colorbar_extend(colorbar, extend) cmap = copy.copy(colorbar.mappable.get_cmap()) cmap.set_over(color_over) cmap.set_under(color_under) colorbar.mappable.set_cmap(cmap) colorbar.mappable.set_clim(*vlim)
[docs] def apply_auto_color_scale( colorbar: matplotlib.colorbar.Colorbar, data_array: np.ndarray | None = None, cutoff_percentile: tuple[float, float] | float = DEFAULT_PERCENTILE, color_over: Any | None = DEFAULT_COLOR_OVER, color_under: Any | None = DEFAULT_COLOR_UNDER, ) -> None: """ Sets the color limits such that outliers are disregarded. This method combines the automatic color scaling from :meth:`auto_range_iqr` with the color bar setting from :meth:`apply_color_scale_limits`. If you want to adjust the color scale based on the configuration file `qcodesrc.json`, use :meth:`auto_color_scale_from_config`, which is used In :func:`qcodes.dataset.plotting.plot_by_id`. Args: colorbar: The matplotlib colorbar to which to apply. data_array: The data on which the statistical analysis is based. If left out, the data associated with the `colorbar` is used cutoff_percentile: Percentile of data that may maximally be clipped on both sides of the distribution. If given a tuple (a,b) the percentile limits will be a and 100-b. color_over: Matplotlib color representing the datapoints clipped by the upper limit. color_under: Matplotlib color representing the datapoints clipped by the lower limit. Raises: RuntimeError: If not mesh data. """ import matplotlib.collections if data_array is None: if not isinstance(colorbar.mappable, matplotlib.collections.QuadMesh): raise RuntimeError("Can only scale mesh data.") data_array = cast(np.ndarray, colorbar.mappable.get_array()) assert data_array is not None new_lim = auto_range_iqr(data_array, cutoff_percentile) apply_color_scale_limits( colorbar, new_lim=new_lim, data_array=data_array, color_over=color_over, color_under=color_under, )
[docs] def auto_color_scale_from_config( colorbar: matplotlib.colorbar.Colorbar, auto_color_scale: bool | None = None, data_array: np.ndarray | None = None, cutoff_percentile: tuple[float, float] | float | None = DEFAULT_PERCENTILE, color_over: Any | None = None, color_under: Any | None = None, ) -> None: """ Sets the color limits such that outliers are disregarded, depending on the configuration file `qcodesrc.json`. If optional arguments are passed the config values are overridden. Args: colorbar: The colorbar to scale. auto_color_scale: Enable smart colorscale. If `False` nothing happens. Default value is read from ``config.plotting.auto_color_scale.enabled``. data_array: Numpy array containing the data to be considered for scaling. cutoff_percentile: The maxiumum percentile that is cut from the data. Default value is read from ``config.plotting.auto_color_scale.cutoff_percentile``. color_over: Matplotlib color representing the datapoints clipped by the upper limit. Default value is read from ``config.plotting.auto_color_scale.color_over``. color_under: Matplotlib color representing the datapoints clipped by the lower limit. Default value is read from ``config.plotting.auto_color_scale.color_under``. """ import qcodes if colorbar is None: _LOG.warning( '"auto_color_scale_from_config" did not receive a colorbar ' "for scaling. Are you trying to scale a plot without " "colorbar?" ) return if auto_color_scale is None: auto_color_scale = qcodes.config.plotting.auto_color_scale.enabled if not auto_color_scale: return if color_over is None: color_over = qcodes.config.plotting.auto_color_scale.color_over if color_under is None: color_under = qcodes.config.plotting.auto_color_scale.color_under if cutoff_percentile is None: cutoff_percentile = cast( tuple[float, float], tuple(qcodes.config.plotting.auto_color_scale.cutoff_percentile), ) apply_auto_color_scale( colorbar, data_array, cutoff_percentile, color_over, color_under )