Repository: wright-group/WrightTools

Issue:
Data.chop: add "drop" kwarg to bypass MultidimensionalAxisError
Looking for easy method for chopping to remove multidimensional array dims we do not care about.
This is based on the issue of accidentally recording superfluous data on an array detector, when all we care about is the non-array data and axes.

e.g. we have interesting channels/variables of shape (n,m,1), but uninteresting channels/variables of shape (n,m,p) that we wish to dispose of.  

@ksunden had a good proposal:
> perhaps we could add a "drop" kwarg to chop which would make this
>```
>data.transform("w1_points", "d1_points")
>data = data.chop("w1_points", "d1_points", drop=("wavelengths", "intensities"))
> ```
> possibly could even be
> ```
> data.transform("w1_points", "d1_points")
> data = data.chop("w1_points", "d1_points", drop=True)
> ```
> where drop=True means "drop all vars and channels which have non-one in the shape where the chopped axes have one in their shape" (perhaps could be worded more succinctly, but that is the idea)

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index cd1d095..e43f1b1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 
 ## [Unreleased]
 
+### Added
+- `Data.squeeze`: squeezing the data object to the shape of the axes.
+
 ### Fixed
 - `interact2D`: fixed bug where use_imshow broke the sliders
 
diff --git a/WrightTools/data/_channel.py b/WrightTools/data/_channel.py
index c495c48..8a12b2b 100644
--- a/WrightTools/data/_channel.py
+++ b/WrightTools/data/_channel.py
@@ -225,3 +225,13 @@ class Channel(Dataset):
         if verbose:
             print("%i outliers removed" % len(outliers))
         return outliers
+
+    def _to_dict(self):
+        out = {}
+        out["name"] = self.natural_name
+        out["values"] = self[:]
+        out["units"] = self.units
+        out["label"] = self.label
+        out["signed"] = self.signed
+        out.update(self.attrs)
+        return out
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
index 86d496d..0750712 100644
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -360,6 +360,78 @@ class Data(Group):
         idx = self._at_to_slice(**at)
         return self._from_slice(idx, name=name, parent=parent)
 
+    def squeeze(self, name=None, parent=None):
+        """Reduce the data to the dimensionality of the (non-trivial) span of the axes.
+        i.e. if the joint shape of the axes has an array dimension with length 1, this
+        array dimension is squeezed.
+
+        channels and variables that span beyond the axes are removed.
+
+        Parameters
+        ----------
+        name : string (optional)
+            name of the new Data.
+        parent : WrightTools Collection instance (optional)
+            Collection to place the new "chop" collection within. Default is
+            None (new parent).
+
+        Returns
+        -------
+        out : wt.Data
+            new data object.  The new data object has dimensions of the
+            (non-trivial) span of the current axes
+
+        Examples
+        --------
+        >>> ...
+
+        See also
+        --------
+        Data.chop: Divide the dataset into its lower-dimensionality components.
+        ...
+        """
+        new = Data(name=name, parent=parent)
+
+        attrs = {
+            k: v
+            for k, v in self.attrs.items()
+            if k
+            not in [
+                "axes",
+                "channel_names",
+                "constants",
+                "name",
+                "source",
+                "item_names",
+                "variable_names",
+            ]
+        }
+        new.attrs.update(attrs)
+
+        joint_shape = wt_kit.joint_shape(*[ai[:] for ai in self.axes])
+        cull_dims = [j == 1 for j in joint_shape]
+        sl = [0 if cull else slice(None) for cull in cull_dims]
+        matches_broadcast_axes = lambda a: all(
+            [a.shape[i] == 1 for i in range(self.ndim) if cull_dims[i]]
+        )
+
+        for v in filter(matches_broadcast_axes, self.variables):
+            kwargs = v._to_dict()
+            kwargs["values"] = v[sl]
+            new.create_variable(**kwargs)
+
+        for c in filter(matches_broadcast_axes, self.channels):
+            kwargs = c._to_dict()
+            kwargs["values"] = c[sl]
+            new.create_channel(**kwargs)
+
+        # inherit constants
+        for c in self.constants:
+            new.create_constant(c.expression)
+
+        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.
 
@@ -509,21 +581,12 @@ class Data(Group):
             out = parent.create_data(name=name)
 
         for v in self.variables:
-            kwargs = {}
-            kwargs["name"] = v.natural_name
+            kwargs = v._to_dict()
             kwargs["values"] = v[idx]
-            kwargs["units"] = v.units
-            kwargs["label"] = v.label
-            kwargs.update(v.attrs)
             out.create_variable(**kwargs)
         for c in self.channels:
-            kwargs = {}
-            kwargs["name"] = c.natural_name
+            kwargs = c._to_dict()
             kwargs["values"] = c[idx]
-            kwargs["units"] = c.units
-            kwargs["label"] = c.label
-            kwargs["signed"] = c.signed
-            kwargs.update(c.attrs)
             out.create_channel(**kwargs)
 
         new_axes = [a.expression for a in self.axes if a[idx].size > 1]
diff --git a/WrightTools/data/_variable.py b/WrightTools/data/_variable.py
index f5ebe5d..2d139a5 100644
--- a/WrightTools/data/_variable.py
+++ b/WrightTools/data/_variable.py
@@ -50,3 +50,12 @@ class Variable(Dataset):
     @label.setter
     def label(self, label):
         self.attrs["label"] = label
+
+    def _to_dict(self):
+        out = {}
+        out["name"] = self.natural_name
+        out["values"] = self[:]
+        out["units"] = self.units
+        out["label"] = self.label
+        out.update(self.attrs)
+        return out


Test patch:
diff --git a/tests/data/squeeze.py b/tests/data/squeeze.py
new file mode 100644
index 0000000..518d478
--- /dev/null
+++ b/tests/data/squeeze.py
@@ -0,0 +1,46 @@
+#! /usr/bin/env python3
+"""Test squeeze."""
+
+
+# --- import -------------------------------------------------------------------------------------
+
+
+import numpy as np
+import WrightTools as wt
+from WrightTools import datasets
+
+
+# --- tests --------------------------------------------------------------------------------------
+
+
+def test_squeeze():
+    d = wt.Data(name="test")
+    d.create_variable("x", values=np.arange(5)[:, None, None])
+    d.create_variable("y", values=np.arange(4)[None, :, None])
+    d.create_variable("redundant_array", values=np.tile(np.arange(3), (5, 4, 1)))
+
+    d.create_channel("keep", values=d.x[:] + d.y[:])
+    d.create_channel("throw_away", values=np.zeros((5, 4, 3)))
+
+    d.transform("x", "y")
+    d = d.squeeze()  # make sure it runs error free
+    assert d.ndim == 2
+    assert d.shape == (5, 4)
+
+
+def test_constants():
+    d = wt.Data(name="test")
+    d.create_variable("x", values=np.array([1]).reshape(1, 1))
+    d.create_constant("x")
+    d.create_variable("y", values=np.linspace(3, 5, 4).reshape(-1, 1))
+    d.create_variable("z", values=np.linspace(0, 1, 6).reshape(1, -1))
+    d.transform("y")
+    ds = d.squeeze()
+    assert "x" in ds.constant_expressions
+    d.print_tree()
+    ds.print_tree()
+
+
+if __name__ == "__main__":
+    test_squeeze()
+    test_constants()