Merge pull request #7 from berndie/train_and_predict

Fix for the final training of the train method for backwards models

mtrf/model.py

@@ -222,7 +222,7 @@ def train(
                     verbose=verbose,
                 )
             best_regularization = list(regularization)[np.argmin(mse)]
-            self._train(stimulus, response, fs, tmin, tmax, best_regularization)
+            self._train(xs, ys, fs, tmin, tmax, best_regularization)
             return r, mse
 
     def _train(self, xs, ys, fs, tmin, tmax, regularization):

tests/test_model.py

@@ -40,6 +40,17 @@ def test_predict():
     prediction, r, mse = trf.predict(stimulus, response, average=False)
     assert r.shape[-1] == mse.shape[-1] == trf.weights.shape[-1]
 
+    # Backwards prediction
+    trf = TRF(-1)
+    trf.train(stimulus, response, fs, tmin, tmax, regularization)
+    for average in [True, list(range(randint(stimulus[0].shape[-1])))]:
+        prediction, r, mse = trf.predict(stimulus, response, average=average)
+        assert len(prediction) == len(stimulus)
+        assert all([p[0].shape == s[0].shape for p, s in zip(prediction, stimulus)])
+        assert np.isscalar(r) and np.isscalar(mse)
+    prediction, r, mse = trf.predict(stimulus, response, average=False)
+    assert r.shape[-1] == mse.shape[-1] == trf.weights.shape[-1]
+
 
 def test_test():
     tmin = np.random.uniform(-0.1, 0.05)