Repository: kutaslab/fitgrid

Issue:
Set epochs table index to EPOCH_ID and TIME during Epochs creation
In `__init__` we reset the index to only keep `EPOCH_ID`. After `snapshots` are created, set index to `EPOCH_ID`, `TIME`. This is needed for plotting of individual epochs.

Comments:

Patch:
diff --git a/fitgrid/epochs.py b/fitgrid/epochs.py
index 1b09227..b14ad8a 100644
--- a/fitgrid/epochs.py
+++ b/fitgrid/epochs.py
@@ -41,12 +41,13 @@ class Epochs:
         levels_to_remove = set(epochs_table.index.names)
         levels_to_remove.discard(EPOCH_ID)
 
-        # so we remove all levels from index except EPOCH_ID
-        epochs_table.reset_index(list(levels_to_remove), inplace=True)
-        assert epochs_table.index.names == [EPOCH_ID]
+        # copy since we are about to modify
+        self.table = epochs_table.copy()
+        # remove all levels from index except EPOCH_ID
+        self.table.reset_index(list(levels_to_remove), inplace=True)
+        assert self.table.index.names == [EPOCH_ID]
 
-        self.table = epochs_table
-        snapshots = epochs_table.groupby(TIME)
+        snapshots = self.table.groupby(TIME)
 
         # check that snapshots across epochs have equal index by transitivity
         prev_group = None
@@ -66,10 +67,13 @@ 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(epochs_table, EPOCH_ID),
+                tools.get_index_duplicates_table(self.table, EPOCH_ID),
             )
 
-        # we're good, set instance variable
+        self.table.reset_index(inplace=True)
+        self.table.set_index([EPOCH_ID, TIME], inplace=True)
+        assert self.table.index.names == [EPOCH_ID, TIME]
+
         self.snapshots = snapshots
 
     def lm(self, LHS='default', RHS=None):


Test patch:
diff --git a/tests/test_epochs.py b/tests/test_epochs.py
index 736db4b..a92bf85 100644
--- a/tests/test_epochs.py
+++ b/tests/test_epochs.py
@@ -2,7 +2,8 @@ import pytest
 import numpy as np
 
 from .context import fitgrid
-from fitgrid import fake_data, epochs, errors
+from fitgrid import fake_data, errors
+from fitgrid.epochs import Epochs
 
 
 def test_epochs_unequal_snapshots():
@@ -13,7 +14,7 @@ def test_epochs_unequal_snapshots():
 
     epochs_table.drop(epochs_table.index[42], inplace=True)
     with pytest.raises(errors.FitGridError) as error:
-        epochs.Epochs(epochs_table)
+        Epochs(epochs_table)
     assert 'differs from previous snapshot' in str(error.value)
 
 
@@ -34,6 +35,23 @@ 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(epochs_table)
+        Epochs(epochs_table)
 
     assert 'differs from previous snapshot' in str(error.value)
+
+
+def test_multiple_indices_end_up_EPOCH_ID_and_TIME():
+
+    from fitgrid import EPOCH_ID, TIME
+
+    epochs_table = 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)
+    # internal table has EPOCH_ID and TIME in index
+    assert epochs.table.index.names == [EPOCH_ID, TIME]
+    # input table is not altered
+    assert epochs_table.index.names == [EPOCH_ID, TIME, 'categorical']