Repository: wright-group/WrightTools

Issue:
data.gradient
[numpy.gradient](https://docs.scipy.org/doc/numpy/reference/generated/numpy.gradient.html) is a fantastic little method

It would be really nice to have a method `data.gradient` that wraps the numpy method for our pretty little channel objects.


Discussion:

Resolving patch:
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -432,11 +432,67 @@ def chop(self, *args, at={}, parent=None, verbose=True) -> wt_collection.Collect
             print("chopped data into %d piece(s)" % len(out), "in", new_axes)
         return out
 
+    def gradient(self, axis, *, channel=0):
+        """
+        Compute the gradient along one axis.
+
+        New channels have names ``<channel name>_<axis name>_gradient``.
+
+        Parameters
+        ----------
+        axis : int or str
+            The axis to differentiate along.
+            If given as an integer, the axis in the underlying array is used,
+            and unitary spacing is assumed.
+            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.
+        channel : int or str
+            The channel to differentiate.
+            Default is the first channel.
+        """
+        # get axis index --------------------------------------------------------------------------
+        if isinstance(axis, int):
+            axis_index = axis
+        elif isinstance(axis, str):
+            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 compute gradient".format(axis)
+                )
+            axis_index = axes[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]
+
+        if self[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)
+        new = self.create_channel(
+            "{}_{}_gradient".format(channel, axis),
+            values=np.empty(self[channel].shape, dtype=rtype),
+        )
+
+        channel = self[channel]
+        if axis == axis_index:
+            new[:] = np.gradient(channel[:], axis=axis_index)
+        else:
+            new[:] = np.gradient(channel[:], self[axis].points, axis=axis_index)
+
     def collapse(self, axis, method="integrate"):
         """
         Collapse the dataset along one axis, adding lower rank channels.
 
-        New channels have names <channel name>_<axis name>_<method>.
+        New channels have names ``<channel name>_<axis name>_<method>``.
 
         Parameters
         ----------
diff --git a/examples/gradient.py b/examples/gradient.py
new file mode 100644
--- /dev/null
+++ b/examples/gradient.py
@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+"""
+Gradient
+========
+
+Demonstration of the gradient method.
+"""
+
+import numpy as np
+import WrightTools as wt
+
+data = wt.data.Data()
+data.create_variable("w1", np.linspace(-10, 10, 100))
+data.create_channel("sig", 1 / (np.pi * (1 + (data.w1[:] - 1) ** 2)))
+data.transform("w1")
+data.gradient("w1")
+
+wt.artists.quick1D(data)
+wt.artists.quick1D(data, channel="sig_w1_gradient")


Test patch:
diff --git a/tests/data/gradient.py b/tests/data/gradient.py
new file mode 100755
--- /dev/null
+++ b/tests/data/gradient.py
@@ -0,0 +1,35 @@
+#! /usr/bin/env python3
+"""Test collapse."""
+
+
+# --- import -------------------------------------------------------------------------------------
+
+
+import WrightTools as wt
+import numpy as np
+
+
+# --- tests --------------------------------------------------------------------------------------
+
+
+def test_gradient():
+    data = wt.Data()
+    data.create_variable("v1", np.arange(0, 6))
+    data.create_variable("v2", np.arange(0, 6) * 2)
+    data.create_channel("ch", np.array([1, 2, 4, 7, 11, 16]))
+    data.transform("v1")
+    data.gradient("v1")
+    data.transform("v2")
+    data.gradient("v2")
+
+    assert data.ch_v1_gradient.shape == (6,)
+    assert np.allclose(data.ch_v1_gradient.points, np.array([1., 1.5, 2.5, 3.5, 4.5, 5.]))
+    assert data.ch_v2_gradient.shape == (6,)
+    assert np.allclose(data.ch_v2_gradient.points, np.array([0.5, 0.75, 1.25, 1.75, 2.25, 2.5]))
+
+
+# --- run -----------------------------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    test_gradient()