Skip to content

[WIP] Robust data children getters #1211

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ddkohler wants to merge 6 commits into
base: master
Choose a base branch
from
Draft

[WIP] Robust data children getters #1211

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
1797fd7
Update _base.py
ddkohler Sep 22, 2025
8a8515b
Update _data.py
ddkohler Sep 22, 2025
ffe992d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 22, 2025
207ecd2
tests
ddkohler Sep 22, 2025
90cd784
tests
ddkohler Sep 22, 2025
048a27b
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 54 additions & 53 deletions WrightTools/artists/_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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])
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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":
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
55 changes: 15 additions & 40 deletions WrightTools/data/_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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[:]
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
Loading