Repository: wright-group/WrightTools

Issue:
[before h5] map_axis edge tolerance
if edge points are within edge of new points within tolerance, the edges are coerced to the new point edges before interpoation


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
@@ -1275,7 +1275,7 @@ def m(self, abs_data, channel=0, this_exp='TG', indices=None, m=None, bounds_err
             else:
                 raise RuntimeError('{0} label_seed not found'.format(indi))
 
-    def map_axis(self, axis, points, input_units='same', verbose=True):
+    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.
 
         Out-of-bounds points are written nan.
@@ -1284,15 +1284,20 @@ def map_axis(self, axis, points, input_units='same', verbose=True):
         ----------
         axis : int or str
             The axis to map onto.
-        points : 1D array-like
-            The new points.
+        points : 1D array-like or int
+            If array, the new points. If int, new points will have the same
+            limits, with int defining the number of evenly spaced points
+            between.
         input_units : str (optional)
             The units of the new points. Default is same, which assumes
             the new points have the same units as the axis.
+        edge_tolerance : float (optional)
+            Axis edge points that are within this amount of the new edge
+            points are coerced to the new edge points before interpolation.
+            Default is 0.
         verbose : bool (optional)
             Toggle talkback. Default is True.
         """
-        points = np.array(points)
         # get axis index --------------------------------------------------------------------------
         if isinstance(axis, int):
             axis_index = axis
@@ -1301,6 +1306,12 @@ def map_axis(self, axis, points, input_units='same', verbose=True):
         else:
             raise TypeError("axis: expected {int, str}, got %s" % type(axis))
         axis = self.axes[axis_index]
+        # get points ------------------------------------------------------------------------------
+        if isinstance(points, int):
+            points = np.linspace(axis.points[0], axis.points[-1], points)
+            input_units = 'same'
+        else:
+            points = np.array(points)
         # transform points to axis units ----------------------------------------------------------
         if input_units == 'same':
             pass
@@ -1312,6 +1323,12 @@ def map_axis(self, axis, points, input_units='same', verbose=True):
             if self.axes[i].points[0] > self.axes[i].points[-1]:
                 self.flip(i)
                 flipped[i] = True
+        # handle edge tolerance -------------------------------------------------------------------
+        for index in [0, -1]:
+            old = axis.points[index]
+            new = points[index]
+            if new - edge_tolerance < old < new + edge_tolerance:
+                axis.points[index] = new
         # interpn data ----------------------------------------------------------------------------
         old_points = [a.points for a in self.axes]
         new_points = [a.points if a is not axis else points for a in self.axes]


Test patch:
diff --git a/tests/data/map_axis.py b/tests/data/map_axis.py
new file mode 100644
--- /dev/null
+++ b/tests/data/map_axis.py
@@ -0,0 +1,41 @@
+"""test map_axis"""
+
+
+# --- import --------------------------------------------------------------------------------------
+
+
+import numpy as np
+
+import WrightTools as wt
+from WrightTools import datasets
+
+
+# --- test ----------------------------------------------------------------------------------------
+
+
+def test_array():
+    p = datasets.JASCO.PbSe_batch_1
+    data = wt.data.from_JASCO(p)
+    assert data.shape == (1801,)
+    new = np.linspace(6000, 8000, 55)
+    data.map_axis(0, new, 'wn')
+    assert data.axes[0].points.all() == new.all()
+
+
+def test_edge_tolerance():
+    ps = datasets.KENT.LDS821_TRSF
+    data = wt.data.from_KENT(ps, ignore=['wm', 'd1', 'd2'])
+    new = np.linspace(1250, 1600, 101)
+    data.map_axis('w2', new, edge_tolerance=1)
+    assert data.w2.points.all() == new.all()
+    assert not np.isnan(data.channels[0].values).any()
+
+
+def test_int():
+    p = datasets.PyCMDS.wm_w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    assert data.shape == (35, 11, 11)
+    data.map_axis(1, 5)
+    assert data.shape == (35, 5, 11)
+    data.map_axis(2, 25)
+    assert data.shape == (35, 5, 25)