Repository: aeye-lab/pymovements
Pull request #670 (2024-01-30 16:18:03)

feat: add trial_columns to EventDataFrame
## Description

For trializing the event detection method it is nice to have info on the trial columns available in the EventDataFrame class itself.

## Implemented changes

- [x] add `trials` argument for passing a list of trial identifiers
- [x] add `trial_columns` for passing a list of columns that refer to trial identifiers
- [x] enforce trial columns to be at the start of the dataframe

## Type of change

- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

- [x] test for correct trial column names
- [x] test for correct trial identifier data

## 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
- [x] 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


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/aeye-lab/pymovements/pull/670?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`250c0f5`)](https://app.codecov.io/gh/aeye-lab/pymovements/commit/250c0f5fef5f4ec3aeebba86ff52076ad136c034?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) 100.00% compared to head [(`49c5aad`)](https://app.codecov.io/gh/aeye-lab/pymovements/pull/670?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) 100.00%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff            @@
##              main      #670   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           54        54           
  Lines         2357      2359    +2     
  Branches       588       588           
=========================================
+ Hits          2357      2359    +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/aeye-lab/pymovements/pull/670?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).


Patch:
diff --git a/src/pymovements/events/frame.py b/src/pymovements/events/frame.py
index 1ccef54..d6c34ae 100644
--- a/src/pymovements/events/frame.py
+++ b/src/pymovements/events/frame.py
@@ -46,6 +46,10 @@ class EventDataFrame:
         List of onsets. (default: None)
     offsets: list[int] | np.ndarray | None
         List of offsets. (default: None)
+    trials: list[int | float | str] | np.ndarray | None
+        List of trial identifiers. (default: None)
+    trial_columns: list[str] | str | None
+        List of trial columns in passed dataframe.
 
     Raises
     ------
@@ -62,16 +66,26 @@ class EventDataFrame:
             name: str | list[str] | None = None,
             onsets: list[int] | np.ndarray | None = None,
             offsets: list[int] | np.ndarray | None = None,
+            trials: list[int | float | str] | np.ndarray | None = None,
+            trial_columns: list[str] | str | None = None,
     ):
+        self.trial_columns: list[str] | None  # otherwise mypy gets confused.
+
         if data is not None:
             checks.check_is_mutual_exclusive(data=data, onsets=onsets)
             checks.check_is_mutual_exclusive(data=data, offsets=offsets)
             checks.check_is_mutual_exclusive(data=data, name=name)
+            checks.check_is_mutual_exclusive(data=data, name=trials)
 
             data = data.clone()
             data = self._add_minimal_schema_columns(data)
             data_dict = data.to_dict()
 
+            if isinstance(trial_columns, str):
+                self.trial_columns = [trial_columns]
+            else:
+                self.trial_columns = trial_columns
+
             self._additional_columns = [
                 column_name for column_name in data_dict.keys()
                 if column_name not in self._minimal_schema
@@ -105,14 +119,27 @@ class EventDataFrame:
                     'onset': pl.Series(onsets, dtype=pl.Int64),
                     'offset': pl.Series(offsets, dtype=pl.Int64),
                 }
+
+                if trials is not None:
+                    data_dict['trial'] = pl.Series('trial', trials)
+                    self.trial_columns = ['trial']
+                else:
+                    self.trial_columns = None
+
             else:
                 data_dict = {
                     'name': pl.Series([], dtype=pl.Utf8),
                     'onset': pl.Series([], dtype=pl.Int64),
                     'offset': pl.Series([], dtype=pl.Int64),
                 }
+                self.trial_columns = None
 
         self.frame = pl.DataFrame(data=data_dict, schema_overrides=self._minimal_schema)
+
+        # Ensure column order: trial columns, name, onset, offset.
+        if self.trial_columns is not None:
+            self.frame = self.frame.select([*self.trial_columns, *self._minimal_schema.keys()])
+
         if 'duration' not in self.frame.columns:
             self._add_duration_property()
 


