Repository: wright-group/WrightTools

Issue:
Plotting for lower rank channels
I will explain by way of example:
I have a data object whose shape is (10,100,1000). Call these dimensions `w1`, `w2`, ans `z`
Within this data object, I may have channels which depend on all three, but may also have channels which depend on only `w1` and `w2`, for instance. These channels would be represented by an HDF5 dataset with shape (10,100,1).

I would like to be able to `data.transform('w1', 'w2')` and call `ax.pcolor(data, channel=ch)`

I cannot do this currently as the arrays passed to matplotlib are 3D arrays (though one dimension is 1)


Discussion:
Transforming to a set of axes which all have a 1 in a given position of the shape should, effectively, ignore this axis where it can (perhaps an error if there is something attempted which cannot work, though mostly these are needed for plotting, as far as I'm concerned)

this should work---place within `pcolor` and other ax methods

```python
if channel.ndim > 2:
    zi = channel.points
else:
    zi = channel.full
```
Not sure this works exactly, for instance, a 10x1x1 should still be 10x1 under these conditions.
Well, it would work, but may need to transpose in some instances (and would also need to do `.points` for the axes)

perhaps I don't understand your goal

if the axes are 3D I don't know how to make a 2D heatmap

Resolving patch:
diff --git a/WrightTools/artists/_base.py b/WrightTools/artists/_base.py
--- a/WrightTools/artists/_base.py
+++ b/WrightTools/artists/_base.py
@@ -108,22 +108,28 @@ def _parse_limits(self, zi=None, data=None, channel_index=None, dynamic_range=Fa
 
     def _parse_plot_args(self, *args, **kwargs):
         plot_type = kwargs.pop("plot_type")
-        if plot_type not in ["pcolor", "pcolormesh"]:
+        if plot_type not in ["pcolor", "pcolormesh", "contourf", "contour"]:
             raise NotImplementedError
         args = list(args)  # offer pop, append etc
         dynamic_range = kwargs.pop("dynamic_range", False)
         if isinstance(args[0], Data):
             data = args.pop(0)
-            if plot_type in ["pcolor", "pcolormesh"]:
-                ndim = 2
-            if not data.ndim == ndim:
-                raise wt_exceptions.DimensionalityError(ndim, data.ndim)
-            # arrays
             channel = kwargs.pop("channel", 0)
             channel_index = wt_kit.get_index(data.channel_names, channel)
-            zi = data.channels[channel_index][:]
-            xi = data.axes[0].full
-            yi = data.axes[1].full
+            squeeze = np.array(data.channels[channel_index].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")
+            squeeze = tuple([0 if i else slice(None) for i in squeeze])
+            zi = data.channels[channel_index].points
+            xi = xa.full[squeeze]
+            yi = ya.full[squeeze]
+            if plot_type in ["pcolor", "pcolormesh", "contourf", "contour"]:
+                ndim = 2
+            if not zi.ndim == ndim:
+                raise wt_exceptions.DimensionalityError(ndim, data.ndim)
             if plot_type in ["pcolor", "pcolormesh"]:
                 X, Y = pcolor_helper(xi, yi)
             else:
@@ -133,8 +139,23 @@ def _parse_plot_args(self, *args, **kwargs):
             kwargs = self._parse_limits(
                 data=data, channel_index=channel_index, dynamic_range=dynamic_range, **kwargs
             )
-            # cmap
-            kwargs = self._parse_cmap(data=data, channel_index=channel_index, **kwargs)
+            if plot_type == "contourf":
+                if "levels" 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:
+                        n = 11
+                    else:
+                        n = 6
+                    kwargs["levels"] = np.linspace(kwargs.pop("vmin"), kwargs.pop("vmax"), n)[1:-1]
+                # colors
+                if "colors" not in kwargs.keys():
+                    kwargs["colors"] = "k"
+                if "alpha" not in kwargs.keys():
+                    kwargs["alpha"] = 0.5
+            if plot_type in ["pcolor", "pcolormesh", "contourf"]:
+                kwargs = self._parse_cmap(data=data, channel_index=channel_index, **kwargs)
         else:
             xi, yi, zi = args[:3]
             if plot_type in ["pcolor", "pcolormesh"]:
@@ -147,7 +168,11 @@ def _parse_plot_args(self, *args, **kwargs):
             data = None
             channel_index = 0
             kwargs = self._parse_limits(zi=args[2], **kwargs)
-            kwargs = self._parse_cmap(**kwargs)
+            if plot_type == "contourf":
+                if "levels" not in kwargs.keys():
+                    kwargs["levels"] = np.linspace(kwargs["vmin"], kwargs["vmax"], 256)
+            if plot_type in ["pcolor", "pcolormesh", "contourf"]:
+                kwargs = self._parse_cmap(**kwargs)
         # labels
         self._apply_labels(
             autolabel=kwargs.pop("autolabel", False),
@@ -157,7 +182,8 @@ def _parse_plot_args(self, *args, **kwargs):
             channel_index=channel_index,
         )
         # decoration
-        self.set_facecolor([0.75] * 3)
+        if plot_type != "contour":
+            self.set_facecolor([0.75] * 3)
         return args, kwargs
 
     def add_sideplot(self, along, pad=0, height=0.75, ymin=0, ymax=1.1):
@@ -200,6 +226,11 @@ def add_sideplot(self, along, pad=0, height=0.75, ymin=0, ymax=1.1):
     def contour(self, *args, **kwargs):
         """Plot contours.
 
+        If a 3D or higher Data object is passed, a lower dimensional
+        channel can be plotted, provided the ``squeeze`` of the channel
+        has ``ndim==2`` and the first two axes do not span dimensions
+        other than those spanned by that channel.
+
         Parameters
         ----------
         data : 2D WrightTools.data.Data object
@@ -224,53 +255,17 @@ def contour(self, *args, **kwargs):
         -------
         matplotlib.contour.QuadContourSet
         """
-        args = list(args)  # offer pop, append etc
-        channel = kwargs.pop("channel", 0)
-        dynamic_range = kwargs.pop("dynamic_range", False)
-        # unpack data object, if given
-        if isinstance(args[0], Data):
-            data = args.pop(0)
-            if not data.ndim == 2:
-                raise wt_exceptions.DimensionalityError(2, data.ndim)
-            # arrays
-            channel_index = wt_kit.get_index(data.channel_names, channel)
-            signed = data.channels[channel_index].signed
-            xi = data.axes[0].full
-            yi = data.axes[1].full
-            zi = data.channels[channel_index][:]
-            args = [xi, yi, zi] + args
-            # limits
-            kwargs = self._parse_limits(
-                data=data, channel_index=channel_index, dynamic_range=dynamic_range, **kwargs
-            )
-            # levels
-            if "levels" not in kwargs.keys():
-                if signed:
-                    n = 11
-                else:
-                    n = 6
-                kwargs["levels"] = np.linspace(kwargs.pop("vmin"), kwargs.pop("vmax"), n)[1:-1]
-            # colors
-            if "colors" not in kwargs.keys():
-                kwargs["colors"] = "k"
-            if "alpha" not in kwargs.keys():
-                kwargs["alpha"] = 0.5
-            # 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,
-            )
-        else:
-            kwargs = self._parse_limits(zi=args[2], dynamic_range=dynamic_range, **kwargs)
-        # call parent
+        args, kwargs = self._parse_plot_args(*args, **kwargs, plot_type="contour")
         return super().contour(*args, **kwargs)
 
     def contourf(self, *args, **kwargs):
         """Plot contours.
 
+        If a 3D or higher Data object is passed, a lower dimensional
+        channel can be plotted, provided the ``squeeze`` of the channel
+        has ``ndim==2`` and the first two axes do not span dimensions
+        other than those spanned by that channel.
+
         Parameters
         ----------
         data : 2D WrightTools.data.Data object
@@ -295,44 +290,7 @@ def contourf(self, *args, **kwargs):
         -------
         matplotlib.contour.QuadContourSet
         """
-        args = list(args)  # offer pop, append etc
-        channel = kwargs.pop("channel", 0)
-        dynamic_range = kwargs.pop("dynamic_range", False)
-        # unpack data object, if given
-        if isinstance(args[0], Data):
-            data = args.pop(0)
-            if not data.ndim == 2:
-                raise wt_exceptions.DimensionalityError(2, data.ndim)
-            # arrays
-            channel_index = wt_kit.get_index(data.channel_names, channel)
-            xi = data.axes[0].full
-            yi = data.axes[1].full
-            zi = data.channels[channel_index][:]
-            args = [xi, yi, zi] + args
-            # limits
-            kwargs = self._parse_limits(
-                data=data, channel_index=channel_index, dynamic_range=dynamic_range, **kwargs
-            )
-            # cmap
-            kwargs = self._parse_cmap(data=data, channel_index=channel_index, **kwargs)
-        else:
-            data = None
-            channel_index = 0
-            kwargs = self._parse_limits(zi=args[2], dynamic_range=dynamic_range, **kwargs)
-            kwargs = self._parse_cmap(**kwargs)
-        # levels
-        if "levels" not in kwargs.keys():
-            vmin = kwargs.pop("vmin", args[2].min())
-            vmax = kwargs.pop("vmax", args[2].max())
-            kwargs["levels"] = np.linspace(vmin, vmax, 256)
-        # 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,
-        )
+        args, kwargs = self._parse_plot_args(*args, **kwargs, plot_type="contourf")
         # Overloading contourf in an attempt to fix aliasing problems when saving vector graphics
         # see https://stackoverflow.com/questions/15822159
         # also see https://stackoverflow.com/a/32911283
@@ -386,6 +344,11 @@ def legend(self, *args, **kwargs):
     def pcolor(self, *args, **kwargs):
         """Create a pseudocolor plot of a 2-D array.
 
+        If a 3D or higher Data object is passed, a lower dimensional
+        channel can be plotted, provided the ``squeeze`` of the channel
+        has ``ndim==2`` and the first two axes do not span dimensions
+        other than those spanned by that channel.
+
         Uses pcolor_helper to ensure that color boundaries are drawn
         bisecting point positions, when possible.
 
@@ -419,6 +382,11 @@ def pcolor(self, *args, **kwargs):
     def pcolormesh(self, *args, **kwargs):
         """Create a pseudocolor plot of a 2-D array.
 
+        If a 3D or higher Data object is passed, a lower dimensional
+        channel can be plotted, provided the ``squeeze`` of the channel
+        has ``ndim==2`` and the first two axes do not span dimensions
+        other than those spanned by that channel.
+
         Uses pcolor_helper to ensure that color boundaries are drawn
         bisecting point positions, when possible.
         Quicker than pcolor
@@ -453,6 +421,11 @@ def pcolormesh(self, *args, **kwargs):
     def plot(self, *args, **kwargs):
         """Plot lines and/or markers.
 
+        If a 2D or higher Data object is passed, a lower dimensional
+        channel can be plotted, provided the ``squeeze`` of the channel
+        has ``ndim==1`` and the first axis does not span dimensions
+        other than that spanned by the channel.
+
         Parameters
         ----------
         data : 1D WrightTools.data.Data object
@@ -480,15 +453,20 @@ def plot(self, *args, **kwargs):
         """
         args = list(args)  # offer pop, append etc
         # unpack data object, if given
-        if hasattr(args[0], "id"):  # TODO: replace once class comparison works...
+        if isinstance(args[0], Data):
             data = args.pop(0)
             channel = kwargs.pop("channel", 0)
-            if not data.ndim == 1:
-                raise wt_exceptions.DimensionalityError(1, data.ndim)
-            # arrays
             channel_index = wt_kit.get_index(data.channel_names, channel)
-            xi = data.axes[0][:]
-            zi = data.channels[channel_index][:].T
+            squeeze = np.array(data.channels[channel_index].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
+            xi = xa[squeeze]
+            if not zi.ndim == 1:
+                raise wt_exceptions.DimensionalityError(1, data.ndim)
             args = [xi, zi] + args
         else:
             data = None
diff --git a/WrightTools/artists/_quick.py b/WrightTools/artists/_quick.py
--- a/WrightTools/artists/_quick.py
+++ b/WrightTools/artists/_quick.py
@@ -65,10 +65,14 @@ def quick1D(
     list of strings
         List of saved image files (if any).
     """
-    # prepare data
-    chopped = data.chop(axis, at=at, verbose=False)
     # channel index
     channel_index = wt_kit.get_index(data.channel_names, channel)
+    shape = data.channels[channel_index].shape
+    collapse = [i for i in range(len(shape)) if shape[i] == 1]
+    at = at.copy()
+    at.update({c: 0 for c in collapse})
+    # prepare data
+    chopped = data.chop(axis, at=at, verbose=False)
     # prepare figure
     fig = None
     if len(chopped) > 10:
@@ -87,7 +91,7 @@ def quick1D(
                 else:
                     fname = data.natural_name
             else:
-                folder_name = "mpl_1D " + wt_kit.TimeStamp().path
+                folder_name = "quick1D " + wt_kit.TimeStamp().path
                 os.mkdir(folder_name)
                 save_directory = folder_name
     # chew through image generation
@@ -191,10 +195,14 @@ def quick2D(
     list of strings
         List of saved image files (if any).
     """
-    # prepare data
-    chopped = data.chop(xaxis, yaxis, at=at, verbose=False)
     # channel index
     channel_index = wt_kit.get_index(data.channel_names, channel)
+    shape = data.channels[channel_index].shape
+    collapse = [i for i in range(len(shape)) if shape[i] == 1]
+    at = at.copy()
+    at.update({c: 0 for c in collapse})
+    # prepare data
+    chopped = data.chop(xaxis, yaxis, at=at, verbose=False)
     # colormap
     # get colormap
     if data.channels[channel_index].signed:
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -358,7 +358,7 @@ def chop(self, *args, at={}, parent=None, verbose=True) -> wt_collection.Collect
                 args[i] = wt_kit.string2identifier(arg)
 
         # normalize the at keys to the natural name
-        for k in list(at.keys()):
+        for k in [ak for ak in at.keys() if type(ak) == str]:
             for op in operators:
                 if op in k:
                     nk = k.replace(op, operator_to_identifier[op])
@@ -369,18 +369,26 @@ def chop(self, *args, at={}, parent=None, verbose=True) -> wt_collection.Collect
         # get output collection
         out = wt_collection.Collection(name="chop", parent=parent)
         # get output shape
-        kept = args + list(at.keys())
+        kept = args + [ak for ak in at.keys() if type(ak) == str]
         kept_axes = [self._axes[self.axis_names.index(a)] for a in kept]
         removed_axes = [a for a in self._axes if a not in kept_axes]
         removed_shape = wt_kit.joint_shape(*removed_axes)
         if removed_shape == ():
             removed_shape = (1,) * self.ndim
+        removed_shape = list(removed_shape)
+        for i in at.keys():
+            if type(i) == int:
+                removed_shape[i] = 1
+        removed_shape = tuple(removed_shape)
         # iterate
         i = 0
         for idx in np.ndindex(removed_shape):
             idx = np.array(idx, dtype=object)
             idx[np.array(removed_shape) == 1] = slice(None)
             for axis, point in at.items():
+                if type(axis) == int:
+                    idx[axis] = point
+                    continue
                 point, units = point
                 destination_units = self._axes[self.axis_names.index(axis)].units
                 point = wt_units.converter(point, units, destination_units)


Test patch:
diff --git a/tests/artists/test_contour.py b/tests/artists/test_contour.py
new file mode 100755
--- /dev/null
+++ b/tests/artists/test_contour.py
@@ -0,0 +1,51 @@
+#! /usr/bin/env python3
+
+
+import WrightTools as wt
+from WrightTools import datasets
+from matplotlib import pyplot as plt
+import shutil
+import os
+
+
+def test_contour():
+    p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+    p = shutil.copy(p, "./test_contour.wt5")
+    data = wt.open(p)
+    os.unlink(p)
+    data.level(0, 2, -3)
+    data.convert("eV")
+    data.ai0.symmetric_root(0.5)
+    chop = data.chop("w1=wm", "w2", at={"d2": [-600, "fs"]})[0]
+
+    fig, gs = wt.artists.create_figure()
+    ax = plt.subplot(gs[0])
+    ax.contourf(chop)
+    ax.contour(chop)
+
+    data.close()
+    chop.close()
+
+
+def test_contour_lower_rank():
+    p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+    p = shutil.copy(p, "./test_contour_lower_rank.wt5")
+    data = wt.open(p)
+    os.unlink(p)
+    data.level(0, 2, -3)
+    data.convert("eV")
+    data.ai0.symmetric_root(0.5)
+    data.collapse("d2", method="sum")
+
+    fig, gs = wt.artists.create_figure()
+    ax = plt.subplot(gs[0])
+    ax.contourf(data, channel="ai0_d2_sum")
+    ax.contour(data, channel="ai0_d2_sum")
+
+    data.close()
+
+
+if __name__ == "__main__":
+    test_contour()
+    test_contour_lower_rank()
+    plt.show()
diff --git a/tests/artists/test_pcolor.py b/tests/artists/test_pcolor.py
new file mode 100755
--- /dev/null
+++ b/tests/artists/test_pcolor.py
@@ -0,0 +1,51 @@
+#! /usr/bin/env python3
+
+
+import WrightTools as wt
+from WrightTools import datasets
+from matplotlib import pyplot as plt
+import shutil
+import os
+
+
+def test_pcolor():
+    p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+    p = shutil.copy(p, "./test_pcolor.wt5")
+    data = wt.open(p)
+    os.unlink(p)
+    data.level(0, 2, -3)
+    data.convert("eV")
+    data.ai0.symmetric_root(0.5)
+    chop = data.chop("w1=wm", "w2", at={"d2": [-600, "fs"]})[0]
+
+    fig, gs = wt.artists.create_figure()
+    ax = plt.subplot(gs[0])
+    ax.pcolor(chop)
+    ax.pcolormesh(chop)
+
+    data.close()
+    chop.close()
+
+
+def test_pcolor_lower_rank():
+    p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+    p = shutil.copy(p, "./test_pcolor_lower_rank.wt5")
+    data = wt.open(p)
+    os.unlink(p)
+    data.level(0, 2, -3)
+    data.convert("eV")
+    data.ai0.symmetric_root(0.5)
+    data.collapse("d2", method="sum")
+
+    fig, gs = wt.artists.create_figure()
+    ax = plt.subplot(gs[0])
+    ax.pcolor(data, channel="ai0_d2_sum")
+    ax.pcolormesh(data, channel="ai0_d2_sum")
+
+    data.close()
+
+
+if __name__ == "__main__":
+    test_pcolor()
+    test_pcolor_lower_rank()
+    plt.show()
diff --git a/tests/artists/test_plot.py b/tests/artists/test_plot.py
new file mode 100755
--- /dev/null
+++ b/tests/artists/test_plot.py
@@ -0,0 +1,51 @@
+#! /usr/bin/env python3
+
+
+import WrightTools as wt
+from WrightTools import datasets
+from matplotlib import pyplot as plt
+import shutil
+import os
+
+
+def test_plot():
+    p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+    p = shutil.copy(p, "./test_plot.wt5")
+    data = wt.open(p)
+    os.unlink(p)
+    data.level(0, 2, -3)
+    data.convert("eV")
+    data.ai0.symmetric_root(0.5)
+    chop = data.chop("w1=wm", at={"d2": [-600, "fs"], "w2": [2.0, "eV"]})[0]
+
+    fig, gs = wt.artists.create_figure()
+    ax = plt.subplot(gs[0])
+    ax.plot(chop)
+
+    data.close()
+    chop.close()
+
+
+def test_plot_lower_rank():
+    p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+    p = shutil.copy(p, "./test_plot_lower_rank.wt5")
+    data = wt.open(p)
+    os.unlink(p)
+    data.transform("w1", "w2", "d2")
+    data.level(0, 2, -3)
+    data.convert("eV")
+    data.ai0.symmetric_root(0.5)
+    data.collapse("d2", method="sum")
+    data.collapse("w2", method="sum")
+
+    fig, gs = wt.artists.create_figure()
+    ax = plt.subplot(gs[0])
+    ax.plot(data, channel="ai0_d2_sum_w2_sum")
+
+    data.close()
+
+
+if __name__ == "__main__":
+    test_plot()
+    test_plot_lower_rank()
+    plt.show()