require sample axis to be axis=0

[BalzaniEdoardo]

This PR fixes a corner-case bug related to basis in mode="conv".

Issue

Before this PR, one could pass at basis initialization any kwargs from nemos.convolve.create_convolutional_predictor. This included axis meaning that one could specify what's the axis along which the convolution is applied.

What happened in basis.compute_features is that the axis is transposed to the first, and the rest of the axis are flattened to the second axis, so that the compute_features returns an 2darray. However, when the basis is composite, the we check that the number of samples matches for all input, and this check relied on the fact that the sample axis is axis=0.

This caused the following corner case:

>>> b = nmo.basis.RaisedCosineBasisLog(5, mode="conv", window_size=100, axis=1)

>>> # this works because there is one input (the check on the sample axis passes)
>>> b.compute_features(np.ones((2, 200)))
Array([[       nan,        nan,        nan, ...,        nan,        nan,
               nan],
       [       nan,        nan,        nan, ...,        nan,        nan,
               nan],
       [       nan,        nan,        nan, ...,        nan,        nan,
               nan],
       ...,
       [ 2.4915924,  2.4915924,  5.5825534, ..., 20.70655  , 39.878582 ,
        39.878582 ],
       [ 2.4915924,  2.4915924,  5.5825534, ..., 20.70655  , 39.878582 ,
        39.878582 ],
       [ 2.4915924,  2.4915924,  5.5825534, ..., 20.706553 , 39.878582 ,
        39.878586 ]], dtype=float32)

>>> # this fails because the check on the sample axis will raise an error since 3 != 2
>>> (b + b).compute_features(np.ones((2, 200)), np.ones((3, 200)))

Solution

Instead of complicating the check logic, or move it after the transposition happens for all 1D basis, I simplified our life and enforced that the sample axis is always the first one. This is always true for pynapple, and requires the user to transpose an array if the data are organized differently;

Now the initialization would raise an error if axis different from 0 is passed.

>>> b = nmo.basis.RaisedCosineBasisLog(5, mode="conv", window_size=100, axis=1)
ValueError: Invalid `axis=1` provided. Basis requires  requires the convolution to be applied along the first axis (`axis=0`).
Please transpose your input so that the desired axis for convolution is the first dimension (axis=0).

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.15%. Comparing base (6b3a1e8) to head (d51f8d3).

Additional details and impacted files

@@               Coverage Diff               @@
##           development     #254      +/-   ##
===============================================
+ Coverage        97.09%   97.15%   +0.06%     
===============================================
  Files               20       20              
  Lines             1857     1863       +6     
===============================================
+ Hits              1803     1810       +7     
+ Misses              54       53       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[billbrod]

My vote would be to have axis just not be allowed at all. Since the behavior cannot be modified, it feels a bit confusing to set it.