Merge pull request #746 from wright-group/smooth1D

revamp kit.smooth_1D (resolves #376)

WrightTools/kit/_array.py

@@ -243,18 +243,52 @@ def share_nans(*arrs) -> tuple:
     return tuple([a + nans for a in arrs])
 
 
-def smooth_1D(arr, n=10) -> np.ndarray:
-    """Smooth 1D data by 'running average'.
+def smooth_1D(arr, n=10, smooth_type="flat") -> np.ndarray:
+    """Smooth 1D data using a window function.
+    
+    Edge effects will be present. 
 
     Parameters
     ----------
-    n : int
-        number of points to average
+    arr : array_like
+        Input array, 1D.
+    n : int (optional)
+        Window length.
+    smooth_type : {'flat', 'hanning', 'hamming', 'bartlett', 'blackman'} (optional)
+        Type of window function to convolve data with.
+        'flat' window will produce a moving average smoothing.
+        
+    Returns
+    -------
+    array_like
+        Smoothed 1D array.
     """
-    for i in range(n, len(arr) - n):
-        window = arr[i - n : i + n].copy()
-        arr[i] = window.mean()
-    return arr
+
+    # check array input
+    if arr.ndim != 1:
+        raise wt_exceptions.DimensionalityError(1, arr.ndim)
+    if arr.size < n:
+        message = "Input array size must be larger than window size."
+        raise wt_exceptions.ValueError(message)
+    if n < 3:
+        return arr
+    # construct window array
+    if smooth_type == "flat":
+        w = np.ones(n, dtype=arr.dtype)
+    elif smooth_type == "hanning":
+        w = np.hanning(n)
+    elif smooth_type == "hamming":
+        w = np.hamming(n)
+    elif smooth_type == "bartlett":
+        w = np.bartlett(n)
+    elif smooth_type == "blackman":
+        w = np.blackman(n)
+    else:
+        message = "Given smooth_type, {0}, not available.".format(str(smooth_type))
+        raise wt_exceptions.ValueError(message)
+    # convolve reflected array with window function
+    out = np.convolve(w / w.sum(), arr, mode="same")
+    return out
 
 
 def svd(a, i=None) -> tuple:

tests/kit/smooth_1D.py

@@ -0,0 +1,35 @@
+"""Test kit.smooth_1D."""
+
+
+# --- import --------------------------------------------------------------------------------------
+
+
+import numpy as np
+
+import WrightTools as wt
+
+
+# --- test ----------------------------------------------------------------------------------------
+
+
+def test_basic_smoothing_functionality():
+    # create arrays
+    x = np.linspace(0, 10, 1000)
+    y = np.sin(x)
+    np.random.seed(seed=12)
+    r = np.random.rand(1000) - .5
+    yr = y + r
+    # iterate through window types
+    windows = ["flat", "hanning", "hamming", "bartlett", "blackman"]
+    for w in windows:
+        out = wt.kit.smooth_1D(yr, n=101, smooth_type=w)
+        check_arr = out - y
+        check_arr = check_arr[50:-50]  # get rid of edge effects
+        assert np.allclose(check_arr, 0, rtol=.2, atol=.2)
+
+
+# --- run -----------------------------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    test_basic_smoothing_functionality()