Repository: kutaslab/fitgrid

Issue:
add design matrix check
if the design matrix, given a formula, doesn't change across time/channels, we can run certain computations much faster

Discussion:

Patch:
diff --git a/fitgrid/epochs.py b/fitgrid/epochs.py
index 9a33eba..f252bc9 100644
--- a/fitgrid/epochs.py
+++ b/fitgrid/epochs.py
@@ -60,7 +60,9 @@ class Epochs:
             raise FitGridError('Duplicate column names not allowed.')
 
         # make our own copy so we are immune to modification to original table
-        table = epochs_table.copy().reset_index().set_index(EPOCH_ID)
+        table = (
+            epochs_table.copy().reset_index().set_index(EPOCH_ID).sort_index()
+        )
         assert table.index.names == [EPOCH_ID]
 
         snapshots = table.groupby(TIME)
diff --git a/fitgrid/tools.py b/fitgrid/tools.py
index d0a2717..fabc9ed 100644
--- a/fitgrid/tools.py
+++ b/fitgrid/tools.py
@@ -182,3 +182,63 @@ class single_threaded:
                     f'to {self.old_n_threads} threads (previous value).'
                 )
                 raise RuntimeError(message)
+
+
+def design_matrix_is_constant(df, columns, time):
+    """Check that values in columns of df do not change within any epoch.
+
+    See Notes for more details.
+
+    Parameters
+    ----------
+    df : pandas.DataFrame
+        dataframe to check
+    columns : list of str
+        list of column names to be checked
+    time : str
+        name of the time column
+
+    Returns
+    -------
+    result : bool
+        True if values in specified columns don't change, False otherwise
+
+
+    Notes
+    -----
+    We check that from timepoint to timepoint, each epoch has the same value in
+    a given column:
+
+                epoch1
+        ==============
+             A |    B
+        ==============
+             1      x
+             1      x
+             1      x
+             1      x
+             1      x
+
+                epoch2
+        ==============
+             A |    B
+        ==============
+             2      y
+             2      y
+             2      y
+             2      y
+             2      y
+
+    This is helpful when performing linear regression on an epochs table where
+    the predictors vary with epochs (as they are expected to) but stay constant
+    from sample to sample, because we can do our modeling much faster.
+    """
+    gb = df.groupby(time)
+    _, group = next(iter(gb))  # first group
+
+    df_columns_values = df[columns].values
+    group_columns_values = group[columns].values
+
+    expected_df_values = np.repeat(group_columns_values, len(gb), axis=0)
+
+    return (df_columns_values == expected_df_values).all()


Test patch:
diff --git a/tests/test_tools.py b/tests/test_tools.py
index e001ff5..b98c7d2 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -87,3 +87,109 @@ def test_single_threaded_change_before():
         assert blas.get_n_threads() == 1
 
     assert blas.get_n_threads() == BEFORE
+
+
+def test_design_matrix_is_constant():
+
+    df = pd.DataFrame(
+        dict(
+            epoch_id=[
+                0,  # epoch 0
+                0,
+                0,
+                0,
+                1,  # epoch 1
+                1,
+                1,
+                1,
+                2,  # epoch 2
+                2,
+                2,
+                2,
+            ],
+            time=[
+                0,  # epoch 0
+                1,
+                2,
+                3,
+                0,  # epoch 1
+                1,
+                2,
+                3,
+                0,  # epoch 2
+                1,
+                2,
+                3,
+            ],
+            b=[
+                'a',  # epoch 0
+                'a',
+                'a',
+                'a',
+                'b',  # epoch 1
+                'b',
+                'b',
+                'b',
+                'c',  # epoch 2
+                'c',
+                'c',
+                'c',
+            ],
+            c=[
+                10,  # epoch 0
+                10,
+                10,
+                10,
+                20,  # epoch 1
+                20,
+                20,
+                20,
+                30,  # epoch 2
+                30,
+                30,
+                30,
+            ],
+        )
+    ).set_index(['epoch_id', 'time'])
+
+    assert tools.design_matrix_is_constant(df, ['b', 'c'], 'time')
+
+    assert not tools.design_matrix_is_constant(
+        df.sample(frac=1), ['b', 'c'], 'time'
+    )
+
+    assert tools.design_matrix_is_constant(
+        df.sample(frac=1).sort_index(), ['b', 'c'], 'time'
+    )
+
+    df['c'] = [
+        0,  # epoch 0
+        1,
+        2,
+        3,
+        4,  # epoch 1
+        5,
+        6,
+        7,
+        8,  # epoch 2
+        9,
+        10,
+        11,
+    ]
+    assert not tools.design_matrix_is_constant(df, ['b', 'c'], 'time')
+
+    df['c'] = [
+        0,  # epoch 0
+        0,
+        0,
+        0,
+        1,  # epoch 1
+        1,
+        10000,
+        1,
+        2,  # epoch 2
+        2,
+        2,
+        2,
+    ]
+    assert not tools.design_matrix_is_constant(df, ['b', 'c'], 'time')