Loss function

\

@@ -0,0 +1,119 @@
+from pathlib import Path
+import tempfile
+import numpy as np
+from numpy.random import randint
+from mtrf.model import TRF, load_sample_data
+
+n = np.random.randint(3, 10)
+stimulus, response, fs = load_sample_data(n_segments=n)
+
+
+def test_loss_function():
+    def mean_absolute_deviation(y, y_pred):
+        return np.abs(y - y_pred).mean()
+
+    def mean_difference(y, y_pred):
+        return np.mean(y - y_pred)
+
+    for loss_function in [mean_absolute_deviation, mean_difference]:
+        trf = TRF(loss_function=loss_function)
+        regularization = [
+            np.random.uniform(0, 1000) for _ in range(np.random.randint(2, 5))
+        ]
+        tmin = np.random.uniform(-0.1, 0.05)
+        tmax = np.random.uniform(0.1, 0.4)
+        loss = trf.train(stimulus, response, fs, tmin, tmax, regularization)
+        assert trf.regularization == regularization[np.argmin(loss)]
+
+
+def test_train():
+    tmin = np.random.uniform(-0.1, 0.05)
+    tmax = np.random.uniform(0.1, 0.4)
+    regularization = np.random.uniform(0, 10)
+    for direction in [1, -1]:
+        trf = TRF(direction=direction)
+        trf.train(stimulus, response, fs, tmin, tmax, regularization)
+        if direction == 1:
+            assert trf.weights.shape[0] == stimulus[0].shape[-1]
+            assert trf.weights.shape[-1] == response[0].shape[-1]
+        if direction == -1:
+            assert trf.weights.shape[-1] == stimulus[0].shape[-1]
+            assert trf.weights.shape[0] == response[0].shape[-1]
+
+
+def test_optimize():
+    for direction in [1, -1]:
+        tmin = np.random.uniform(-0.1, 0.05)
+        tmax = np.random.uniform(0.1, 0.4)
+        trf = TRF(direction=direction)
+        regularization = [np.random.uniform(0, 10) for _ in range(randint(2, 5))]
+        loss = trf.train(stimulus, response, fs, tmin, tmax, regularization)
+    assert len(loss) == len(regularization)
+
+
+def test_predict():
+    tmin = np.random.uniform(-0.1, 0.05)
+    tmax = np.random.uniform(0.1, 0.4)
+    regularization = np.random.uniform(0, 10)
+    trf = TRF()
+    trf.train(stimulus, response, fs, tmin, tmax, regularization)
+    for average in [True, list(range(randint(response[0].shape[-1])))]:
+        prediction, loss = trf.predict(stimulus, response, average=average)
+        assert len(prediction) == len(response)
+        assert all([p[0].shape == r[0].shape for p, r in zip(prediction, response)])
+        assert np.isscalar(loss)
+    prediction, loss = trf.predict(stimulus, response, average=False)
+    assert loss.shape[-1] == trf.weights.shape[-1]
+    trf = TRF(direction=-1)
+    trf.train(stimulus, response, fs, tmin, tmax, regularization)
+    for average in [True, list(range(randint(stimulus[0].shape[-1])))]:
+        prediction, loss = 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(loss)
+    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)
+    tmax = np.random.uniform(0.1, 0.4)
+    reg = [np.random.uniform(0, 10) for _ in range(randint(2, 10))]
+    trf = TRF()
+    loss, best_regularization = trf.test(stimulus, response, fs, tmin, tmax, reg)
+    assert len(loss) == len(best_regularization) == n
+
+
+def test_save_load():
+    tmpdir = Path(tempfile.gettempdir())
+    tmin = np.random.uniform(-0.1, 0.05)
+    tmax = np.random.uniform(0.1, 0.4)
+    direction = np.random.choice([1, -1])
+    regularization = np.random.uniform(0, 10)
+    trf1 = TRF(direction=direction)
+    trf1.train(stimulus, response, fs, tmin, tmax, regularization)
+    trf1.save(tmpdir / "test.trf")
+    trf2 = TRF()
+    trf2.load(tmpdir / "test.trf")
+    np.testing.assert_equal(trf1.weights, trf2.weights)
+
+
+def test_no_preload(decimal=10):
+    for direction in [1, -1]:
+        tmin = np.random.uniform(-0.1, 0.05)
+        tmax = np.random.uniform(0.1, 0.2)
+        regularization = [np.random.uniform(0, 10) for _ in range(2)]
+        stimulus, response, fs = load_sample_data(n_segments=4)
+        # testing forward model
+        trf1 = TRF()
+        trf1.train(stimulus, response, fs, tmin, tmax, regularization)
+        prediction1, loss1 = trf1.predict(stimulus, response, average=False)
+        trf2 = TRF(preload=False)
+        trf2.train(stimulus, response, fs, tmin, tmax, regularization)
+        prediction2, loss2 = trf2.predict(stimulus, response, average=False)
+        # assert all close
+        np.testing.assert_almost_equal(trf1.weights, trf2.weights, decimal=decimal)
+        np.testing.assert_almost_equal(trf1.bias, trf2.bias, decimal=decimal)
+        for pred1, pred2 in zip(prediction1, prediction2):
+            np.testing.assert_almost_equal(pred1, pred2, decimal=decimal)
+        np.testing.assert_almost_equal(loss1, loss2, decimal=decimal)

mtrf/matrices.py

@@ -132,24 +132,19 @@ def covariance_matrices(x, y, lags, zeropad=True, preload=True):
     x, y, _ = _check_data(x, y)
     if zeropad is False:
         y = truncate(y, lags[0], lags[-1])
-
-    x_lag_size = x[0].shape[-1] * len(lags) + 1
-    cov_xx_size = (x_lag_size, x_lag_size)
-    cov_xy_size = (x_lag_size, y[0].shape[-1])
-
-    if preload is True:
-        cov_xx = np.zeros((len(x), *cov_xx_size))
-        cov_xy = np.zeros((len(y), *cov_xy_size))
-        for i_x in range(len(x)):
-            x_lag = lag_matrix(x[i_x], lags, zeropad)
-            cov_xx[i_x] = x_lag.T @ x_lag
-            cov_xy[i_x] = x_lag.T @ y[i_x]
-    else:
-        cov_xx, cov_xy = 0, 0
-        for i_x in range(len(x)):
-            x_lag = lag_matrix(x[i_x], lags, zeropad)
+    cov_xx, cov_xy = 0, 0
+    for i, (x_i, y_i) in enumerate(zip(x, y)):
+        x_lag = lag_matrix(x_i, lags, zeropad)
+        if preload is True:
+            if i == 0:
+                cov_xx = np.zeros((len(x), x_lag.shape[-1], x_lag.shape[-1]))
+                cov_xy = np.zeros((len(y), x_lag.shape[-1], y_i.shape[-1]))
+            cov_xx[i] = x_lag.T @ x_lag
+            cov_xy[i] = x_lag.T @ y_i
+        else:
             cov_xx += x_lag.T @ x_lag
-            cov_xy += x_lag.T @ y[i_x]
+            cov_xy += x_lag.T @ y_i
+    if preload is False:
         cov_xx, cov_xy = cov_xx / len(x), cov_xy / len(x)
 
     return cov_xx, cov_xy