Repository: wright-group/WrightTools

Issue:
recover level



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
@@ -4,12 +4,10 @@
 # --- import --------------------------------------------------------------------------------------
 
 
-from __future__ import absolute_import, division, print_function, unicode_literals
-
 import collections
 import operator
 import functools
-import posixpath
+import warnings
 
 import numpy as np
 
@@ -653,73 +651,44 @@ def heal(self, channel=0, method='linear', fill_value=np.nan,
             print('channel {0} healed in {1} seconds'.format(
                 channel.name, np.around(timer.interval, decimals=3)))
 
-    def level(self, channel, axis, npts, verbose=True):
-        """For a channel, subtract the average value of several points at the edge of a given axis.
+    def level(self, channel, axis, npts, *, verbose=True):
+        """Subtract the average value of npts at the edge of a given axis.
 
         Parameters
         ----------
         channel : int or str
             Channel to level.
-        axis : int or str
+        axis : int
             Axis to level along.
         npts : int
             Number of points to average for each slice. Positive numbers
-            take leading points and negative numbers take trailing points.
+            take points at leading indicies and negative numbers take points
+            at trailing indicies.
         verbose : bool (optional)
             Toggle talkback. Default is True.
         """
-        raise NotImplementedError
-        # channel ---------------------------------------------------------------------------------
-        if isinstance(channel, int):
-            channel_index = channel
-        elif isinstance(channel, str):
-            channel_index = self.channel_names.index(channel)
-        else:
-            raise TypeError("channel: expected {int, str}, got %s" % type(channel))
-        channel = self.channels[channel_index]
-        # axis ------------------------------------------------------------------------------------
-        if isinstance(axis, int):
-            axis_index = axis
-        elif isinstance(axis, str):
-            axis_index = self.axis_names.index(axis)
-        else:
-            raise TypeError("axis: expected {int, str}, got %s" % type(axis))
-        # verify npts not zero --------------------------------------------------------------------
+        warnings.warn('level', category=wt_exceptions.EntireDatasetInMemoryWarning)
+        channel_index = wt_kit.get_index(self.channel_names, channel)
+        channel = self.channels[0]
+        # verify npts not zero
         npts = int(npts)
         if npts == 0:
-            print('cannot level if no sampling range is specified')
-            return
-        # level -----------------------------------------------------------------------------------
-        channel = self.channels[channel_index]
-        values = channel[:]
-        # transpose so the axis of interest is last
-        transpose_order = range(len(values.shape))
-        # replace axis_index with zero
-        transpose_order = [len(values.shape) - 1 if i ==
-                           axis_index else i for i in transpose_order]
-        transpose_order[len(values.shape) - 1] = axis_index
-        values = values.transpose(transpose_order)
-        # subtract
-        for index in np.ndindex(values[..., 0].shape):
-            if npts > 0:
-                offset = np.nanmean(values[index][:npts])
-            elif npts < 0:
-                offset = np.nanmean(values[index][npts:])
-            values[index] = values[index] - offset
-        # transpose back
-        values = values.transpose(transpose_order)
-        # return
-        channel[:] = values
-        channel._null = 0.
-        # print
+            raise wt_exceptions.ValueError('npts must not be zero')
+        # get subtrahend
+        ss = [slice(None)] * self.ndim
+        if npts > 0:
+            ss[axis] = slice(0, npts, None)
+        else:
+            ss[axis] = slice(-npts, None, None)
+        subtrahend = np.nanmean(channel[ss], axis=axis)
+        if self.ndim > 1:
+            subtrahend = np.expand_dims(subtrahend, axis=axis)
+        # level
+        channel[:] = channel[:] - subtrahend  # verbose
+        # finish
+        channel._null = 0
         if verbose:
-            axis = self._axes[axis_index]
-            if npts > 0:
-                points = axis[:npts]
-            if npts < 0:
-                points = axis[npts:]
-            print('channel', channel.name, 'offset by', axis.name, 'between',
-                  int(points.min()), 'and', int(points.max()), axis.units)
+            print('channel {0} leveled along axis {1}'.format(channel.natural_name, axis))
 
     def map_axis(self, axis, points, input_units='same', edge_tolerance=0., verbose=True):
         """Map points of an axis to new points using linear interpolation.
diff --git a/WrightTools/exceptions.py b/WrightTools/exceptions.py
--- a/WrightTools/exceptions.py
+++ b/WrightTools/exceptions.py
@@ -11,6 +11,7 @@
 
 # --- custom exceptions ---------------------------------------------------------------------------
 
+
 class WrightToolsException(Exception):
     """WrightTools Base Exception."""
     pass
@@ -68,6 +69,11 @@ def __init__(self, axis, operation):
         super().__init__(self, message)
 
 
+class ValueError(ValueError, WrightToolsException):
+    """Raised when an argument has the right type but an inappropriate value."""
+    pass
+
+
 class UnitsError(WrightToolsException):
     """Units Error."""
 
@@ -87,14 +93,26 @@ def __init__(self, expected, recieved):
 
 # --- custom warnings -----------------------------------------------------------------------------
 
+
 class WrightToolsWarning(Warning):
     """WrightTools Base Warning."""
     pass
 
 
+class EntireDatasetInMemoryWarning(WrightToolsWarning):
+    """Warn when an entire dataset is taken into memory at once.
+
+    Such operations may lead to memory overflow errors for large datasets.
+
+    Warning ignored by default."""
+    pass
+
+
+warnings.simplefilter("ignore", category=EntireDatasetInMemoryWarning)  # ignore by default
+
+
 class VisibleDeprecationWarning(WrightToolsWarning):
     """VisibleDepreciationWarning."""
-
     pass
 
 


Test patch:
diff --git a/tests/data/level.py b/tests/data/level.py
new file mode 100644
--- /dev/null
+++ b/tests/data/level.py
@@ -0,0 +1,47 @@
+"""Test level."""
+
+
+# --- import --------------------------------------------------------------------------------------
+
+
+import WrightTools as wt
+from WrightTools import datasets
+
+import numpy as np
+
+
+# --- test ----------------------------------------------------------------------------------------
+
+
+def test_1D():
+    p = datasets.Cary50.CuPCtS_H2O_vis
+    data = wt.collection.from_Cary(p)[0]
+    data.level(0, 0, 5)
+    assert np.isclose(data.abs[:5].mean(), 0.)
+    data.close()
+
+
+def test_2D():
+    p = datasets.COLORS.v2p2_WL_wigner
+    data = wt.data.from_COLORS(p)
+    data.level('ai0', 1, -3)
+    print(data.ai0[:, -3:].max())
+    assert np.allclose(data.ai0[:, -3:], [0.], atol=5)  # very noisy data
+    data.close()
+
+
+def test_3D():
+    p = datasets.COLORS.v2p1_MoS2_TrEE_movie
+    data = wt.data.from_COLORS(p)
+    data.level('ai0', 1, 1)
+    assert np.allclose(data.ai0[:, :1], [0.], atol=1e-3)
+    data.close()
+
+
+# --- run -----------------------------------------------------------------------------------------
+
+
+if __name__ == '__main__':
+    test_1D()
+    test_2D()
+    test_3D()