fixed identifiability constraint basis and improved test

flatironinstitute · Oct 10, 2024 · d4f505a · d4f505a
1 parent 4722933
commit d4f505a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/src/nemos/basis.py b/src/nemos/basis.py
@@ -1292,7 +1292,7 @@ def _compute_features(self, *xi: ArrayLike) -> FeatureMatrix:
         # the numpy conversion is important, there is some in-place
         # array modification in basis.
         hstack_pynapple = support_pynapple(conv_type="numpy")(np.hstack)
-        return hstack_pynapple(
+        X = hstack_pynapple(
             (
                 self._basis1._compute_features(
                     *xi[: self._basis1._n_input_dimensionality]
@@ -1302,6 +1302,9 @@ def _compute_features(self, *xi: ArrayLike) -> FeatureMatrix:
                 ),
             ),
         )
+        if self.identifiability_constraints:
+            X = self._apply_identifiability_constraints(X)
+        return X
 
     def _set_kernel(self, *xi: ArrayLike) -> Basis:
         """Call fit on the added basis.
@@ -1425,11 +1428,14 @@ def _compute_features(self, *xi: ArrayLike) -> FeatureMatrix:
 
         """
         kron = support_pynapple(conv_type="numpy")(row_wise_kron)
-        return kron(
+        X = kron(
             self._basis1._compute_features(*xi[: self._basis1._n_input_dimensionality]),
             self._basis2._compute_features(*xi[self._basis1._n_input_dimensionality :]),
             transpose=False,
         )
+        if self.identifiability_constraints:
+            X = self._apply_identifiability_constraints(X)
+        return X
 
 
 class SplineBasis(Basis, abc.ABC):

diff --git a/tests/test_basis.py b/tests/test_basis.py
@@ -5584,7 +5584,9 @@ def test_identifiability_constraints_apply(self, basis_cls1, basis_cls2):
         bas2 = basis_cls2(5, **kwargs2)
         bas_add = bas1 + bas2
         bas_add.identifiability_constraints = True
-        bas_add.compute_features(np.arange(10), np.arange(10))
+        X1 = bas_add.compute_features(np.arange(10), np.arange(10))
+        X2 = bas_add(np.arange(10), np.arange(10))
+        assert np.all(X1 == X2)
 
     @pytest.mark.parametrize("n_basis_input1", [1, 2, 3])
     @pytest.mark.parametrize("n_basis_input2", [1, 2, 3])
@@ -6148,7 +6150,9 @@ def test_identifiability_constraints_apply(self, basis_cls1, basis_cls2):
         bas2 = basis_cls2(5, **kwargs2)
         bas_prod = bas1 * bas2
         bas_prod.identifiability_constraints = True
-        bas_prod.compute_features(np.arange(10), np.arange(10))
+        X1 = bas_prod.compute_features(np.arange(10), np.arange(10))
+        X2 = bas_prod(np.arange(10), np.arange(10))
+        assert np.all(X1 == X2)
 
     @pytest.mark.parametrize("n_basis_input1", [1, 2, 3])
     @pytest.mark.parametrize("n_basis_input2", [1, 2, 3])