Repository: aeye-lab/pymovements
Pull request #87 (2023-01-25 15:39:39)

Add multi-channel support for synthetic.step_function
## Description

This adds multi channel support for synthetic.step_function
Fixes issue #80 

Documentation updates will be added soon.

## Implemented changes

- [x] add support for passing list of tuples as `values`
- [x] add support of scalar `start_value` in case of multi-channel `values`
- [ ] update docstring to multi-channel specification
- [ ] add at least one multi-channel example to docstring

## Type of change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [x] This change is or requires a documentation update

## How Has This Been Tested?

- [x] single step, 2 channel values, scalar start value
- [x] single step, 2 channel values, 2 channel start value
- [x] single step, 4 channel values, 4 channel start value
- [x] single step, values with varying number of channels -> `ValueError`
- [x] single step, 2 channel values, 3 channel start_value ->  `ValueError`

## Checklist:

- [x] My code follows the style guidelines of this project
- [x]  I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] Any dependent changes have been merged and published in downstream modules
- [x] I have checked my code and corrected any misspellings



Patch:
diff --git a/src/pymovements/synthetic/step_function.py b/src/pymovements/synthetic/step_function.py
index 0b0daa5..f5efe77 100644
--- a/src/pymovements/synthetic/step_function.py
+++ b/src/pymovements/synthetic/step_function.py
@@ -9,8 +9,8 @@ import numpy as np
 def step_function(
     length: int,
     steps: list[int],
-    values: list[float],
-    start_value: float = 0,
+    values: list[float | tuple[float, ...]],
+    start_value: float | tuple[float, ...] = 0,
 ) -> np.ndarray:
     """
     Create a synthetic eye gaze by using a simple step function.
@@ -47,16 +47,45 @@ def step_function(
         ... )
         array([0., 0., 1., 1., 1., 2., 2., 2., 2., 3.])
     """
+    # Check that steps and values have equal length.
     if len(steps) != len(values):
         raise ValueError(
             'length of steps not equal to length of values'
             f' ({len(steps)} != {len(values)})',
         )
 
+    # Check that steps are sorted in ascending order.
     if sorted(steps) != steps:
-        raise ValueError('steps need to be sorted in ascending order.')
+        raise ValueError('steps must be sorted in ascending order.')
 
-    arr = np.ones(length) * start_value
+    # Infer number of channels from values.
+    if isinstance(values[0], (int, float)):
+        n_channels = 1
+    else:
+        n_channels = len(values[0])
+
+    # Check that all values have equal number of channels.
+    if n_channels > 1 and any(len(value) != n_channels for value in values):
+        raise ValueError('all values must have equal number of channels.')
+
+    # Make sure start value corresponds to number of channels.
+    if n_channels > 1:
+
+        # If start value is a scalar, create tuple with length of number of channels.
+        if isinstance(start_value, (int, float)):
+            start_value = tuple(start_value for _ in range(n_channels))
+
+        # Raise error if length of start value doesn't match n_channels.
+        elif len(start_value) != n_channels:
+            raise ValueError(
+                'start_value must be scalar or must have same number of channels as values.'
+            )
+
+    # Initialize output array with start value.
+    if n_channels == 1:
+        arr = np.ones(length) * start_value
+    else:
+        arr = np.tile(start_value, (length, 1))
 
     # Iterate through all steps except the last.
     for begin, end, value in zip(steps[:-1], steps[1:], values[:-1]):


Test patch:
diff --git a/tests/synthetic/test_step_function.py b/tests/synthetic/test_step_function.py
index ae70064..4568ef5 100644
--- a/tests/synthetic/test_step_function.py
+++ b/tests/synthetic/test_step_function.py
@@ -30,6 +30,24 @@ from pymovements.synthetic import step_function
             {'value': np.concatenate([np.zeros(10), np.ones(40), np.zeros(40), np.ones(10) * 20])},
             id='length_100_3_steps',
         ),
+        pytest.param(
+            {'length': 10, 'steps': [5], 'values': [(1, 2)], 'start_value': 10},
+            {'value': np.concatenate([np.tile(10, (5, 2)), np.tile((1, 2), (5, 1))])},
+            id='length_10_2_channel_single_step_with_single_start_value',
+        ),
+        pytest.param(
+            {'length': 10, 'steps': [5], 'values': [(1, 2)], 'start_value': (11, 12)},
+            {'value': np.concatenate([np.tile((11, 12), (5, 1)), np.tile((1, 2), (5, 1))])},
+            id='length_10_2_channel_single_step_with_2_channel_start_value',
+        ),
+        pytest.param(
+            {'length': 10, 'steps': [5], 'values': [(1, 2, 3, 4)], 'start_value': (11, 12, 13, 14)},
+            {'value': np.concatenate([
+                np.tile((11, 12, 13, 14), (5, 1)),
+                np.tile((1, 2, 3, 4), (5, 1)),
+            ])},
+            id='length_10_4_channel_single_step_with_start_value',
+        ),
         pytest.param(
             {'length': 100, 'steps': [10, 50, 90], 'values': [1, 0], 'start_value': 0},
             {'exception': ValueError},
@@ -40,6 +58,16 @@ from pymovements.synthetic import step_function
             {'exception': ValueError},
             id='steps_not_sorted_raises_value_error',
         ),
+        pytest.param(
+            {'length': 10, 'steps': [3, 5], 'values': [(1, 2), (3, 5, 6)]},
+            {'exception': ValueError},
+            id='varying_number_of_channels_raises_value_error',
+        ),
+        pytest.param(
+            {'length': 10, 'steps': [5], 'values': [(1, 2)], 'start_value': (1, 2, 3)},
+            {'exception': ValueError},
+            id='number_of_channels_unequal_start_value_channels_raises_value_error',
+        ),
     ],
 )
 def test_step_function(params, expected):
@@ -51,4 +79,4 @@ def test_step_function(params, expected):
         return
 
     arr = step_function(**params)
-    assert np.array_equal(arr, expected['value']), f'arr = {arr}, expected = {expected["value"]}'
+    assert np.array_equal(arr, expected['value']), f"arr = {arr}, expected = {expected['value']}"