Repository: wright-group/WrightTools

Issue:
Make clip argument chunkwiseable
Clip fails with chunked data if you pass an array, such that you wish to clip based on different values along some given axis, something that is useful for, e.g. clipping to the value at a given index of a frame in a higher dimensional scan

Could also apply to other similar chunked functions


Discussion:
We already have this for other parts or the codebase, including `__imul__`, here:


```python
    def f(dataset, s, value):
            if hasattr(value, "shape"):
                dataset[s] *= value[wt_kit.valid_index(s, value.shape)]
            else:
                 dataset[s] *= value
```

Resolving patch:
diff --git a/WrightTools/_dataset.py b/WrightTools/_dataset.py
--- a/WrightTools/_dataset.py
+++ b/WrightTools/_dataset.py
@@ -190,16 +190,18 @@ def argmax(self):
 
             def f(dataset, s):
                 arr = dataset[s]
-                amin = np.nanargmax(arr)
+                try:
+                    amin = np.nanargmax(arr)
+                except ValueError:
+                    amin = 0
                 idx = np.unravel_index(amin, arr.shape)
-                return tuple(i + ss.indices(0)[0] for i, ss in zip(idx, s))
-
-            idxs = list(self.chunkwise(f).values())
-            curmax = idxs[0]
-            for idx in idxs[1:]:
-                if self[idx] > self[curmax]:
-                    curmax = idx
-            self.attrs["argmax"] = curmax
+                val = arr[idx]
+                return (tuple(i + (ss.start if ss.start else 0) for i, ss in zip(idx, s)), val)
+
+            chunk_res = self.chunkwise(f)
+            idxs = [i[0] for i in chunk_res.values()]
+            vals = [i[1] for i in chunk_res.values()]
+            self.attrs["argmax"] = idxs[np.nanargmax(vals)]
         return tuple(self.attrs["argmax"])
 
     def argmin(self):
@@ -208,16 +210,18 @@ def argmin(self):
 
             def f(dataset, s):
                 arr = dataset[s]
-                amin = np.nanargmin(arr)
+                try:
+                    amin = np.nanargmin(arr)
+                except ValueError:
+                    amin = 0
                 idx = np.unravel_index(amin, arr.shape)
-                return tuple(i + ss.indices(0)[0] for i, ss in zip(idx, s))
-
-            idxs = list(self.chunkwise(f).values())
-            curmin = idxs[0]
-            for idx in idxs[1:]:
-                if self[idx] < self[curmin]:
-                    curmin = idx
-            self.attrs["argmin"] = curmin
+                val = arr[idx]
+                return (tuple(i + (ss.start if ss.start else 0) for i, ss in zip(idx, s)), val)
+
+            chunk_res = self.chunkwise(f)
+            idxs = [i[0] for i in chunk_res.values()]
+            vals = [i[1] for i in chunk_res.values()]
+            self.attrs["argmin"] = idxs[np.nanargmin(vals)]
         return tuple(self.attrs["argmin"])
 
     def chunkwise(self, func, *args, **kwargs):
@@ -266,6 +270,12 @@ def clip(self, min=None, max=None, replace=np.nan):
             min = self.min()
 
         def f(dataset, s, min, max, replace):
+            if hasattr(min, "shape"):
+                min = min[wt_kit.valid_index(s, min.shape)]
+            if hasattr(max, "shape"):
+                max = max[wt_kit.valid_index(s, max.shape)]
+            if hasattr(replace, "shape"):
+                replace = replace[wt_kit.valid_index(s, replace.shape)]
             arr = dataset[s]
             if replace == "value":
                 dataset[s] = np.clip(arr, min, max)


Test patch:
diff --git a/tests/data/get_zenith.py b/tests/data/get_zenith.py
--- a/tests/data/get_zenith.py
+++ b/tests/data/get_zenith.py
@@ -11,7 +11,7 @@ def test_zenith():
     data = wt.data.from_PyCMDS(p)
     data.wm.convert("nm")
     data.w1.convert("nm")
-    assert np.allclose(list(data.get_zenith()), [499.9893, 1575., 6349.206349])
+    assert np.allclose(list(data.get_zenith()), [495.0566, 1560., 6369.426752])
 
 
 if __name__ == "__main__":
diff --git a/tests/dataset/argmax.py b/tests/dataset/argmax.py
--- a/tests/dataset/argmax.py
+++ b/tests/dataset/argmax.py
@@ -25,6 +25,6 @@ def test_PyCMDS_wm_w2_w1_000():
     assert data["wm"].argmax() == (34, 0, 0)
     assert data["w2"].argmax() == (0, 0, 0)
     assert data["w1"].argmax() == (0, 0, 10)
-    assert data.signal_diff.argmax() == (14, 5, 5)
-    assert data.signal_mean.argmax() == (14, 5, 5)
+    assert data.signal_diff.argmax() == (10, 8, 6)
+    assert data.signal_mean.argmax() == (14, 7, 4)
     data.close()
diff --git a/tests/dataset/argmin.py b/tests/dataset/argmin.py
--- a/tests/dataset/argmin.py
+++ b/tests/dataset/argmin.py
@@ -26,5 +26,5 @@ def test_PyCMDS_wm_w2_w1_000():
     assert data["w2"].argmin() == (0, 10, 0)
     assert data["w1"].argmin() == (0, 0, 0)
     assert data.signal_diff.argmin() == (0, 2, 0)
-    assert data.signal_mean.argmin() == (0, 1, 3)
+    assert data.signal_mean.argmin() == (30, 1, 9)
     data.close()