Repository: wright-group/WrightTools

Issue:
recover collapse
collapse will create NEW channels by default---overwriting can be a kwarg (perhaps at a later date)

three levels of behavior for specifying the axis:

1. give an integer, simply collapse along that dimension (easy, first to implement)
2. give an expression corresponding to a "1D" axis, collapse along the non-trivial dimension of that axis (also easy)
3. give an expression corresponding to a n-dimensional axis (difficult, not implemented at first)

I move that we spawn separate issues for overwrite-existing-channel and collapse-n-dimensions and simply cover the easy ground now


Discussion:
Moved to 3.1.0 because it is desired for the Cobalamin paper

Resolving patch:
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -353,16 +353,24 @@ def chop(self, *args, at={}, parent=None, verbose=True):
 
     def collapse(self, axis, method="integrate"):
         """
-        Collapse the dataset along one axis.
+        Collapse the dataset along one axis, adding lower rank channels.
+
+        New channels have names <channel name>_<axis name>_<method>.
 
         Parameters
         ----------
         axis : int or str
             The axis to collapse along.
+            If given as an integer, the axis in the underlying array is used.
+            If given as a string, the axis must exist, and be a 1D array-aligned axis.
+            (i.e. have a shape with a single value which is not ``1``)
+            The axis to collapse along is inferred from the shape of the axis.
         method : {'integrate', 'average', 'sum', 'max', 'min'} (optional)
             The method of collapsing the given axis. Method may also be list
             of methods corresponding to the channels of the object. Default
             is integrate. All methods but integrate disregard NANs.
+            Can also be a list, allowing for different treatment for varied channels.
+            In this case, None indicates that no change to that channel should occur.
 
         See Also
         --------
@@ -371,38 +379,92 @@ def collapse(self, axis, method="integrate"):
         split
             Split the dataset while maintaining its dimensionality.
         """
-        raise NotImplementedError
         # get axis index --------------------------------------------------------------------------
         if isinstance(axis, int):
             axis_index = axis
         elif isinstance(axis, str):
-            axis_index = self.axis_names.index(axis)
+            index = self.axis_names.index(axis)
+            axes = [i for i in range(self.ndim) if self.axes[index].shape[i] > 1]
+            if len(axes) > 1:
+                raise wt_exceptions.MultidimensionalAxisError(axis, "collapse")
+            elif len(axes) == 0:
+                raise wt_exceptions.ValueError(
+                    "Axis {} is a single point, cannot collapse".format(axis)
+                )
+            axis_index = axes[0]
         else:
-            raise TypeError("axis: expected {int, str}, got %s" % type(axis))
+            raise wt_exceptions.TypeError("axis: expected {int, str}, got %s" % type(axis))
+
+        new_shape = list(self.shape)
+        new_shape[axis_index] = 1
         # methods ---------------------------------------------------------------------------------
         if isinstance(method, list):
             if len(method) == len(self.channels):
                 methods = method
             else:
-                print("method argument incompatible in data.collapse")
+                raise wt_exceptions.ValueError(
+                    "method argument must have same number of elements as there are channels"
+                )
+            for m in methods:
+                if m not in [
+                    "sum",
+                    "max",
+                    "maximum",
+                    "min",
+                    "minimum",
+                    "ave",
+                    "average",
+                    "mean",
+                    "int",
+                    "integrate",
+                ]:
+                    raise wt_exceptions.ValueError("method '{}' not recognized".format(m))
         elif isinstance(method, str):
             methods = [method for _ in self.channels]
+
+        warnings.warn("collapse", category=wt_exceptions.EntireDatasetInMemoryWarning)
+
         # collapse --------------------------------------------------------------------------------
-        for method, channel in zip(methods, self.channels):
-            if method in ["int", "integrate"]:
-                channel[:] = np.trapz(y=channel[:], x=self._axes[axis_index][:], axis=axis_index)
-            elif method == "sum":
-                channel[:] = np.nansum(channel[:], axis=axis_index)
+        for method, channel in zip(methods, self.channel_names):
+            if method is None:
+                continue
+
+            if self[channel].shape[axis_index] == 1:
+                continue  # Cannot collapse any further, don't clutter data object
+
+            new_shape = list(self[channel].shape)
+            new_shape[axis_index] = 1
+
+            new = self.create_channel(
+                "{}_{}_{}".format(channel, axis, method),
+                shape=new_shape,
+                units=self[channel].units,
+            )
+
+            channel = self[channel]
+
+            if method == "sum":
+                res = np.nansum(channel[:], axis=axis_index)
+                res.shape = new_shape
+                new[:] = res
             elif method in ["max", "maximum"]:
