Repository: wright-group/WrightTools

Issue:
moments as method of data
fit.py is being removed


Discussion:
- Example syntax: `data.moment(axis=0, channel=0, moment=1)`
- adds a channel to the data object named like we do in `data.collapse`
- moment and channel kwargs can accept tuple, e.g. `(1,3)`, to compute multiple channels/moments at once.
- `axis` and `channel` can accept strings.

Resolving patch:
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -490,9 +490,130 @@ def gradient(self, axis, *, channel=0):
         else:
             new[:] = np.gradient(channel[:], self[axis].points, axis=axis_index)
 
-    def collapse(self, axis, method="integrate"):
+    def moment(self, axis, channel=0, moment=1):
+        """Take the nth moment the dataset along one axis, adding lower rank channels.
+
+        New channels have names ``<channel name>_<axis name>_moment_<moment num>``.
+
+        Moment 0 is the integral of the slice.
+        Moment 1 is the weighted average or "Center of Mass", normalized by the integral
+        Moment 2 is the variance, the central moment about the center of mass,
+            normalized by the integral
+        Moments 3+ are central moments about the center of mass, normalized by the integral
+            and by the standard deviation to the power of the moment.
+
+        Moments, especially higher order moments, are susceptible to noise and baseline.
+        It is recommended when used with real data to use :meth:`WrightTools.data.Channel.clip`
+        in conjunction with moments to reduce effects of noise.
+
+        Parameters
+        ----------
+        axis : int or str
+            The axis to take the moment along.
+            If given as an integer, the axis with that index is used.
+            If given as a string, the axis with that name is used.
+            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 take the moment.
+            If given as an integer, the channel with that index is used.
+            If given as a string, the channel with that name is used.
+            The channel must have values along the axis
+            (i.e. its shape must not be ``1`` in the dimension for which the axis is not ``1``)
+            Default is 0, the first channel.
+        moment : int or tuple of int
+            The moments to take.
+            One channel will be created for each number given.
+            Default is 1, the center of mass.
+
+        See Also
+        --------
+        collapse
+            Reduce dimensionality by some mathematical operation
+        clip
+            Set values above/below a threshold to a particular value
         """
-        Collapse the dataset along one axis, adding lower rank channels.
+        # get axis index --------------------------------------------------------------------------
+        index = wt_kit.get_index(self.axis_names, 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]
+
+        warnings.warn("moment", category=wt_exceptions.EntireDatasetInMemoryWarning)
+
+        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 moment".format(
+                    channel, axis
+                )
+            )
+
+        new_shape = list(self[channel].shape)
+        new_shape[axis_index] = 1
+
+        channel = self[channel]
+        axis_inp = axis
+        axis = self.axes[index]
+        x = axis[:]
+        if np.any(np.isnan(x)):
+            raise wt_exceptions.ValueError("Axis '{}' includes NaN".format(axis_inp))
+        y = np.nan_to_num(channel[:])
+
+        try:
+            moments = tuple(moment)
+        except TypeError:
+            moments = (moment,)
+
+        if 0 in moments:
+            # Sort axis, so that integrals come out with expected sign
+            # only matters for integral, all others normalize by integral
+            sli = [slice(None) for _ in range(x.ndim)]
+            sort = np.argsort(x.flat)
+            sli[axis_index] = sort
+            sli = tuple(sli)
+            x = x[sli]
+            y = y[sli]
+
+        for moment in moments:
+            about = 0
+            norm = 1
+            if moment > 0:
+                norm = np.trapz(y, x, axis=axis_index)
+                norm = np.array(norm)
+                norm.shape = new_shape
+            if moment > 1:
+                about = np.trapz(x * y, x, axis=axis_index)
+                about = np.array(about)
+                about.shape = new_shape
+                about /= norm
+            if moment > 2:
+                sigma = np.trapz((x - about) ** 2 * y, x, axis=axis_index)
+                sigma = np.array(sigma)
+                sigma.shape = new_shape
+                sigma /= norm
+                sigma **= 0.5
+                norm *= sigma ** moment
+
+            values = np.trapz((x - about) ** moment * y, x, axis=axis_index)
+            values = np.array(values)
+            values.shape = new_shape
+            values /= norm
+            self.create_channel(
+                "{}_{}_{}_{}".format(channel.natural_name, axis_inp, "moment", moment),
+                values=values,
+            )
+
+    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>``.
 
@@ -517,6 +638,8 @@ def collapse(self, axis, method="integrate"):
             Divide the dataset into its lower-dimensionality components.
         split
             Split the dataset while maintaining its dimensionality.
+        moment
+            Take the moment along a particular axis
         """
         # get axis index --------------------------------------------------------------------------
         if isinstance(axis, int):


Test patch:
diff --git a/tests/data/moment.py b/tests/data/moment.py
new file mode 100755
--- /dev/null
+++ b/tests/data/moment.py
@@ -0,0 +1,48 @@
+#! /usr/bin/env python3
+import numpy as np
+import WrightTools as wt
+from WrightTools import datasets
+
+
+def test_moment_lognormal():
+    data = wt.Data()
+    data.create_variable("x", np.linspace(0, 1000, 100000))
+    data.transform("x")
+    x = data.x[:]
+    data.create_channel("y", 1 / (x * np.sqrt(2 * np.pi)) * np.exp(-np.log(x) ** 2 / 2))
+
+    data.moment(0, moment=(0, 1, 2, 3, 4))
+
+    vals = [
+        1,
+        np.exp(1 / 2),
+        (np.exp(1) - 1) * np.exp(1),
+        (np.exp(1) + 2) * np.sqrt(np.exp(1) - 1),
+        np.exp(4) + 2 * np.exp(3) + 3 * np.exp(2) - 6,
+    ]
+
+    for i, val in zip(range(5), vals):
+        assert np.isclose(data[f"y_0_moment_{i}"][0], val, rtol=10 ** (-4 + i))
+
+
+def test_moment_nd():
+    p = datasets.PyCMDS.w1_wa_000
+    data = wt.data.from_PyCMDS(p)
+    data.convert("nm")
+    data.create_variable("x", (data.wa[:] - data.w1[:]).mean(axis=0, keepdims=True), units="nm")
+
+    data.transform("w1", "x")
+
+    data.level(0, 1, -10)
+    data.moment(1, moment=0)
+    assert data.channels[-1].shape == (25, 1)
+    data.moment(1, moment=1)
+    assert data.channels[-1].shape == (25, 1)
+    data.moment(1, moment=2)
+    assert data.channels[-1].shape == (25, 1)
+    data.close()
+
+
+if __name__ == "__main__":
+    test_moment_nd()
+    test_moment_lognormal()