tests for design matrix code

Author: berkgercek

brainbox/modeling/neural_model.py

@@ -192,9 +192,12 @@ def fit(self, train_idx=None, printcond=True):
             raise IndexError('Not all train indices in the trials of design matrix')
 
         # Store training and test indices for self so that .score() method will know what to
-        # operate on
+        # operate on. If all data indices are in train indices, train and test are the same set.
         self.traininds = train_idx
-        self.testinds = self.trialsdf.index[~self.trialsdf.index.isin(train_idx)]
+        if not np.all(np.isin(self.design.trialsdf.index, train_idx)):
+            self.testinds = self.design.trialsdf.index[~self.trialsdf.index.isin(train_idx)]
+        else:
+            self.testinds = train_idx
 
         # Mask for training data
         trainmask = np.isin(self.design.trlabels, train_idx).flatten()

brainbox/tests/fixtures/design_matrix_test.npy

@@ -0,0 +1,62 @@
+import unittest
+import numpy as np
+import pandas as pd
+import brainbox.modeling.design_matrix as bdm
+import brainbox.modeling.utils as mut
+from pathlib import Path
+
+
+class TestModeling(unittest.TestCase):
+
+    def setUp(self):
+        # Params for test
+        self.binwidth = 0.02
+        self.rng = np.random.default_rng(seed=112358)
+
+        # Generate fake trial start, stimon, feedback, and end times
+        starts = np.array([0, 1.48, 2.93, 4.67, 6.01, 7.31, 8.68, 9.99, 11.43, 12.86])
+        ends = np.array([1.35, 2.09, 3.53, 5.23, 6.58, 7.95, 9.37, 11.31, 12.14, 13.26])
+        stons = starts + 0.1
+        fdbks = stons + self.rng.normal(loc=0.1, scale=0.05, size=10)
+
+        # Figure out how many bins each trial is and generate non-monotonic trace of fake wheel
+        trlens = self.binf(ends - starts)
+        fakewheels = [np.cumsum(self.rng.normal(loc=0.01, scale=0.01, size=(x))) for x in trlens]
+
+        # Store trialsdf for later use
+        self.trialsdf = pd.DataFrame({'trial_start': starts,
+                                      'trial_end': ends,
+                                      'stim_onset': stons,
+                                      'feedback': fdbks,
+                                      'wheel_traces': fakewheels})
+
+    def binf(self, x):
+        return np.ceil(x / self.binwidth).astype(int)
+
+    def test_dm_construct(self):
+        """
+        Check whether or not design matrix construction works as intended
+        """
+
+        # Design matrix instance
+        self.design = bdm.DesignMatrix(self.trialsdf,
+                                       vartypes={'trial_start': 'timing',
+                                                 'trial_end': 'timing',
+                                                 'stim_onset': 'timing',
+                                                 'feedback': 'timing',
+                                                 'wheel_traces': 'continuous'})
+
+        # Separate bases for wheel and timing
+        tbases = mut.raised_cosine(0.2, 3, self.binf)
+        wbases = mut.raised_cosine(0.1, 2, self.binf)
+        # Add covariates one by one. Add different offsets on timings to test
+        self.design.add_covariate_timing('start', 'trial_start', tbases, offset=0.02)
+        self.design.add_covariate_timing('stim_on', 'stim_onset', tbases)
+        self.design.add_covariate_timing('feedback', 'feedback', tbases, offset=-0.02)
+        self.design.add_covariate('wheelpos', self.trialsdf.wheel_traces, wbases, offset=-0.1)
+        self.design.compile_design_matrix()  # Finally compile
+        # Load target DM
+        npy_file = Path(__file__).parent.joinpath('fixtures', 'design_matrix_test.npy')
+        if npy_file.exists():
+            ref_dm = np.load(npy_file)
+            self.assertEqual(self.design.dm, ref_dm)