-                channel[:] = np.nanmax(channel[:], axis=axis_index)
+                res = np.nanmax(channel[:], axis=axis_index)
+                res.shape = new_shape
+                new[:] = res
             elif method in ["min", "minimum"]:
-                channel[:] = np.nanmin(channel[:], axis=axis_index)
+                res = np.nanmin(channel[:], axis=axis_index)
+                res.shape = new_shape
+                new[:] = res
             elif method in ["ave", "average", "mean"]:
-                channel[:] = np.nanmean(channel[:], axis=axis_index)
+                res = np.nanmean(channel[:], axis=axis_index)
+                res.shape = new_shape
+                new[:] = res
+            elif method in ["int", "integrate"]:
+                res = np.trapz(y=channel[:], x=self._axes[axis_index][:], axis=axis_index)
+                res.shape = new_shape
+                new[:] = res
             else:
-                print("method not recognized in data.collapse")
-        # cleanup ---------------------------------------------------------------------------------
-        self._axes.pop(axis_index)
+                raise wt_exceptions.ValueError("method '{}' not recognized".format(m))
 
     def convert(self, destination_units, *, convert_variables=False, verbose=True):
         """Convert all compatable axes to given units.
@@ -473,7 +535,7 @@ def create_channel(self, name, values=None, shape=None, units=None, **kwargs):
             Created channel.
         """
         if name in self.channel_names:
-            wt_exceptions.ObjectExistsWarning.warn(name)
+            warnings.warn(name, wt_exceptions.ObjectExistsWarning)
             return self[name]
 
         require_kwargs = {}
