diff --git a/WrightTools/artists/_base.py b/WrightTools/artists/_base.py index 3b16a23f..cbc73348 100644 --- a/WrightTools/artists/_base.py +++ b/WrightTools/artists/_base.py @@ -10,10 +10,11 @@ from matplotlib.colors import Normalize import matplotlib.pyplot as plt from mpl_toolkits.axes_grid1 import make_axes_locatable +from typing import Literal from .. import exceptions as wt_exceptions from .. import kit as wt_kit -from ..data import Data +from ..data import Data, Channel from ._colors import colormaps @@ -35,42 +36,50 @@ class Axes(matplotlib.axes.Axes): is_sideplot = False def _apply_labels( - self, autolabel="none", xlabel=None, ylabel=None, data=None, channel_index=0 + self, + autolabel: Literal["both", "x", "y"] = None, + xlabel: str = None, + ylabel: str = None, + data: Data = None, + channel: Channel = None, ): """Apply x and y labels to axes. + Autolabels are ignored if xlabel or ylabel are provided. + + NOTE: if autolabel is requested, data is required! + + NOTE: if autolabel is requested and data is 1D, channel is required! Parameters ---------- - autolabel : {'none', 'both', 'x', 'y'} (optional) - Label(s) to apply from data. Default is none. - xlabel : string (optional) - x label. Default is None. + autolabel : {'both', 'x', 'y'} (optional) + Label(s) to apply from data. Default is no labels. + xlabel : string or bool (optional) + x label. Default is None. if a string, xlabel overrides autolabel ylabel : string (optional) - y label. Default is None. + y label. Default is None. if a string, autolabel is overwritten with ylabel data : WrightTools.data.Data object (optional) - data to read labels from. Default is None. - channel_index : integer (optional) - Channel index. Default is 0. + data to autolabel from. Default is None. + channel : WrightTools.data.Channel object (optional) + Channel index. Default is None. """ # read from data if autolabel in ["xy", "both", "x"] and not xlabel: xlabel = data.axes[0].label if autolabel in ["xy", "both", "y"] and not ylabel: if data.ndim == 1: - ylabel = data.channels[channel_index].label + ylabel = channel.label elif data.ndim == 2: ylabel = data.axes[1].label # apply if xlabel: - if isinstance(xlabel, bool): - xlabel = data.axes[0].label - self.set_xlabel(xlabel, fontsize=18) + self.set_xlabel(xlabel) if ylabel: - if isinstance(ylabel, bool): - ylabel = data.axes[1].label - self.set_ylabel(ylabel, fontsize=18) + self.set_ylabel(ylabel) - def _parse_limits(self, zi=None, data=None, channel_index=None, dynamic_range=False, **kwargs): + def _parse_limits( + self, zi=None, data: Data = None, channel: Channel = None, dynamic_range=False, **kwargs + ): if "norm" in kwargs: return kwargs if zi is not None: @@ -82,17 +91,17 @@ def _parse_limits(self, zi=None, data=None, channel_index=None, dynamic_range=Fa vmin = np.nanmin(zi) vmax = np.nanmax(zi) elif data is not None: - signed = data.channels[channel_index].signed - null = data.channels[channel_index].null + signed = channel.signed + null = channel.null if signed and dynamic_range: - vmin = -data.channels[channel_index].minor_extent + null - vmax = +data.channels[channel_index].minor_extent + null + vmin = -channel.minor_extent + null + vmax = +channel.minor_extent + null elif signed and not dynamic_range: - vmin = -data.channels[channel_index].major_extent + null - vmax = +data.channels[channel_index].major_extent + null + vmin = -channel.major_extent + null + vmax = +channel.major_extent + null else: vmin = null - vmax = data.channels[channel_index].max() + vmax = channel.max() # don't overwrite if "vmin" not in kwargs.keys(): kwargs["vmin"] = vmin @@ -108,15 +117,14 @@ def _parse_plot_args(self, *args, **kwargs): dynamic_range = kwargs.pop("dynamic_range", False) if isinstance(args[0], Data): data = args.pop(0) - channel = kwargs.pop("channel", 0) - channel_index = wt_kit.get_index(data.channel_names, channel) - squeeze = np.array(data.channels[channel_index].shape) == 1 + channel_: Channel = data.get_channel(kwargs.pop("channel", 0)) + squeeze = np.array(channel_.shape) == 1 xa = data.axes[0] ya = data.axes[1] for sq, xs, ys in zip(squeeze, xa.shape, ya.shape): if sq and (xs != 1 or ys != 1): raise wt_exceptions.ValueError("Cannot squeeze axis to fit channel") - zi = data.channels[channel_index].points + zi = channel_.points if not zi.ndim == 2: raise wt_exceptions.DimensionalityError(2, zi.ndim) squeeze = tuple([0 if i else slice(None) for i in squeeze]) @@ -151,14 +159,14 @@ def _parse_plot_args(self, *args, **kwargs): args = [xi, yi, zi] + args # limits kwargs = self._parse_limits( - data=data, channel_index=channel_index, dynamic_range=dynamic_range, **kwargs + data=data, channel=channel_, dynamic_range=dynamic_range, **kwargs ) if plot_type == "contourf": if "levels" not in kwargs.keys() and "norm" not in kwargs.keys(): kwargs["levels"] = np.linspace(kwargs["vmin"], kwargs["vmax"], 256) elif plot_type == "contour": if "levels" not in kwargs.keys(): - if data.channels[channel_index].signed: + if channel_.signed: n = 11 else: n = 6 @@ -169,14 +177,13 @@ def _parse_plot_args(self, *args, **kwargs): if "alpha" not in kwargs.keys(): kwargs["alpha"] = 0.5 if plot_type in ["pcolor", "pcolormesh", "contourf", "imshow"]: - kwargs = _parse_cmap(data=data, channel_index=channel_index, **kwargs) + kwargs = _parse_cmap(data=data, signed=channel_.signed, **kwargs) else: if plot_type == "imshow": kwargs = self._parse_limits(zi=args[0], **kwargs) else: kwargs = self._parse_limits(zi=args[2], **kwargs) - data = None - channel_index = 0 + data = channel_ = None if plot_type == "contourf": if "levels" not in kwargs.keys(): kwargs["levels"] = np.linspace(kwargs["vmin"], kwargs["vmax"], 256) @@ -188,7 +195,7 @@ def _parse_plot_args(self, *args, **kwargs): xlabel=kwargs.pop("xlabel", None), ylabel=kwargs.pop("ylabel", None), data=data, - channel_index=channel_index, + channel=channel_, ) if plot_type != "contour": @@ -507,16 +514,15 @@ def scatter(self, *args, **kwargs): Use `cmap` instead to control colors." ) - channel = kwargs.pop("channel", 0) - channel_index = wt_kit.get_index(data.channel_names, channel) + channel = data.get_channel(kwargs.pop("channel", 0)) limits = {} - limits = self._parse_limits(data=data, channel_index=channel_index, **limits) + limits = self._parse_limits(data=data, channel=channel, **limits) norm = Normalize(**limits) - cmap = _parse_cmap(data, channel_index=channel_index, **kwargs)["cmap"] + cmap = _parse_cmap(data, signed=channel.signed, **kwargs)["cmap"] - z = data.channels[channel_index][:] + z = channel[:] # fill x, y, z to joint shape shape = wt_kit.joint_shape(z, *coords) @@ -539,7 +545,7 @@ def full(arr, shape): xlabel=kwargs.pop("xlabel", None), ylabel=kwargs.pop("ylabel", None), data=data, - channel_index=channel_index, + channel=channel, ) return super().scatter(*args, **kwargs) @@ -621,29 +627,27 @@ def plot(self, *args, **kwargs): # unpack data object, if given if isinstance(args[0], Data): data = args.pop(0) - channel = kwargs.pop("channel", 0) - channel_index = wt_kit.get_index(data.channel_names, channel) - squeeze = np.array(data.channels[channel_index].shape) == 1 + channel = data.get_channel(kwargs.pop("channel", 0)) + squeeze = np.array(channel.shape) == 1 xa = data.axes[0] for sq, xs in zip(squeeze, xa.shape): if sq and xs != 1: raise wt_exceptions.ValueError("Cannot squeeze axis to fit channel") squeeze = tuple([0 if i else slice(None) for i in squeeze]) - zi = data.channels[channel_index].points + zi = channel.points xi = xa[squeeze] if not zi.ndim == 1: raise wt_exceptions.DimensionalityError(1, zi.ndim) args = [xi, zi] + args else: - data = None - channel_index = 0 + data = channel = None # labels self._apply_labels( autolabel=kwargs.pop("autolabel", False), xlabel=kwargs.pop("xlabel", None), ylabel=kwargs.pop("ylabel", None), data=data, - channel_index=channel_index, + channel=channel, ) # call parent return super().plot(*args, **kwargs) @@ -734,13 +738,10 @@ def _order_for_imshow(xi, yi): raise TypeError(f"Axes are not 1D: {xi.shape}, {yi.shape}") -def _parse_cmap(data=None, channel_index=None, **kwargs): +def _parse_cmap(data=None, signed=None, **kwargs): if "cmap" in kwargs.keys(): if isinstance(kwargs["cmap"], str): kwargs["cmap"] = colormaps[kwargs["cmap"]] elif data: - if data.channels[channel_index].signed: - kwargs["cmap"] = colormaps["signed"] - return kwargs - kwargs["cmap"] = colormaps["default"] + kwargs["cmap"] = colormaps["signed"] if signed else colormaps["default"] return kwargs diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py index bddf090b..c8b0155f 100644 --- a/WrightTools/data/_data.py +++ b/WrightTools/data/_data.py @@ -717,22 +717,20 @@ def gradient(self, axis, *, channel=0): else: raise wt_exceptions.TypeError("axis: expected {int, str}, got %s" % type(axis)) - channel_index = wt_kit.get_index(self.channel_names, channel) - channel = self.channel_names[channel_index] + channel = self.get_channel(channel) - if self[channel].shape[axis_index] == 1: + if channel.shape[axis_index] == 1: raise wt_exceptions.ValueError( "Channel '{}' has a single point along Axis '{}', cannot compute gradient".format( channel, axis ) ) - rtype = np.result_type(self[channel].dtype, float) + rtype = np.result_type(channel.dtype, float) new = self.create_channel( - "{}_{}_gradient".format(channel, axis), - values=np.empty(self[channel].shape, dtype=rtype), + "{}_{}_gradient".format(channel.natural_name, axis), + values=np.empty(channel.shape, dtype=rtype), ) - channel = self[channel] if axis == axis_index: new[:] = np.gradient(channel[:], axis=axis_index) else: @@ -882,20 +880,18 @@ def moment(self, axis, channel=0, moment=1, *, resultant=None): warnings.warn("moment", category=wt_exceptions.EntireDatasetInMemoryWarning) - channel_index = wt_kit.get_index(self.channel_names, channel) - channel = self.channel_names[channel_index] + channel: Channel = self.get_channel(channel) - if self[channel].shape[axis_index] == 1: + if channel.shape[axis_index] == 1: raise wt_exceptions.ValueError( "Channel '{}' has a single point along Axis '{}', cannot compute moment".format( channel, axis ) ) - new_shape = list(self[channel].shape) + new_shape = list(channel.shape) new_shape[axis_index] = 1 - channel = self[channel] axis_inp = axis axis = self.axes[index] x = axis[:] @@ -1268,13 +1264,7 @@ def get_nadir(self, channel=0) -> tuple: Coordinates in units for each axis. """ # get channel - if isinstance(channel, int): - channel_index = channel - elif isinstance(channel, str): - channel_index = self.channel_names.index(channel) - else: - raise TypeError("channel: expected {int, str}, got %s" % type(channel)) - channel = self.channels[channel_index] + channel = self.get_channel(channel) # get indicies idx = channel.argmin() # finish @@ -1294,13 +1284,7 @@ def get_zenith(self, channel=0) -> tuple: Coordinates in units for each axis. """ # get channel - if isinstance(channel, int): - channel_index = channel - elif isinstance(channel, str): - channel_index = self.channel_names.index(channel) - else: - raise TypeError("channel: expected {int, str}, got %s" % type(channel)) - channel = self.channels[channel_index] + channel: Channel = self.get_channel(channel) # get indicies idx = channel.argmax() # finish @@ -1336,15 +1320,8 @@ def heal(self, channel=0, method="linear", fill_value=np.nan, verbose=True): warnings.warn("heal", category=wt_exceptions.EntireDatasetInMemoryWarning) timer = wt_kit.Timer(verbose=False) with timer: - # channel - if isinstance(channel, int): - channel_index = channel - elif isinstance(channel, str): - channel_index = self.channel_names.index(channel) - else: - raise TypeError("channel: expected {int, str}, got %s" % type(channel)) - channel = self.channels[channel_index] - values = self.channels[channel_index][:] + channel = self.get_channel(channel) + values = channel[:] points = [axis[:] for axis in self._axes] xi = tuple(np.meshgrid(*points, indexing="ij")) # 'undo' gridding @@ -1358,7 +1335,7 @@ def heal(self, channel=0, method="linear", fill_value=np.nan, verbose=True): tup = tuple([arr[i] for i in range(len(arr) - 1)]) # grid data out = griddata(tup, arr[-1], xi, method=method, fill_value=fill_value) - self.channels[channel_index][:] = out + channel[:] = out # print if verbose: print( @@ -1384,8 +1361,7 @@ def level(self, channel, axis, npts, *, verbose=True): Toggle talkback. Default is True. """ warnings.warn("level", category=wt_exceptions.EntireDatasetInMemoryWarning) - channel_index = wt_kit.get_index(self.channel_names, channel) - channel = self.channels[channel_index] + channel = self.get_channel(channel) # verify npts not zero npts = int(npts) if npts == 0: @@ -1443,8 +1419,7 @@ def map_variable( New data object. """ # get variable index - variable_index = wt_kit.get_index(self.variable_names, variable) - variable = self.variables[variable_index] + variable = self.get_var(variable) # get points if isinstance(points, int): points = np.linspace(variable.min(), variable.max(), points)