Source code for autoarray.plot.wrap.base.colorbar

import matplotlib.pyplot as plt
import as cm
import numpy as np
from typing import List, Optional

from autoconf import conf

from autoarray.plot.wrap.base.abstract import AbstractMatWrap
from autoarray.plot.wrap.base.units import Units

from autoarray import exc

[docs]class Colorbar(AbstractMatWrap): def __init__( self, manual_tick_labels: Optional[List[float]] = None, manual_tick_values: Optional[List[float]] = None, manual_alignment: Optional[str] = None, manual_unit: Optional[str] = None, manual_log10: bool = False, **kwargs, ): """ Customizes the colorbar of the plotted figure. This object wraps the following Matplotlib method: - plt.colorbar: The colorbar object `cb` that is created is also customized using the following methods: - cb.set_yticklabels: Parameters ---------- manual_tick_labels Manually override the colorbar tick labels to an input list of float. manual_tick_values If the colorbar tick labels are manually specified the locations on the colorbar they appear running 0 -> 1. manual_alignment The vertical alignment of the colorbar tick labels, specified via the matplotlib method `set_yticklabels` and input `va`. manual_unit The unit label that appears next to the colorbar tick labels, which if not input uses a default unit label specified as `cb_unit` in the config file `config/visualize/general.yaml. """ super().__init__(**kwargs) self.manual_tick_labels = manual_tick_labels self.manual_tick_values = manual_tick_values self.manual_alignment = manual_alignment self.manual_unit = manual_unit self.manual_log10 = manual_log10 @property def cb_unit(self): if self.manual_unit is None: return conf.instance["visualize"]["general"]["units"]["cb_unit"] return self.manual_unit def tick_values_from(self, norm=None, use_log10: bool = False): if ( sum( x is not None for x in [self.manual_tick_values, self.manual_tick_labels] ) == 1 ): raise exc.PlottingException( "You can only manually specify the colorbar tick labels and values if both are input." ) if self.manual_tick_values is not None: return self.manual_tick_values if norm is not None: min_value = norm.vmin max_value = norm.vmax if use_log10: if min_value < self.log10_min_value: min_value = self.log10_min_value log_mid_value = (np.log10(max_value) + np.log10(min_value)) / 2.0 mid_value = 10**log_mid_value else: mid_value = (max_value + min_value) / 2.0 return [min_value, mid_value, max_value] def tick_labels_from( self, units: Units, manual_tick_values: List[float], cb_unit=None, ): if manual_tick_values is None: return None convert_factor = units.colorbar_convert_factor or 1.0 if self.manual_tick_labels is not None: manual_tick_labels = self.manual_tick_labels else: manual_tick_labels = [ np.round(value * convert_factor, 2) for value in manual_tick_values ] if self.manual_log10: manual_tick_labels = [ "{:.0e}".format(label) for label in manual_tick_labels ] manual_tick_labels = [ label.replace("1e", "$10^{") + "}$" for label in manual_tick_labels ] manual_tick_labels = [ label.replace("{-0", "{-").replace("{+0", "{+").replace("+", "") for label in manual_tick_labels ] if units.colorbar_label is None: if cb_unit is None: cb_unit = self.cb_unit else: cb_unit = units.colorbar_label middle_index = (len(manual_tick_labels) - 1) // 2 manual_tick_labels[ middle_index ] = rf"{manual_tick_labels[middle_index]}{cb_unit}" return manual_tick_labels
[docs] def set( self, units: Units, ax=None, norm=None, cb_unit=None, use_log10: bool = False ): """ Set the figure's colorbar, optionally overriding the tick labels and values with manual inputs. """ tick_values = self.tick_values_from(norm=norm, use_log10=use_log10) tick_labels = self.tick_labels_from( manual_tick_values=tick_values, units=units, cb_unit=cb_unit, ) if tick_values is None and tick_labels is None: cb = plt.colorbar(ax=ax, **self.config_dict) else: cb = plt.colorbar(ticks=tick_values, ax=ax, **self.config_dict) labels=tick_labels, va=self.manual_alignment or "center" ) return cb
[docs] def set_with_color_values( self, units: Units, cmap: str, color_values: np.ndarray, ax=None, norm=None, use_log10: bool = False, ): """ Set the figure's colorbar using an array of already known color values. This method is used for producing the color bar on a Voronoi mesh plot, which is unable to use the in-built Matplotlib colorbar method. Parameters ---------- cmap The colormap used to map normalized data values to RGBA colors (see color_values The values of the pixels on the Voronoi mesh which are used to create the colorbar. """ mappable = cm.ScalarMappable(norm=norm, cmap=cmap) mappable.set_array(color_values) tick_values = self.tick_values_from(norm=norm, use_log10=use_log10) tick_labels = self.tick_labels_from( manual_tick_values=tick_values, units=units, ) if tick_values is None and tick_labels is None: cb = plt.colorbar( mappable=mappable, ax=ax, **self.config_dict, ) else: cb = plt.colorbar( mappable=mappable, ax=ax, ticks=tick_values, **self.config_dict, ) labels=tick_labels, va=self.manual_alignment or "center" ) return cb