Merge pull request #301 from flatironinstitute/negative_axis_id

Negative axis

Author: BalzaniEdoardo

src/nemos/basis/_basis.py

@@ -665,6 +665,9 @@ def split_by_feature(
             - **Key**: Label of the basis.
             - **Value**: the array reshaped to: ``(..., n_inputs, n_basis_funcs, ...)``
         """
+        # convert axis to positive ints
+        axis = axis if axis >= 0 else x.ndim + axis
+
         if x.shape[axis] != self.n_output_features:
             raise ValueError(
                 "`x.shape[axis]` does not match the expected number of features."

src/nemos/convolve.py

@@ -90,6 +90,8 @@ def _shift_time_axis_and_convolve(array: NDArray, eval_basis: NDArray, axis: int
     -----
     This function supports arrays of any dimensionality greater or equal than 1.
     """
+    # convert axis
+    axis = axis if axis >= 0 else array.ndim + axis
     # move time axis to first
     new_axis = (jnp.arange(array.ndim) + axis) % array.ndim
     array = jnp.transpose(array, new_axis)

src/nemos/utils.py

@@ -156,6 +156,8 @@ def _pad_dimension(
         "acausal": ((pad_size) // 2, pad_size - (pad_size) // 2),
         "anti-causal": (0, pad_size),
     }
+    # convert negative axis in jax jit compilable way
+    axis = axis * (axis >= 0) + (array.ndim + axis) * (axis < 0)
 
     pad_width = (
         ((0, 0),) * axis

tests/test_basis.py

@@ -5156,6 +5156,8 @@ def test_duplicate_keys(bas1, bas2, bas3, basis_class_specific_params):
     "x, axis, expectation, exp_shapes",  # num output is 5*2 + 6*3 = 28
     [
         (np.ones((1, 28)), 1, does_not_raise(), [(1, 2, 5), (1, 3, 6)]),
+        (np.ones((1, 28)), -1, does_not_raise(), [(1, 2, 5), (1, 3, 6)]),
+        (np.ones((1, 28, 2)), -2, does_not_raise(), [(1, 2, 5, 2), (1, 3, 6, 2)]),
         (np.ones((28,)), 0, does_not_raise(), [(2, 5), (3, 6)]),
         (np.ones((2, 2, 28)), 2, does_not_raise(), [(2, 2, 2, 5), (2, 2, 3, 6)]),
         (

tests/test_convolution.py

@@ -14,9 +14,11 @@ class TestShiftTimeAxisAndConvolve:
         "time_series, check_func, axis",
         [
             (np.zeros((1, 20)), lambda x: x.ndim == 3, 1),
+            (np.zeros((1, 20)), lambda x: x.ndim == 3, -1),
             (np.zeros((20,)), lambda x: x.ndim == 2, 0),
             (np.zeros((20, 1)), lambda x: x.ndim == 3, 0),
             (np.zeros((1, 20, 1)), lambda x: x.ndim == 4, 0),
+            (np.zeros((1, 20, 1)), lambda x: x.ndim == 4, -3),
         ],
     )
     def test_output_ndim(self, time_series, check_func, axis):