Repository: kutaslab/fitgrid

Issue:
Make Epochs and FitGrid channel aware
`Epochs` and `FitGrid` objects should know what channels (dependent variables) are inside them. `FitGrid`: easy, it's the columns. `Epochs`: the value of `fitgrid.CHANNELS` at the time of creation.

This makes plotting averages or running models easier since we have a reasonable default for target columns and don't have to pass a list of channels as parameter to every method.

Additionally, this allows for multiple `Epochs` objects with different channels.

The disadvantage is that it's inexplicit, you set an object property indirectly by setting a module-level constant... Maybe pass a list of channels on creation?

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/kutaslab/fitgrid/pull/30?src=pr&el=h1) Report
> Merging [#30](https://codecov.io/gh/kutaslab/fitgrid/pull/30?src=pr&el=desc) into [dev](https://codecov.io/gh/kutaslab/fitgrid/commit/18722a8e8b95ee02f3578e425967adc63778bbe8?src=pr&el=desc) will **increase** coverage by `0.19%`.
> The diff coverage is `79.41%`.

[![Impacted file tree graph](https://codecov.io/gh/kutaslab/fitgrid/pull/30/graphs/tree.svg?width=650&token=d9xVBW7gSn&height=150&src=pr)](https://codecov.io/gh/kutaslab/fitgrid/pull/30?src=pr&el=tree)

```diff
@@           Coverage Diff            @@
##             dev     #30      +/-   ##
========================================
+ Coverage   87.6%   87.8%   +0.19%     
========================================
  Files          8       8              
  Lines        234     246      +12     
========================================
+ Hits         205     216      +11     
- Misses        29      30       +1
```


| [Impacted Files](https://codecov.io/gh/kutaslab/fitgrid/pull/30?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [fitgrid/fitgrid.py](https://codecov.io/gh/kutaslab/fitgrid/pull/30/diff?src=pr&el=tree#diff-Zml0Z3JpZC9maXRncmlkLnB5) | `88.13% <100%> (+0.1%)` | :arrow_up: |
| [fitgrid/fake\_data.py](https://codecov.io/gh/kutaslab/fitgrid/pull/30/diff?src=pr&el=tree#diff-Zml0Z3JpZC9mYWtlX2RhdGEucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [fitgrid/epochs.py](https://codecov.io/gh/kutaslab/fitgrid/pull/30/diff?src=pr&el=tree#diff-Zml0Z3JpZC9lcG9jaHMucHk=) | `81.35% <75%> (+1.76%)` | :arrow_up: |
| [fitgrid/io.py](https://codecov.io/gh/kutaslab/fitgrid/pull/30/diff?src=pr&el=tree#diff-Zml0Z3JpZC9pby5weQ==) | `71.42% <80%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/kutaslab/fitgrid/pull/30?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/kutaslab/fitgrid/pull/30?src=pr&el=footer). Last update [18722a8...677a6dd](https://codecov.io/gh/kutaslab/fitgrid/pull/30?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/fitgrid/epochs.py b/fitgrid/epochs.py
index 40004d2..459530e 100644
--- a/fitgrid/epochs.py
+++ b/fitgrid/epochs.py
@@ -22,23 +22,42 @@ class Epochs:
 
     """
 
-    def __init__(self, epochs_table):
+    def __init__(self, epochs_table, channels='default'):
 
-        from . import EPOCH_ID, TIME
+        from . import EPOCH_ID, TIME, CHANNELS
+
+        if channels == 'default':
+            channels = CHANNELS
+
+        # channels must be a list of strings
+        if not isinstance(channels, list) or not all(
+            isinstance(item, str) for item in channels
+        ):
+            raise FitGridError('channels should be a list of strings.')
+
+        # all channels must be present as epochs table columns
+        missing_channels = set(channels) - set(epochs_table.columns)
+        if missing_channels:
+            raise FitGridError(
+                'channels should all be present in the epochs table, '
+                f'the following are missing: {missing_channels}'
+            )
 
         if not isinstance(epochs_table, pd.DataFrame):
             raise FitGridError('epochs_table must be a Pandas DataFrame.')
 
         # these index columns are required for consistency checks
-        assert (
-            TIME in epochs_table.index.names
-            and EPOCH_ID in epochs_table.index.names
-        )
+        for item in (EPOCH_ID, TIME):
+            if item not in epochs_table.index.names:
+                raise FitGridError(
+                    f'{item} must be a column in the epochs table index.'
+                )
 
-        self.table = epochs_table.copy().reset_index().set_index(EPOCH_ID)
-        assert self.table.index.names == [EPOCH_ID]
+        # make our own copy so we are immune to modification to original table
+        table = epochs_table.copy().reset_index().set_index(EPOCH_ID)
+        assert table.index.names == [EPOCH_ID]
 
-        snapshots = self.table.groupby(TIME)
+        snapshots = table.groupby(TIME)
 
         # check that snapshots across epochs have equal index by transitivity
         prev_group = None
@@ -58,13 +77,16 @@ class Epochs:
         if not prev_group.index.is_unique:
             raise FitGridError(
                 f'Duplicate values in {EPOCH_ID} index not allowed:',
-                tools.get_index_duplicates_table(self.table, EPOCH_ID),
+                tools.get_index_duplicates_table(table, EPOCH_ID),
             )
 
+        # checks passed, set instance variables
+        self.channels = channels
+        self.table = table
         self.snapshots = snapshots
         self.epoch_index = snapshots.get_group(0).index.copy()
 
-    def lm(self, LHS='default', RHS=None):
+    def lm(self, LHS=None, RHS=None):
         """Run ordinary least squares linear regression on the epochs.
 
         Parameters
@@ -82,10 +104,8 @@ class Epochs:
 
         """
 
-        if LHS == 'default':
-            from . import CHANNELS
-
-            LHS = CHANNELS
+        if LHS is None:
+            LHS = self.channels
 
         # validate LHS
         if not (
@@ -94,7 +114,13 @@ class Epochs:
         ):
             raise FitGridError('LHS must be a list of strings.')
 
-        assert set(LHS).issubset(set(self.table.columns))
+        # all LHS items must be present in the epochs_table
+        missing = set(LHS) - set(self.table.columns)
+        if missing:
+            raise FitGridError(
+                'Items in LHS should all be present in the epochs table, '
+                f'the following are missing: {missing}'
+            )
 
         # validate RHS
         if RHS is None:
@@ -128,13 +154,7 @@ class Epochs:
         from . import CHANNELS
 
         if channels is None:
-            if set(CHANNELS).issubset(set(self.table.columns)):
-                channels = CHANNELS
-            else:
-                raise FitGridError(
-                    f'Default channels {CHANNELS} missing in epochs table,'
-                    ' please pass list of channels or set CHANNELS.'
-                )
+            channels = self.channels
 
         from . import plots
 
diff --git a/fitgrid/fake_data.py b/fitgrid/fake_data.py
index e4b06fe..f0ab176 100644
--- a/fitgrid/fake_data.py
+++ b/fitgrid/fake_data.py
@@ -37,12 +37,12 @@ def generate(n_epochs=10, n_samples=100, n_categories=2, n_channels=32):
     = 2` produces 20 epochs, 10 per category.
     """
 
-    df = _generate(n_epochs, n_samples, n_categories, n_channels)
-    return Epochs(df)
+    df, channels = _generate(n_epochs, n_samples, n_categories, n_channels)
+    return Epochs(df, channels)
 
 
 def _generate(n_epochs, n_samples, n_categories, n_channels):
-    """Return Pandas DataFrame with fake EEG data."""
+    """Return Pandas DataFrame with fake EEG data, and a list of channels."""
 
     from . import EPOCH_ID, TIME
 
@@ -60,13 +60,14 @@ def _generate(n_epochs, n_samples, n_categories, n_channels):
         'continuous': np.random.uniform(size=total),
     }
 
+    channels = [f'channel{i}' for i in range(n_channels)]
     eeg = {
-        f'channel{i}': np.random.normal(loc=0, scale=30, size=total)
-        for i in range(n_channels)
+        channel: np.random.normal(loc=0, scale=30, size=total)
+        for channel in channels
     }
 
     data = {**indices, **predictors, **eeg}
 
     df = pd.DataFrame(data).set_index([EPOCH_ID, TIME]).sort_index()
 
-    return df
+    return df, channels
diff --git a/fitgrid/fitgrid.py b/fitgrid/fitgrid.py
index 3a7688d..a67ddde 100644
--- a/fitgrid/fitgrid.py
+++ b/fitgrid/fitgrid.py
@@ -130,7 +130,7 @@ class FitGrid:
         self._grid = _grid
         self.epoch_index = epoch_index
         self.tester = _grid.iloc[0, 0]
-        # TODO the grid should be aware of the channel names
+        self.channels = list(_grid.columns)
 
     def __getitem__(self, slicer):
         """Slice grid on time and channels, return new grid with that shape.
diff --git a/fitgrid/io.py b/fitgrid/io.py
index 20f98f3..60bfb58 100644
--- a/fitgrid/io.py
+++ b/fitgrid/io.py
@@ -3,7 +3,7 @@ from .epochs import Epochs
 from .errors import FitGridError
 
 
-def epochs_from_hdf(hdf_filename):
+def epochs_from_hdf(hdf_filename, channels='default'):
     """Construct Epochs object from an HDF5 file containing an epochs table.
 
     The HDF5 file should contain columns with names defined by EPOCH_ID and
@@ -28,19 +28,19 @@ def epochs_from_hdf(hdf_filename):
 
     # time and epoch id already present in index
     if EPOCH_ID in df.index.names and TIME in df.index.names:
-        return Epochs(df)
+        return Epochs(df, channels)
 
     # time and epoch id present in columns, set index
     if EPOCH_ID in df.columns and TIME in df.columns:
         df.set_index([EPOCH_ID, TIME], inplace=True)
-        return Epochs(df)
+        return Epochs(df, channels)
 
     raise FitGridError(
         f'Dataset has to contain {EPOCH_ID} and {TIME} as columns or indices.'
     )
 
 
-def epochs_from_dataframe(dataframe):
+def epochs_from_dataframe(dataframe, channels='default'):
     """Construct Epochs object from a Pandas DataFrame epochs table.
 
     The DataFrame should contain columns with names defined by EPOCH_ID and
@@ -50,10 +50,12 @@ def epochs_from_dataframe(dataframe):
     ----------
     dataframe : pandas DataFrame
         a pandas DataFrame object
+    channels : list of str, optional, defaults to CHANNELS
+        list of string channel names
 
     Returns
     -------
     epochs : Epochs
         an Epochs object with the data
     """
-    return Epochs(dataframe)
+    return Epochs(dataframe, channels)


Test patch:
diff --git a/tests/test_epochs.py b/tests/test_epochs.py
index 16ddf99..c0a40f0 100644
--- a/tests/test_epochs.py
+++ b/tests/test_epochs.py
@@ -1,24 +1,30 @@
 import pytest
 import numpy as np
+from .context import fitgrid
+from fitgrid import fake_data, errors
+from fitgrid.epochs import Epochs
+
 import matplotlib
 
 matplotlib.use('Agg')
 
 
-from .context import fitgrid
-from fitgrid import fake_data, errors
-from fitgrid.epochs import Epochs
+def test_epochs_lm_bad_inputs():
+
+    epochs = fake_data.generate()
+    with pytest.raises(errors.FitGridError):
+        epochs.lm(LHS=['bad_channel'], RHS='categorical')
 
 
 def test_epochs_unequal_snapshots():
 
-    epochs_table = fake_data._generate(
+    epochs_table, channels = fake_data._generate(
         n_epochs=10, n_samples=100, n_categories=2, n_channels=32
     )
 
     epochs_table.drop(epochs_table.index[42], inplace=True)
     with pytest.raises(errors.FitGridError) as error:
-        Epochs(epochs_table)
+        Epochs(epochs_table, channels)
     assert 'differs from previous snapshot' in str(error.value)
 
 
@@ -28,7 +34,7 @@ def test__raises_error_on_epoch_index_mismatch():
     from fitgrid import TIME
 
     # strategy: generate epochs, but insert meaningless time index
-    epochs_table = fake_data._generate(
+    epochs_table, channels = fake_data._generate(
         n_epochs=10, n_samples=100, n_categories=2, n_channels=32
     )
 
@@ -39,7 +45,7 @@ def test__raises_error_on_epoch_index_mismatch():
 
     # now time index is equal to row number in the table overall
     with pytest.raises(errors.FitGridError) as error:
-        Epochs(epochs_table)
+        Epochs(epochs_table, channels)
 
     assert 'differs from previous snapshot' in str(error.value)
 
@@ -48,13 +54,13 @@ def test_multiple_indices_end_up_EPOCH_ID():
 
     from fitgrid import EPOCH_ID, TIME
 
-    epochs_table = fake_data._generate(
+    epochs_table, channels = fake_data._generate(
         n_epochs=10, n_samples=100, n_categories=2, n_channels=32
     )
     epochs_table.reset_index(inplace=True)
     epochs_table.set_index([EPOCH_ID, TIME, 'categorical'], inplace=True)
 
-    epochs = Epochs(epochs_table)
+    epochs = Epochs(epochs_table, channels)
     # internal table has EPOCH_ID in index
     assert epochs.table.index.names == [EPOCH_ID]
     # input table is not altered
diff --git a/tests/test_fake_data.py b/tests/test_fake_data.py
index e368f86..e5e738d 100644
--- a/tests/test_fake_data.py
+++ b/tests/test_fake_data.py
@@ -7,7 +7,7 @@ def test__generate():
     n_epochs = 10
     n_categories = 2
     n_channels = 32
-    table = fg.fake_data._generate(
+    table, channels = fg.fake_data._generate(
         n_samples=n_samples,
         n_epochs=n_epochs,
         n_categories=n_categories,
@@ -37,4 +37,5 @@ def test__generate():
         assert len(category.index.unique(fg.EPOCH_ID)) == n_epochs
 
     # test number of columns (2 predictors + n_channels)
+    assert len(channels) == n_channels
     assert len(table.columns) == 2 + n_channels
diff --git a/tests/test_fitgrid.py b/tests/test_fitgrid.py
index bd3a294..8667a18 100644
--- a/tests/test_fitgrid.py
+++ b/tests/test_fitgrid.py
@@ -7,6 +7,13 @@ import matplotlib
 matplotlib.use('Agg')
 
 
+def test__correct_channels_in_fitgrid():
+    epochs = fitgrid.generate()
+    LHS = ['channel0', 'channel1', 'channel2']
+    grid = epochs.lm(LHS=LHS, RHS='categorical + continuous')
+    assert grid.channels == LHS
+
+
 def test__method_returning_dataframe_expands_correctly():
 
     epochs = fitgrid.generate()
@@ -56,16 +63,17 @@ def test__epoch_id_substitution():
     from fitgrid import EPOCH_ID
 
     # create data with unusual index (shifted by 5)
-    data = fitgrid.fake_data._generate(10, 100, 2, 32)
+    data, channels = fitgrid.fake_data._generate(10, 100, 2, 32)
     unusual_index = np.arange(20) + 5
     data.index.set_levels(unusual_index, level=EPOCH_ID, inplace=True)
-    epochs = fitgrid.epochs_from_dataframe(data)
+    epochs = fitgrid.epochs_from_dataframe(data, channels)
 
     # remember epoch_index
     epoch_index = epochs.snapshots.get_group(0).index
     assert (epoch_index == unusual_index).all()
 
-    LHS = ['channel0', 'channel1']
+    # take just two channels for speed
+    LHS = channels[:2]
     grid = epochs.lm(LHS=LHS, RHS='categorical + continuous')
 
     # one additional level
diff --git a/tests/test_io.py b/tests/test_io.py
index 509de40..104e54f 100644
--- a/tests/test_io.py
+++ b/tests/test_io.py
@@ -3,7 +3,8 @@ from .context import fitgrid
 
 def test__epochs_from_hdf():
     TEST_FILE = 'tests/data/fake_epochs.h5'
-    fitgrid.epochs_from_hdf(TEST_FILE)
+    channels = [f'channel{i}' for i in range(32)]
+    fitgrid.epochs_from_hdf(TEST_FILE, channels)
 
 
 def test__epochs_from_dataframe_good_data():
@@ -11,7 +12,7 @@ def test__epochs_from_dataframe_good_data():
     # should silently create an Epochs object from normally generated data
     # this is a helper function, all further testing of Epochs creation is
     # done in test_epochs.py
-    table = fitgrid.fake_data._generate(
+    table, channels = fitgrid.fake_data._generate(
         n_epochs=10, n_samples=100, n_categories=2, n_channels=32
     )
-    fitgrid.epochs_from_dataframe(table)
+    fitgrid.epochs_from_dataframe(table, channels)