Repository: aeye-lab/pymovements
Pull request #118 (2023-02-09 08:47:37)

Trafos: Fix savitzky golay method in pos2vel()
scipy import was wrong and derivative mode was not activated


Patch:
diff --git a/src/pymovements/transforms.py b/src/pymovements/transforms.py
index cf40737..25012dd 100644
--- a/src/pymovements/transforms.py
+++ b/src/pymovements/transforms.py
@@ -6,7 +6,7 @@ from __future__ import annotations
 from typing import Any
 
 import numpy as np
-import scipy
+from scipy.signal import savgol_filter
 
 from pymovements.utils import checks
 
@@ -206,12 +206,12 @@ def pos2vel(
     elif method == 'savitzky_golay':
         # transform to velocities
         if arr.ndim == 1:
-            v = scipy.signal.savgol_filter(x=arr, **kwargs)
+            v = savgol_filter(x=arr, deriv=1, **kwargs)
         else:  # we already checked for error cases
 
             for channel_id in range(arr.shape[1]):
-                v[:, channel_id] = scipy.signal.savgol_filter(
-                    x=arr[:, channel_id], **kwargs,
+                v[:, channel_id] = savgol_filter(
+                    x=arr[:, channel_id], deriv=1, **kwargs,
                 )
         v = v * sampling_rate
 


Test patch:
diff --git a/tests/test_transformations.py b/tests/test_transformations.py
index af40d44..ca93a71 100644
--- a/tests/test_transformations.py
+++ b/tests/test_transformations.py
@@ -443,33 +443,45 @@ def test_pos2vel_raises_error(kwargs, expected_error):
     ],
 )
 def test_pos2vel_returns(method, kwargs, padding, expected_value):
-    actual_value = pos2vel(**kwargs, method=method)
+    actual_value = pos2vel(method=method, **kwargs)
     assert (actual_value[padding[0]:padding[1]] == expected_value[padding[0]:padding[1]]).all()
 
 
 @pytest.mark.parametrize(
-    'method, expected_value',
+    'params, expected_value',
     [
         pytest.param(
-            'preceding', np.array([2.0, 0.0] * (100 // 2)),
+            {'method': 'preceding', 'sampling_rate': 1},
+            np.array([2.0, 0.0] * (100 // 2)),
             id='method_preceding_alternating_velocity',
         ),
         pytest.param(
-            'neighbors', np.ones((100,)),
-            id='method_neighbors_alternating_velocity',
+            {'method': 'neighbors', 'sampling_rate': 1},
+            np.ones((100,)),
+            id='method_neighbors_linear_velocity',
         ),
         pytest.param(
-            'smooth', np.ones((100,)),
-            id='method_smooth_alternating_velocity',
+            {'method': 'smooth', 'sampling_rate': 1},
+            np.ones((100,)),
+            id='method_smooth_linear_velocity',
+        ),
+        pytest.param(
+            {'method': 'savitzky_golay', 'window_length': 7, 'polyorder': 2, 'sampling_rate': 1},
+            np.concatenate([
+                np.array([0.71428571, 0.80952381, 0.9047619]),
+                np.ones((94,)),
+                np.array([0.9047619, 0.80952381, 0.71428571]),
+            ]),
+            id='method_savitzky_golay_linear_velocity',
         ),
     ],
 )
-def test_pos2vel_stepped_input_returns(method, expected_value):
+def test_pos2vel_stepped_input_returns(params, expected_value):
     N = 100
     x = np.linspace(0, N - 2, N // 2)
     x = np.repeat(x, 2)
 
-    actual_value = pos2vel(x, sampling_rate=1, method=method)
+    actual_value = pos2vel(x, **params)
 
     lpad, rpad = 1, -1
-    assert (actual_value[lpad:rpad] == expected_value[lpad:rpad]).all()
+    assert np.allclose(actual_value[lpad:rpad], expected_value[lpad:rpad])