Repository: wright-group/WrightTools

Issue:
`Data.ichop`: chop iterator for improved performance
The main idea is chopping creates an iterator rather than a wt Collection.  

One of the first places this could see use is a refactor of artists.quicknD methods.  This could be a good step towards making that code more modular and easier to maintain.

Not sure if this is worthwhile, but we could make quicknD functions that themselves are image iterators.  This would expose the looping controls to the user, and make the core plotting code of quicknD functions easily reused.



Discussion:

Resolving patch:
diff --git a/WrightTools/artists/_base.py b/WrightTools/artists/_base.py
--- a/WrightTools/artists/_base.py
+++ b/WrightTools/artists/_base.py
@@ -34,17 +34,6 @@ class Axes(matplotlib.axes.Axes):
     transposed = False
     is_sideplot = False
 
-    def _parse_cmap(self, data=None, channel_index=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"]
-        return kwargs
-
     def _apply_labels(
         self, autolabel="none", xlabel=None, ylabel=None, data=None, channel_index=0
     ):
@@ -180,7 +169,7 @@ 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 = self._parse_cmap(data=data, channel_index=channel_index, **kwargs)
+                kwargs = _parse_cmap(data=data, channel_index=channel_index, **kwargs)
         else:
             if plot_type == "imshow":
                 kwargs = self._parse_limits(zi=args[0], **kwargs)
@@ -192,7 +181,7 @@ def _parse_plot_args(self, *args, **kwargs):
                 if "levels" not in kwargs.keys():
                     kwargs["levels"] = np.linspace(kwargs["vmin"], kwargs["vmax"], 256)
             if plot_type in ["pcolor", "pcolormesh", "contourf", "imshow"]:
-                kwargs = self._parse_cmap(**kwargs)
+                kwargs = _parse_cmap(**kwargs)
         # labels
         self._apply_labels(
             autolabel=kwargs.pop("autolabel", False),
@@ -527,7 +516,7 @@ def scatter(self, *args, **kwargs):
             limits = self._parse_limits(data=data, channel_index=channel_index, **limits)
             norm = Normalize(**limits)
 
-            cmap = self._parse_cmap(data, channel_index=channel_index, **kwargs)["cmap"]
+            cmap = _parse_cmap(data, channel_index=channel_index, **kwargs)["cmap"]
 
             z = data.channels[channel_index][:]
 
@@ -745,3 +734,15 @@ def _order_for_imshow(xi, yi):
             raise TypeError(f"x and y must be orthogonal; shapes are: {xi.shape}, {yi.shape}")
     else:
         raise TypeError(f"Axes are not 1D: {xi.shape}, {yi.shape}")
+
+
+def _parse_cmap(data=None, channel_index=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"]
+    return kwargs
diff --git a/WrightTools/artists/_quick.py b/WrightTools/artists/_quick.py
--- a/WrightTools/artists/_quick.py
+++ b/WrightTools/artists/_quick.py
@@ -3,15 +3,14 @@
 # --- import --------------------------------------------------------------------------------------
 
 
-from contextlib import closing
-import os
-
+import pathlib
 import numpy as np
-
 import matplotlib.pyplot as plt
 
+from contextlib import closing
+from functools import reduce
+
 from ._helpers import _title, create_figure, plot_colorbar, savefig
-from ._colors import colormaps
 from .. import kit as wt_kit
 
 
@@ -68,46 +67,38 @@ def quick1D(
     channel_index = wt_kit.get_index(data.channel_names, channel)
     shape = data.channels[channel_index].shape
     # remove dimensions that do not involve the channel
-    channel_slice = [0 if size == 1 else slice(None, None) for size in shape]
+    channel_slice = [0 if size == 1 else slice(None) for size in shape]
     sliced_constants = [
         data.axis_expressions[i] for i in range(len(shape)) if not channel_slice[i]
     ]
+    removed_shape = data._chop_prep(axis, at=at)[0]
+
+    len_chopped = reduce(int.__mul__, removed_shape) // reduce(int.__mul__, shape)
+    if len_chopped > 10 and not autosave:
+        print(f"expecting {len_chopped} figures.  Forcing autosave.")
+        autosave = True
+    if autosave:
+        save_directory, filepath_seed = _filepath_seed(
+            save_directory, fname if fname else data.natural_name, len_chopped, "quick1D"
+        )
+        pathlib.Path.mkdir(save_directory)
     # prepare data
-    with closing(data._from_slice(channel_slice).chop(axis, at=at, verbose=False)) as chopped:
-        # prepare figure
-        fig = None
-        if len(chopped) > 10:
-            if not autosave:
-                print("more than 10 images will be generated: forcing autosave")
-                autosave = True
-        # prepare output folders
-        if autosave:
-            if save_directory:
-                pass
-            else:
-                if len(chopped) == 1:
-                    save_directory = os.getcwd()
-                    if fname:
-                        pass
-                    else:
-                        fname = data.natural_name
-                else:
-                    folder_name = "quick1D " + wt_kit.TimeStamp().path
-                    os.mkdir(folder_name)
-                    save_directory = folder_name
-        # determine ymin and ymax for global axis scale
-        data_channel = data.channels[channel_index]
-        ymin, ymax = data_channel.min(), data_channel.max()
-        dynamic_range = ymax - ymin
-        ymin -= dynamic_range * 0.05
-        ymax += dynamic_range * 0.05
-        if np.sign(ymin) != np.sign(data_channel.min()):
-            ymin = 0
-        if np.sign(ymax) != np.sign(data_channel.max()):
-            ymax = 0
-        # chew through image generation
+    with closing(data._from_slice(channel_slice)) as sliced:
         out = []
-        for i, d in enumerate(chopped.values()):
+        for constant in sliced_constants:
+            sliced.remove_constant(constant)
+        # chew through image generation
+        for i, d in enumerate(sliced.ichop(axis, at=at)):
+            # determine ymin and ymax for global axis scale
+            data_channel = data.channels[channel_index]
+            ymin, ymax = data_channel.min(), data_channel.max()
+            dynamic_range = ymax - ymin
+            ymin -= dynamic_range * 0.05
+            ymax += dynamic_range * 0.05
+            if np.sign(ymin) != np.sign(data_channel.min()):
+                ymin = 0
+            if np.sign(ymax) != np.sign(data_channel.max()):
+                ymax = 0
             # unpack data -------------------------------------------------------------------------
             axis = d.axes[0]
             xi = axis.full
@@ -140,24 +131,18 @@ def quick1D(
                     # ignore these constants; no relation to the data
                     continue
                 ls.append(constant.label)
-                if constant.units is not None:
-                    if axis.units_kind == constant.units_kind:
-                        constant.convert(axis.units)
-                        plt.axvline(constant.value, color="k", linewidth=4, alpha=0.25)
-            title = ", ".join(ls)
-            _title(fig, data.natural_name, subtitle=title)
+                if constant.units and (axis.units_kind == constant.units_kind):
+                    constant.convert(axis.units)
+                    plt.axvline(constant.value, color="k", linewidth=4, alpha=0.25)
+            _title(fig, data.natural_name, subtitle=", ".join(ls))
             # save --------------------------------------------------------------------------------
             if autosave:
-                if fname:
-                    file_name = fname + " " + str(i).zfill(3)
-                else:
-                    file_name = str(i).zfill(3)
-                fpath = os.path.join(save_directory, file_name + ".png")
-                savefig(fpath, fig=fig, facecolor="white")
+                filepath = filepath_seed.format(i)
+                savefig(filepath, fig=fig, facecolor="white")
                 plt.close()
                 if verbose:
-                    print("image saved at", fpath)
-                out.append(fpath)
+                    print("image saved at", str(filepath))
+                out.append(str(filepath))
     return out
 
 
@@ -209,7 +194,8 @@ def quick2D(
     contours_local : boolean (optional)
         Toggle plotting black contour lines locally. Default is True.
     autosave : boolean (optional)
-         Toggle autosave. Default is False.
+        Toggle autosave. Default is False when the number of plots is 10 or less.
+        When the number of plots is greater than 10, saving is forced.
     save_directory : string (optional)
          Location to save image(s). Default is None (auto-generated).
     fname : string (optional)
@@ -226,35 +212,31 @@ def quick2D(
     channel_index = wt_kit.get_index(data.channel_names, channel)
     shape = data.channels[channel_index].shape
     # remove axes that are independent of channel
-    channel_slice = [0 if size == 1 else slice(None, None) for size in shape]
+    channel_slice = [0 if size == 1 else slice(None) for size in shape]
     sliced_constants = [
         data.axis_expressions[i] for i in range(len(shape)) if not channel_slice[i]
     ]
-    with closing(
-        data._from_slice(channel_slice).chop(xaxis, yaxis, at=at, verbose=False)
-    ) as chopped:
-        # fname
-        if fname is None:
-            fname = data.natural_name
-        # autosave
-        if len(chopped) > 10:
-            if not autosave:
-                print("more than 10 images will be generated: forcing autosave")
-                autosave = True
-        # output folder
-        if autosave:
-            if save_directory:
-                pass
-            else:
-                if len(chopped) == 1:
-                    save_directory = os.getcwd()
-                else:
-                    folder_name = "quick2D " + wt_kit.TimeStamp().path
-                    os.mkdir(folder_name)
-                    save_directory = folder_name
-        # loop through image generation
+    # determine saving, prepare for saving
+    removed_shape = data._chop_prep(xaxis, yaxis, at=at)[0]
+    len_chopped = reduce(int.__mul__, removed_shape) // reduce(int.__mul__, shape)
+    if len_chopped > 10 and not autosave:
+        print(f"expecting {len_chopped} figures.  Forcing autosave.")
+        autosave = True
+    if autosave:
+        save_directory, filepath_seed = _filepath_seed(
+            save_directory, fname if fname else data.natural_name, len_chopped, "quick2D"
+        )
+        pathlib.Path.mkdir(save_directory)
+
+    kwargs = {}
+    if cmap:
+        kwargs["cmap"] = cmap
+
+    with closing(data._from_slice(channel_slice)) as sliced:
         out = []
-        for i, d in enumerate(chopped.values()):
+        for constant in sliced_constants:
+            sliced.remove_constant(constant)
+        for i, d in enumerate(sliced.ichop(xaxis, yaxis, at=at)):
             # unpack data -------------------------------------------------------------------------
             xaxis = d.axes[0]
             xlim = xaxis.min(), xaxis.max()
@@ -277,55 +259,18 @@ def quick2D(
                 width="single", nrows=1, cols=[1, "cbar"], aspects=[[[0, 0], aspect]]
             )
             ax = plt.subplot(gs[0])
-            ax.patch.set_facecolor("w")
-            # levels ------------------------------------------------------------------------------
-            if channel.signed:
-                if local:
-                    limit = channel.mag()
-                else:
-                    data_channel = data.channels[channel_index]
-                    if dynamic_range:
-                        limit = min(
-                            abs(data_channel.null - data_channel.min()),
-                            abs(data_channel.null - data_channel.max()),
-                        )
-                    else:
-                        limit = data_channel.mag()
-                levels = np.linspace(-limit + channel.null, limit + channel.null, 200)
-            else:
-                if local:
-                    levels = np.linspace(channel.null, np.nanmax(zi), 200)
-                else:
-                    data_channel = data.channels[channel_index]
-                    if data_channel.max() < data_channel.null:
-                        levels = np.linspace(data_channel.min(), data_channel.null, 200)
-                    else:
-                        levels = np.linspace(data_channel.null, data_channel.max(), 200)
+            # ax.patch.set_facecolor("w")
             # colors ------------------------------------------------------------------------------
+            levels = determine_levels(channel, data.channels[channel_index], dynamic_range, local)
             if pixelated:
-                ax.pcolor(
-                    d, channel=channel_index, cmap=cmap, vmin=levels.min(), vmax=levels.max()
-                )
+                ax.pcolor(d, channel=channel_index, vmin=levels.min(), vmax=levels.max(), **kwargs)
             else:
-                ax.contourf(d, channel=channel_index, cmap=cmap, levels=levels)
+                ax.contourf(d, channel=channel_index, levels=levels, **kwargs)
             # contour lines -----------------------------------------------------------------------
             if contours:
-                # get contour levels
-                # force top and bottom contour to be data range then clip them out
-                if channel.signed:
-                    if contours_local:
-                        limit = channel.mag()
-                    else:
-                        limit = data_channel.mag()
-                    contour_levels = np.linspace(
-                        -limit + channel.null, limit + channel.null, contours + 2
-                    )[1:-1]
-                else:
-                    if contours_local:
-                        limit = channel.max()
-                    else:
-                        limit = data_channel.max()
-                    contour_levels = np.linspace(channel.null, limit, contours + 2)[1:-1]
+                contour_levels = determine_contour_levels(
+                    channel, data.channels[channel_index], contours_local
+                )
                 ax.contour(d, channel=channel_index, levels=contour_levels)
             # decoration --------------------------------------------------------------------------
             plt.xticks(rotation=45, fontsize=14)
@@ -346,7 +291,7 @@ def quick2D(
                     # ignore these constants; no relation to the data
                     continue
                 ls.append(constant.label)
-                if constant.units is not None:
+                if constant.units:
                     # x axis
                     if xaxis.units_kind == constant.units_kind:
                         constant.convert(xaxis.units)
@@ -355,23 +300,76 @@ def quick2D(
                     if yaxis.units_kind == constant.units_kind:
                         constant.convert(yaxis.units)
                         plt.axhline(constant.value, color="k", linewidth=4, alpha=0.25)
-            title = ", ".join(ls)
-            _title(fig, data.natural_name, subtitle=title)
+            _title(fig, data.natural_name, subtitle=", ".join(ls))
             # colorbar
             cax = plt.subplot(gs[1])
             cbar_ticks = np.linspace(levels.min(), levels.max(), 11)
-            plot_colorbar(cax=cax, ticks=cbar_ticks, label=channel.natural_name, cmap=cmap)
+            plot_colorbar(cax=cax, ticks=cbar_ticks, label=channel.natural_name, **kwargs)
             plt.sca(ax)
             # save figure -------------------------------------------------------------------------
             if autosave:
-                if fname:
-                    file_name = fname + " " + str(i).zfill(3)
-                else:
-                    file_name = str(i).zfill(3)
-                fpath = os.path.join(save_directory, file_name + ".png")
-                savefig(fpath, fig=fig, facecolor="white")
+                filepath = filepath_seed.format(i)
+                savefig(filepath, fig=fig, facecolor="white")
                 plt.close()
                 if verbose:
-                    print("image saved at", fpath)
-                out.append(fpath)
+                    print("image saved at", str(filepath))
+                out.append(str(filepath))
+            # else:
+            #     out.append(fig)
     return out
+
+
+def determine_levels(local_channel, global_channel, dynamic_range, local):
+    if local_channel.signed:
+        if local:
+            limit = local_channel.mag()
+        else:
+            if dynamic_range:
+                limit = min(
+                    abs(global_channel.null - global_channel.min()),
+                    abs(global_channel.null - global_channel.max()),
+                )
+            else:
+                limit = global_channel.mag()
+        levels = np.linspace(-limit + local_channel.null, limit + local_channel.null, 200)
+    else:
+        if local:
+            levels = np.linspace(local_channel.null, local_channel.max(), 200)
+        else:
+            if global_channel.max() < global_channel.null:
+                levels = np.linspace(global_channel.min(), global_channel.null, 200)
+            else:
+                levels = np.linspace(global_channel.null, global_channel.max(), 200)
+    return levels
+
+
+def determine_contour_levels(local_channel, global_channel, contours, local):
+    # force top and bottom contour to be data range then clip them out
+    if local_channel.signed:
+        if local:
+            limit = local_channel.mag()
+        else:
+            limit = global_channel.mag()
+        levels = np.linspace(
+            -limit + local_channel.null, limit + local_channel.null, contours + 2
+        )[1:-1]
+    else:
+        if local:
+            limit = local_channel.max()
+        else:
+            limit = global_channel.max()
+        levels = np.linspace(local_channel.null, limit, contours + 2)[1:-1]
+    return levels
+
+
+def _filepath_seed(save_directory, fname, nchops, artist):
+    """the big ugly logic block to determine the autosave filepaths"""
+    if isinstance(save_directory, str):
+        save_directory = pathlib.Path(save_directory)
+    elif save_directory is None:
+        save_directory = pathlib.Path.cwd()
+    # create a folder if multiple images
+    if nchops > 1:
+        save_directory = save_directory / f"{artist} {wt_kit.TimeStamp().path}"
+        pathlib.Path.mkdir(save_directory)
+    return save_directory, fname + " {0:0>3}.png"
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -3,6 +3,7 @@
 # --- import --------------------------------------------------------------------------------------
 
 from __future__ import annotations
+from contextlib import closing
 
 import collections
 import operator
@@ -16,6 +17,7 @@
 
 import scipy
 from scipy.interpolate import griddata, interp1d
+from typing import Generator
 
 from .._group import Group
 from .. import collection as wt_collection
@@ -431,8 +433,11 @@ def squeeze(self, name=None, parent=None):
         new.transform(*self.axis_expressions)
         return new
 
-    def chop(self, *args, at=None, parent=None, verbose=True) -> wt_collection.Collection:
-        """Divide the dataset into its lower-dimensionality components.
+    def ichop(self, *args, at=None, autoclose=True) -> Generator:
+        """
+        Similar to chop, but an iterable is produced insted of a collection.
+        Useful for workflows where only one of the chopped children is needed at a time.
+        As there is no collection, you can not assign the chopped elements to a parent.
 
         Parameters
         ----------
@@ -444,45 +449,57 @@ def chop(self, *args, at=None, parent=None, verbose=True) -> wt_collection.Colle
             Choice of position along an axis. Keys are axis names, values are lists
             ``[position, input units]``. If exact position does not exist,
             the closest valid position is used.
-        parent : WrightTools Collection instance (optional)
-            Collection to place the new "chop" collection within. Default is
-            None (new parent).
-        verbose : bool (optional)
-            Toggle talkback. Default is True.
+        autoclose : bool (optional)
+            If True, ichop will close each datafile after each iteration.
+            Set to False when you wish to store spawned Datas beyond each iteration.
+            Also consider using chop instead.
 
         Returns
         -------
-        WrightTools Collection
-            Collection of chopped data objects.
+        Generator
+            Generator spawning chopped elements for each `__next__` call
 
         Examples
         --------
-        >>> data.axis_names
-        ['d2', 'w1', 'w2']
-
-        Get all w1 wigners.
-
-        >>> datas = data.chop('d2', 'w1')
-        >>> len(datas)
-        51
-
-        Get 2D frequency at d2=0 fs.
+        loop through the select chop elements:
+        ```
+        is_interesting = lambda d: d.signal[:].mean() > 0.5
+        for data in filter(is_interesting, data.ichop("w1", "w2")):
+            ...
+        ```
 
-        >>> datas = data.chop('w1', 'w2', at={'d2': [0, 'fs']})
-        >>> len(datas)
-        0
-        >>> datas[0].axis_names
-        ['w1', 'w2']
-        >>> datas[0].d2[:]
-        0.
+        make a dictionary similar to a `chop` collection:
+        ```
+        chop = {f"chop{i:0>3}":d for i, d in enumerate(data.ichop(..., autoclose=False))}
+        ```
 
         See Also
         --------
-        collapse
-            Collapse the dataset along one axis.
-        split
-            Split the dataset while maintaining its dimensionality.
+        * Data.at
+            make a lower-dimension dataset at specific coordinates
+        * Data.chop
+            Divide the dataset into a collection of lower-dimensionality components.
         """
+        removed_shape, at_axes, at_idx, transform_expression = self._chop_prep(*args, at=at)
+
+        i_digits = int(np.log10(np.prod(removed_shape))) + 1
+        # iterate
+        for i, idx in enumerate(np.ndindex(removed_shape)):
+            name = f"chop{i:0>{i_digits}}"
+            idx = np.array(idx, dtype=object)
+            idx[np.array(removed_shape) == 1] = slice(None)
+            idx[at_axes] = at_idx[at_axes]
+            if autoclose:
+                with closing(self._from_slice(idx, name=name)) as a_slice:
+                    a_slice.transform(*transform_expression)
+                    yield a_slice
+            else:
+                a_slice = self._from_slice(idx, name=name)
+                a_slice.transform(*transform_expression)
+                yield a_slice
+
+    def _chop_prep(self, *args, at=None):
+        """calculates important objects for performing chop"""
         # parse args
         args = list(args)
         for i, arg in enumerate(args):
@@ -521,6 +538,66 @@ def chop(self, *args, at=None, parent=None, verbose=True) -> wt_collection.Colle
                 removed_shape[ax.shape.index(ax.size)] = 1
         removed_shape = tuple(removed_shape)
 
+        return removed_shape, at_axes, at_idx, transform_expression
+
+    def chop(self, *args, at=None, parent=None, verbose=True) -> wt_collection.Collection:
+        """Divide the dataset into its lower-dimensionality components.
+
+        Parameters
+        ----------
+        axis : str or int (args)
+            Axes of the returned data objects. Strings refer to the names of
+            axes in this object, integers refer to their index. Provide multiple
+            axes to return multidimensional data objects.
+        at : dict (optional)
+            Choice of position along an axis. Keys are axis names, values are lists
+            ``[position, input units]``. If exact position does not exist,
+            the closest valid position is used.
+        parent : WrightTools Collection instance (optional)
+            Collection to place the new "chop" collection within. Default is
+            None (new parent).
+        verbose : bool (optional)
+            Toggle talkback. Default is True.
+
+        Returns
+        -------
+        WrightTools Collection
+            Collection of chopped data objects.
+
+        Examples
+        --------
+        >>> data.axis_names
+        ['d2', 'w1', 'w2']
+
+        Get all w1 wigners.
+
+        >>> datas = data.chop('d2', 'w1')
+        >>> len(datas)
+        51
+
+        Get 2D frequency at d2=0 fs.
+
+        >>> datas = data.chop('w1', 'w2', at={'d2': [0, 'fs']})
+        >>> len(datas)
+        0
+        >>> datas[0].axis_names
+        ['w1', 'w2']
+        >>> datas[0].d2[:]
+        0.
+
+        See Also
+        --------
+        collapse
+            Collapse the dataset along one axis.
+        split
+            Split the dataset while maintaining its dimensionality.
+        ichop
+            iterator version of chop
+        at
+            select a specific smaller dimension of data at specific coordinates
+        """
+        removed_shape, at_axes, at_idx, transform_expression = self._chop_prep(*args, at=at)
+
         # get output collection
         out = wt_collection.Collection(name="chop", parent=parent)
         i_digits = int(np.log10(np.prod(removed_shape))) + 1


Test patch:
diff --git a/tests/data/ichop.py b/tests/data/ichop.py
new file mode 100644
--- /dev/null
+++ b/tests/data/ichop.py
@@ -0,0 +1,128 @@
+#! /usr/bin/env python3
+"""
+Test ichop.
+
+Many tests of `chop` also validate ichop functionality, so tests are less
+encompasing here and focus on iterator behavior.
+"""
+
+
+# --- import -------------------------------------------------------------------------------------
+
+
+import WrightTools as wt
+from WrightTools import datasets
+import pathlib
+
+
+# --- tests --------------------------------------------------------------------------------------
+
+
+def test_2D_to_1D():
+    p = datasets.PyCMDS.w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    for i, d in enumerate(data.ichop("w2")):
+        assert d.w2.size == 81
+        assert d.axis_expressions == ("w2",)
+        assert d.constant_expressions == ("w1",)
+        for k in data.variable_names:
+            assert d[k].label == data[k].label
+    assert i == data.w1.size - 1
+    data.close()
+
+
+def test_3D_to_1D():
+    p = datasets.PyCMDS.wm_w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    for i, d in enumerate(data.ichop("w2")):
+        assert d.w2.size == 11
+        assert d.axis_expressions == ("w2",)
+    assert i == (data.wm.size * data.w1.size) - 1
+    data.close()
+
+
+def test_3D_to_2D():
+    p = datasets.PyCMDS.wm_w2_w1_001
+    data = wt.data.from_PyCMDS(p)
+    for i, d in enumerate(data.ichop("wm", "w2")):
+        assert d.wm.size == 29
+        assert d.w2.size == 11
+        assert d.axis_expressions == ("wm", "w2")
+    assert i == data.w1.size - 1
+    data.close()
+
+
+def test_3D_to_2D_at():
+    p = datasets.PyCMDS.wm_w2_w1_001
+    data = wt.data.from_PyCMDS(p)
+    ichop = data.ichop("wm", "w2", at={"w1": [1590, "wn"]})
+    d = ichop.__next__()
+    assert d.wm.size == 29
+    assert d.w2.size == 11
+    assert d.axis_expressions == ("wm", "w2")
+
+    # no next because only one data object
+    try:
+        d = ichop.__next__()
+    except StopIteration:
+        pass
+    else:
+        raise AssertionError
+    data.close()
+
+
+def test_closing_context():
+    """
+    the iterator uses the close method on the last iteration,
+    so we shouldn't have persistent files
+    """
+    p = datasets.PyCMDS.wm_w2_w1_001
+    data = wt.data.from_PyCMDS(p)
+    ichop = data.ichop("wm", "w2")
+    d1 = ichop.__next__()
+    assert pathlib.Path(d1.filepath).exists()
+    d2 = ichop.__next__()
+    assert not pathlib.Path(d1.filepath).exists()
+    d2.close()
+    ichop = data.ichop("wm", "w2", autoclose=False)
+    d1 = ichop.__next__()
+    d2 = ichop.__next__()
+    assert pathlib.Path(d1.filepath).exists()
+    d1.close()
+    d2.close()
+    data.close()
+
+
+def test_filtering():
+    p = datasets.PyCMDS.wm_w2_w1_001
+    data = wt.data.from_PyCMDS(p)
+    data.print_tree()
+    mean = data.channels[0][:].mean()
+    is_interesting = lambda d: d.channels[0][:].mean() > mean
+    interesting = {
+        f"{i}": d
+        for i, d in enumerate(filter(is_interesting, data.ichop("w1", "w2", autoclose=False)))
+    }
+    assert len(interesting) == 19
+
+
+def test_collection_like():
+    p = datasets.PyCMDS.w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    collection = {f"chop{i:0>3}": d for i, d in enumerate(data.ichop("w2", autoclose=False))}
+    assert len(collection) == data.w1.size
+    collection["chop002"].__repr__()
+    wt.close()
+
+
+# --- run -----------------------------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    test_2D_to_1D()
+    test_3D_to_1D()
+    test_3D_to_2D()
+    test_3D_to_2D_at()
+    test_closing_context()
+    test_collection_like()
+    test_filtering()