[BUG] Error in finite difference coefficient calculation

[NolanBrb]

When trying to test PySINDy with basic examples, I run into an error linked with the use of AxesArrays. It seems that during coefficient calculation, in the process of calling the transpose function, the code tries to get an item from the array with a non-valid key.

Transpose doesn't work for AxesArray

import pysindy as ps
import numpy as np

x = np.linspace(-1,1,10)**2
t = np.linspace(0,1,10)

model = ps.SINDy()
model.fit(x, t=t)
model.print()

Error message:

---

KeyError Traceback (most recent call last)
Cell In[11], line 8
5 t = np.linspace(0,1,10)
7 model = ps.SINDy()
----> 8 model.fit(x, t=t)
9 model.print()

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\pysindy.py:246, in SINDy.fit(self, x, t, x_dot, u)
240 u = validate_control_variables(
241 x,
242 u,
243 trim_last_point=(self.discrete_time and x_dot is None),
244 )
245 self.n_control_features_ = u[0].shape[u[0].ax_coord]
--> 246 x, x_dot = self._process_trajectories(x, t, x_dot)
248 # Append control variables
249 if u is not None:

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\pysindy.py:491, in SINDy._process_trajectories(self, x, t, x_dot)
488 x = [xi[:-1] for xi in x]
489 else:
490 x, x_dot = zip(
--> 491 *[
492 self.feature_library.calc_trajectory(
493 self.differentiation_method, xi, ti
494 )
495 for xi, ti in _zip_like_sequence(x, t)
496 ]
497 )
498 return x, x_dot

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\pysindy.py:492, in (.0)
488 x = [xi[:-1] for xi in x]
489 else:
490 x, x_dot = zip(
491 *[
--> 492 self.feature_library.calc_trajectory(
493 self.differentiation_method, xi, ti
494 )
495 for xi, ti in _zip_like_sequence(x, t)
496 ]
497 )
498 return x, x_dot

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\feature_library\base.py:66, in BaseFeatureLibrary.calc_trajectory(self, diff_method, x, t)
65 def calc_trajectory(self, diff_method, x, t):
---> 66 x_dot = diff_method(x, t=t)
67 x = AxesArray(diff_method.smoothed_x_, x.axes)
68 return x, AxesArray(x_dot, x.axes)

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\differentiation\base.py:53, in BaseDifferentiation.call(self, x, t)
52 def call(self, x, t=1):
---> 53 return self._differentiate(x, t)

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\differentiation\finite_difference.py:270, in FiniteDifference._differentiate(self, x, t)
268 t = AxesArray(np.array(t), axes={"ax_time": 0})
269 coeffs = self._coefficients(t)
--> 270 interior = self._accumulate(coeffs, x)
271 s[self.axis] = slice((self.n_stencil - 1) // 2, -(self.n_stencil - 1) // 2)
272 x_dot[tuple(s)] = interior

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\differentiation\finite_difference.py:230, in FiniteDifference._accumulate(self, coeffs, x)
222 x = AxesArray(x, {"ax_unk": list(range(x.ndim))})
223 x_expanded = AxesArray(
224 np.transpose(x[tuple(s)], axes=trans), x.insert_axis(0, "ax_offset")
225 )
226 return np.transpose(
227 np.einsum(
228 "ij...,ij->j...",
229 x_expanded,
--> 230 np.transpose(coeffs),
231 ),
232 np.roll(np.arange(x.ndim), self.axis),
233 )

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\utils\axes.py:420, in AxesArray.array_function(self, func, types, args, kwargs)
418 if not all(issubclass(t, AxesArray) for t in types):
419 return NotImplemented
--> 420 return HANDLED_FUNCTIONS[func](*args, **kwargs)

File c:\Users\nolan\miniconda3\envs\numerical-physics\Lib\site-packages\pysindy\utils\axes.py:544, in transpose(a, axes)
542 old_reverse = a._ax_map.reverse_map
543 for new_ind, old_ind in enumerate(axes):
--> 544 compat_dict_append(new_axes, old_reverse[old_ind], new_ind)
546 return AxesArray(out, new_axes)

KeyError: 2

PySINDy/Python version information:

1.7.6.dev465+g69e51e8 3.11.2 | packaged by conda-forge | (main, Mar 31 2023, 17:45:57) [MSC v.1934 64 bit (AMD64)]