[BUG] SINDyDerivative does not handle PDEs

[yb6599]

I am trying to differentiate pde data of shape (256, 100, 1) and len(t_train)=100 and I keep getting an error with the message "Invalid shape of X". Here is the simplified version of the code and the corresponding error.

Reproducing code example:

import numpy as np
import pysindy as ps

X = np.random.random(size=(256, 100, 1))
t = np.arange(0, 10, 0.1)
differentiator = ps.SINDyDerivative(kind='kalman')
differentiator._differentiate(X, t)

Error message:

{
    "name": "ValueError",
    "message": "Invalid shape of X.",
    "stack": "---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File d:\\UW\\Research\\gen-experiments\\src\\gen_experiments\\debug.py:30
     28 t = np.arange(0, 10, 0.1)
     29 differentiator = ps.SINDyDerivative(kind='kalman')
---> 30 differentiator._differentiate(X, t)

File d:\\UW\\Research\\gen-experiments\\env\\Lib\\site-packages\\pysindy\\differentiation\\sindy_derivative.py:79, in SINDyDerivative._differentiate(self, x, t)
     74     t = arange(x.shape[0]) * t
     76 differentiator = methods[self.kwargs[\"kind\"]](
     77     **{k: v for k, v in self.kwargs.items() if k != \"kind\"}
     78 )
---> 79 x_dot = differentiator.d(x, t, axis=0)
     80 if self.save_smooth:
     81     self.smoothed_x_ = differentiator.x(x, t, axis=0)

File d:\\UW\\Research\\gen-experiments\\env\\Lib\\site-packages\\derivative\\differentiation.py:194, in Derivative.d(self, X, t, axis)
    176 def d(self, X, t, axis=1):
    177     \"\"\"
    178     Compute the derivative of measurements X taken at times t.
    179 
   (...)
    192         ValueError: Requires that X.shape[axis] equals len(t). If X is flat, requires that len(X) equals len(t).
    193     \"\"\"
--> 194     X, flat = _align_axes(X, t, axis)
    196     if X.shape[1] == 1:
    197         dX = X

File d:\\UW\\Research\\gen-experiments\\env\\Lib\\site-packages\\derivative\\differentiation.py:245, in _align_axes(X, t, axis)
    243         raise ValueError(\"Invalid axis.\")
    244 else:
--> 245     raise ValueError(\"Invalid shape of X.\")
    247 if X.shape[1] != len(t):
    248     raise ValueError(\"Desired X axis size does not match t size.\")

ValueError: Invalid shape of X."
}

PySINDy/Python version information:

1.7.6.dev215+g9c73768 3.11.7 (tags/v3.11.7:fa7a6f2, Dec 4 2023, 19:24:49) [MSC v.1937 64 bit (AMD64)]

[Jacob-Stevens-Haas]

I think this is because BaseDifferentiation subclasses don't all handle the axis kwarg correctly, e.g.

ps.SINDyDerivative(kind="kalman")(X, t, axis=1)
# TypeError: __init__() got an unexpected keyword argument 'axis'

ps.FiniteDifference()(X, t)
# Error, expected, because no way that FiniteDifference knows to differentiate the second axis

ps.SpectralDerivative()(X, t)
# no error, but should, because Spectral is assuming it should differentiate the 0th axis

ps.FiniteDifference(axis=1)(X, t)
# no error, correct

ps.SmoothedFiniteDifference(axis=1)(X, t)
# correct, behaves like FiniteDifference

[yb6599]

This is not an issue on pysindy version 1.7.6.dev278+g098d231, but is an issue on pysindy version 1.7.6.dev315+ga43e217