Test patch:
diff --git a/tests/unit/events/frame_test.py b/tests/unit/events/frame_test.py
index 03624e0..8b49340 100644
--- a/tests/unit/events/frame_test.py
+++ b/tests/unit/events/frame_test.py
@@ -181,6 +181,217 @@ def test_event_dataframe_init_expected(args, kwargs, expected_df_data, expected_
     assert_frame_equal(event_df.frame, expected_df)
 
 
+@pytest.mark.parametrize(
+    ('args', 'kwargs', 'expected_df'),
+    [
+        pytest.param(
+            [], {'onsets': [0], 'offsets': [1]},
+            pl.DataFrame({'name': [''], 'onset': [0], 'offset': [1], 'duration': [1]}),
+            id='no_arg_lists_with_single_event_kwarg',
+        ),
+        pytest.param(
+            [pl.DataFrame()], {},
+            pl.DataFrame(
+                {}, schema={
+                    'name': pl.Utf8, 'onset': pl.Int64, 'offset': pl.Int64, 'duration': pl.Int64,
+                },
+            ),
+            id='dataframe_arg_no_kwargs',
+        ),
+        pytest.param(
+            [], {'name': 'bar', 'onsets': [0], 'offsets': [1]},
+            pl.DataFrame({'name': ['bar'], 'onset': [0], 'offset': [1], 'duration': [1]}),
+            id='lists_with_single_named_event',
+        ),
+        pytest.param(
+            [], {'name': 'bar', 'onsets': [0, 2], 'offsets': [1, 3]},
+            pl.DataFrame(
+                {'name': ['bar', 'bar'], 'onset': [0, 2], 'offset': [1, 3], 'duration': [1, 1]},
+            ),
+            id='lists_with_two_events_same_name',
+        ),
+        pytest.param(
+            [], {'name': ['foo', 'bar'], 'onsets': [0, 2], 'offsets': [1, 4]},
+            pl.DataFrame(
+                {'name': ['foo', 'bar'], 'onset': [0, 2], 'offset': [1, 4], 'duration': [1, 2]},
+            ),
+            id='lists_with_two_differently_named_events',
+        ),
+        pytest.param(
+            [], {'name': ['foo'], 'onsets': [0], 'offsets': [1], 'trials': [1]},
+            pl.DataFrame(
+                {'trial': [1], 'name': ['foo'], 'onset': [0], 'offset': [1], 'duration': [1]},
+            ),
+            id='lists_one_event_trial_column_at_start',
+        ),
+        pytest.param(
+            [], {
+                'data': pl.DataFrame(
+                    data={
+                        'trial': [1], 'name': ['foo'], 'onset': [0], 'offset': [1],
+                    },
+                ),
+                'trial_columns': 'trial',
+            },
+            pl.DataFrame(
+                {'trial': [1], 'name': ['foo'], 'onset': [0], 'offset': [1], 'duration': [1]},
+            ),
+            id='data_one_event_trial_column_at_start',
+        ),
+        pytest.param(
+            [], {
+                'data': pl.DataFrame(
+                    data={
+                        'name': ['foo'], 'onset': [0], 'offset': [1], 'trial': [1],
+                    },
+                ),
+                'trial_columns': 'trial',
+            },
+            pl.DataFrame(
+                {'trial': [1], 'name': ['foo'], 'onset': [0], 'offset': [1], 'duration': [1]},
+            ),
+            id='data_one_event_trial_column_enforce_start',
+        ),
+    ],
+)
+def test_event_dataframe_init_expected_df(args, kwargs, expected_df):
+    event_df = pm.EventDataFrame(*args, **kwargs)
+
+    assert_frame_equal(event_df.frame, expected_df)
+
+
+@pytest.mark.parametrize(
+    ('kwargs', 'expected_trial_column_list'),
+    [
+        pytest.param(
+            {'data': pl.DataFrame()},
+            None,
+            id='empty_df_no_trial_columns',
+        ),
+        pytest.param(
+            {'onsets': [0], 'offsets': [1]},
+            None,
+            id='single_row_no_trial_columns',
+        ),
+        pytest.param(
+            {'onsets': [0], 'offsets': [1], 'trials': None},
+            None,
+            id='single_row_trials_list',
+        ),
+        pytest.param(
+            {'onsets': [0], 'offsets': [1], 'trials': ['A']},
+            ['trial'],
+            id='single_row_trials_list',
+        ),
+        pytest.param(
+            {'data': pl.DataFrame({'onset': [0], 'offset': [1], 'trial': ['A']})},
+            None,
+            id='single_row_trial_column_not_specified',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame({'onset': [0], 'offset': [1], 'trial': ['A']}),
+                'trial_columns': ['trial'],
+            },
+            ['trial'],
+            id='single_row_trial_column_specified',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame({'onset': [0], 'offset': [1], 'group': [1], 'trial': ['C']}),
+                'trial_columns': ['group', 'trial'],
+            },
+            ['group', 'trial'],
+            id='single_row_two_trial_columns',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame({'onset': [0], 'offset': [1], 'trial': ['A']}),
+                'trial_columns': 'trial',
+            },
+            ['trial'],
+            id='single_row_trial_column_str',
+        ),
+    ],
+)
+def test_event_dataframe_init_expected_trial_column_list(kwargs, expected_trial_column_list):
+    events = pm.EventDataFrame(**kwargs)
+
+    assert events.trial_columns == expected_trial_column_list
+
+
+@pytest.mark.parametrize(
+    ('kwargs', 'expected_trial_column_data'),
+    [
+        pytest.param(
+            {'onsets': [0], 'offsets': [1], 'trials': ['A']},
+            pl.Series('trial', ['A']),
+            id='single_row_trials_list',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame({'onset': [0], 'offset': [1], 'trial': ['C']}),
+                'trial_columns': 'trial',
+            },
+            pl.Series('trial', ['C']),
+            id='single_row_trial_column_str',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame({'onset': [0], 'offset': [1], 'trial': ['B']}),
+                'trial_columns': ['trial'],
+            },
+            pl.Series('trial', ['B']),
+            id='single_row_trial_column_list_single',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame({'onset': [0], 'offset': [1], 'group': [1], 'trial': ['C']}),
+                'trial_columns': ['group', 'trial'],
+            },
+            pl.DataFrame({'group': [1], 'trial': ['C']}),
+            id='single_row_two_trial_columns',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame(
+                    {'onset': [0, 2], 'offset': [1, 3], 'trial': [1, 1]},
+                ),
+                'trial_columns': 'trial',
+            },
+            pl.DataFrame({'trial': [1, 1]}),
+            id='two_rows_one_trial',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame(
+                    {'onset': [0, 2], 'offset': [1, 3], 'trial': [1, 2]},
+                ),
+                'trial_columns': 'trial',
+            },
+            pl.DataFrame({'trial': [1, 2]}),
+            id='two_rows_one_trial',
+        ),
+        pytest.param(
+            {
+                'data': pl.DataFrame(
+                    {'onset': [0, 2], 'offset': [1, 3], 'trial': 1},
+                ),
+                'trial_columns': 'trial',
+            },
+            pl.DataFrame({'trial': [1, 1]}),
+            id='two_rows_plain_trial',
+        ),
+    ],
+)
+def test_event_dataframe_init_expected_trial_column_data(kwargs, expected_trial_column_data):
+    events = pm.EventDataFrame(**kwargs)
+
+    if isinstance(expected_trial_column_data, pl.Series):
+        expected_trial_column_data = pl.DataFrame(expected_trial_column_data)
+    assert_frame_equal(events.frame[events.trial_columns], expected_trial_column_data)
+
+
 def test_event_dataframe_columns_same_as_frame():
     init_kwargs = {'onsets': [0], 'offsets': [1]}
     event_df = pm.EventDataFrame(**init_kwargs)