Fix banded tests

Author: vinaysraghavan

naplib/encoding/banded_trf.py

@@ -160,7 +160,7 @@ def fit(self, data=None, X=['aud'], y='resp'):
         self.feature_alphas_ = []
         self.alpha_paths = np.zeros((len(X), len(self.alphas)))
 
-        for i, current_feat in enumerate(X):
+        for i in range(len(X)):
             best_alpha = None
             max_r = -np.inf
             r_history = []
@@ -262,6 +262,7 @@ def predict(self, data=None, X=['aud']):
 
         X_mats = self._prepare_matrix(X, self.feature_alphas_)
         n_trials = len(X_mats)
+        n_feat_dim = X_mats[0].shape[1]
 
         if n_trials != len(self.model_):
             raise ValueError(
@@ -274,6 +275,12 @@ def predict(self, data=None, X=['aud']):
             # Expand (trials, features) -> (trials, 1_target, features)
             all_coefs = all_coefs[:, np.newaxis, :]
 
+        if n_feat_dim < all_coefs.shape[2]:
+            all_coefs = all_coefs[:, :, :n_feat_dim]
+            print('Using reduced TRF')
+        elif n_feat_dim > all_coefs.shape[2]:
+            raise ValueError('Too many features for trained model.')
+
         preds = []
         for i in range(n_trials):
             # Indices for all trials except the current one
@@ -282,11 +289,8 @@ def predict(self, data=None, X=['aud']):
             # Average coefficients and intercepts from the other trials
             loto_coef = np.mean(all_coefs[loto_indices], axis=0)
 
-            # Only use first coef features if using feature subset
-            sliced_coef = loto_coef[:, :X_mats[i].shape[1]]
-            
             # Predict for the current trial
-            preds.append(X_mats[i] @ sliced_coef.T)
+            preds.append(X_mats[i] @ loto_coef.T)
 
         return preds
 

tests/encoding/test_banded_trf.py

@@ -38,7 +38,7 @@ def test_banded_trf_loto_consistency(synth_data):
     """Test that coef_ property handles the 4D reshape correctly."""
     model = BandedTRF(tmin=synth_data['tmin'], tmax=synth_data['tmax'], 
                       sfreq=synth_data['sfreq'], alphas=[0.1, 10.0])
-    model.fit(data=synth_data['data'], feature_order=synth_data['feature_order'], target='resp')
+    model.fit(data=synth_data['data'], X=synth_data['feature_order'], target='resp')
 
     # Shape calculation: 2 targets, 2 features, 4 delays, 3 trials.
     # ndelays = (0.03 * 100) - (0 * 100) + 1 = 4.
@@ -47,7 +47,7 @@ def test_banded_trf_loto_consistency(synth_data):
 def test_predict_masking_logic(synth_data):
     """Verify that partial feature prediction works with multi-channel targets."""
     model = BandedTRF(tmin=synth_data['tmin'], tmax=synth_data['tmax'], sfreq=synth_data['sfreq'])
-    model.fit(data=synth_data['data'], feature_order=synth_data['feature_order'], target='resp')
+    model.fit(data=synth_data['data'], X=synth_data['feature_order'], target='resp')
 
     # Full prediction: should match target shape (samples, channels)
     preds_all = model.predict(synth_data['data'])
@@ -63,7 +63,7 @@ def test_predict_masking_logic(synth_data):
 def test_summary_p_values(synth_data):
     """Verify summary table computes stats across channels correctly."""
     model = BandedTRF(tmin=synth_data['tmin'], tmax=synth_data['tmax'], sfreq=synth_data['sfreq'])
-    model.fit(data=synth_data['data'], feature_order=synth_data['feature_order'], target='resp')
+    model.fit(data=synth_data['data'], X=synth_data['feature_order'], target='resp')
 
     df = model.summary()
     assert isinstance(df, pd.DataFrame)
@@ -81,7 +81,7 @@ def test_unfitted_attribute_error():
 def test_predict_trial_mismatch(synth_data):
     """LOTO requires the same number of trials for predict as fit."""
     model = BandedTRF(tmin=synth_data['tmin'], tmax=synth_data['tmax'], sfreq=synth_data['sfreq'])
-    model.fit(data=synth_data['data'], feature_order=synth_data['feature_order'], target='resp')
+    model.fit(data=synth_data['data'], X=synth_data['feature_order'], target='resp')
 
     # Try predicting with only 2 trials instead of 3
     short_data = synth_data['data'][:2]