@@ -519,7 +581,7 @@ def create_variable(self, name, values=None, shape=None, units=None, **kwargs):
             New child variable.
         """
         if name in self.variable_names:
-            wt_exceptions.ObjectExistsWarning.warn(name)
+            warnings.warn(name, wt_exceptions.ObjectExistsWarning)
             return self[name]
         if values is None:
             if shape is None:
diff --git a/WrightTools/exceptions.py b/WrightTools/exceptions.py
--- a/WrightTools/exceptions.py
+++ b/WrightTools/exceptions.py
@@ -76,6 +76,15 @@ class ValueError(ValueError, WrightToolsException):
     pass
 
 
+class TypeError(TypeError, WrightToolsException):
+    """Raised when an operation or function is applied to an object of inappropriate type.
+
+    The associated value is a string giving details about the type mismatch.
+    """
+
+    pass
+
+
 class UnitsError(WrightToolsException):
     """Units Error."""
 


Test patch:
diff --git a/tests/data/collapse.py b/tests/data/collapse.py
new file mode 100755
--- /dev/null
+++ b/tests/data/collapse.py
@@ -0,0 +1,133 @@
+#! /usr/bin/env python3
+"""Test collapse."""
+
+
+# --- import -------------------------------------------------------------------------------------
+
+
+import WrightTools as wt
+import numpy as np
+
+
+# --- tests --------------------------------------------------------------------------------------
+
+
+def test_integrate():
+    data = wt.Data()
+    data.create_variable("v1", np.arange(0, 10)[None, :])
+    data.create_variable("v2", np.arange(0, 10)[:, None])
+    data.create_channel("ch", data.v1[:] * data.v2[:])
+    data.transform("v2", "v1")
+    data.collapse("v1", "integrate")
+    data.collapse("v2", "int")
+
+    assert data.ch_v2_int.shape == (1, 10)
+    assert data.ch_v1_integrate.shape == (10, 1)
+    assert data.ch_v1_integrate_v2_int.shape == (1, 1)
+    assert np.allclose(
+        data.ch_v2_int.points,
+        np.array([0., 40.5, 81., 121.5, 162., 202.5, 243., 283.5, 324., 364.5]),
+    )
+    assert np.allclose(
+        data.ch_v1_integrate.points,
+        np.array([0., 40.5, 81., 121.5, 162., 202.5, 243., 283.5, 324., 364.5]),
+    )
+    assert np.allclose(data.ch_v1_integrate_v2_int.points, np.array(1640.25))
+    assert len(data.channel_names) == 4
+
+
+def test_average():
+    data = wt.Data()
+    data.create_variable("v1", np.arange(0, 10)[None, :])
+    data.create_variable("v2", np.arange(0, 10)[:, None])
+    data.create_channel("ch", data.v1[:] * data.v2[:])
+    data.transform("v2", "v1")
+    data.collapse(0, "average")
+    data.collapse(1, "mean")
+
+    assert data.ch_1_mean.shape == (10, 1)
+    assert data.ch_0_average.shape == (1, 10)
+    assert data.ch_0_average_1_mean.shape == (1, 1)
+    assert np.allclose(
+        data.ch_1_mean.points, np.array([0., 4.5, 9., 13.5, 18., 22.5, 27., 31.5, 36., 40.5])
+    )
+    assert np.allclose(
+        data.ch_0_average.points, np.array([0., 4.5, 9., 13.5, 18., 22.5, 27., 31.5, 36., 40.5])
+    )
+    assert np.allclose(data.ch_0_average_1_mean.points, np.array(20.25))
+    assert len(data.channel_names) == 4
+
+
+def test_sum():
+    data = wt.Data()
+    data.create_variable("v1", np.arange(0, 10)[None, :])
+    data.create_variable("v2", np.arange(0, 10)[:, None])
+    data.create_channel("ch", data.v1[:] * data.v2[:])
+    data.transform("v2", "v1")
+    data.collapse("v1", "sum")
+    data.collapse("v2", "sum")
+
+    assert data.ch_v2_sum.shape == (1, 10)
+    assert data.ch_v1_sum.shape == (10, 1)
+    assert data.ch_v1_sum_v2_sum.shape == (1, 1)
+    assert np.allclose(
+        data.ch_v2_sum.points, np.array([0., 45., 90., 135., 180., 225., 270., 315., 360., 405.])
+    )
+    assert np.allclose(
+        data.ch_v1_sum.points, np.array([0., 45., 90., 135., 180., 225., 270., 315., 360., 405.])
+    )
+    assert np.allclose(data.ch_v1_sum_v2_sum.points, np.array(2025))
+    assert len(data.channel_names) == 4
+
+
+def test_max():
+    data = wt.Data()
+    data.create_variable("v1", np.arange(0, 10)[None, :])
+    data.create_variable("v2", np.arange(0, 10)[:, None])
+    data.create_channel("ch", data.v1[:] * data.v2[:])
+    data.transform("v2", "v1")
+    data.collapse("v1", "maximum")
+    data.collapse("v2", "max")
+
+    assert data.ch_v2_max.shape == (1, 10)
+    assert data.ch_v1_maximum.shape == (10, 1)
+    assert data.ch_v1_maximum_v2_max.shape == (1, 1)
+    assert np.allclose(
+        data.ch_v2_max.points, np.array([0., 9., 18., 27., 36., 45., 54., 63., 72., 81.])
+    )
+    assert np.allclose(
+        data.ch_v1_maximum.points, np.array([0., 9., 18., 27., 36., 45., 54., 63., 72., 81.])
+    )
+    assert np.allclose(data.ch_v1_maximum_v2_max.points, np.array(81))
+    assert len(data.channel_names) == 4
+
+
+def test_min():
+    data = wt.Data()
+    data.create_variable("v1", np.arange(0, 10)[None, :])
+    data.create_variable("v2", np.arange(0, 10)[:, None])
+    data.create_channel("ch", data.v1[:] * data.v2[:])
+    data.transform("v2", "v1")
+    data.collapse("v1", "minimum")
+    data.collapse("v2", "min")
+
+    assert data.ch_v2_min.shape == (1, 10)
+    assert data.ch_v1_minimum.shape == (10, 1)
+    assert data.ch_v1_minimum_v2_min.shape == (1, 1)
+    assert np.allclose(data.ch_v2_min.points, np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]))
+    assert np.allclose(
+        data.ch_v1_minimum.points, np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
+    )
+    assert np.allclose(data.ch_v1_minimum_v2_min.points, np.array(0))
+    assert len(data.channel_names) == 4
+
+
+# --- run -----------------------------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    test_integrate()
+    test_average()
+    test_sum()
+    test_max()
+    test_min()