Use x_dot kwarg in model.fit in PDEs with WeakPDELibrary()

[Jacob-Stevens-Haas]

SINDy.fit() doesn't accept bare numpy array as x_dot.

Thanks to @georgemilosh for finding! SINDy.fit() without an x_dot argument creates it by differentiating x. So internally, it's expecting an AxesArray in order to run concat_sample_axis. However, we also accept users providing x_dot as an argument. When we do, the AxesArray object is never created. This is because _comprehend_and_validate_inputs only adds axes to x and u, not x_dot. Solution is probably just adding x_dot to that function and adding a test for fit called with x_dot.

Discussed in #444

^{Originally posted by georgemilosh December 27, 2023}
So it seems that the PDELibrary does not throw an error, although I have to manually flatten u_dot:

import numpy as np
from scipy.io import loadmat
from sklearn.metrics import mean_squared_error

import pysindy as ps

# Ignore matplotlib deprecation warnings
import warnings

warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)

# Seed the random number generators for reproducibility
np.random.seed(100)

data = loadmat("data/burgers.mat")
time = np.ravel(data["t"])
x = np.ravel(data["x"])
u = np.real(data["usol"])
dt = time[1] - time[0]
dx = x[1] - x[0]
rmse = mean_squared_error(u, np.zeros(u.shape), squared=False)
# add 20% noise (note the impact on derivatives depends on step size...)
u = u + np.random.normal(0, rmse / 5.0, u.shape)


u = np.reshape(u, (len(x), len(time), 1))

# Define weak form PDE library
library_functions = [lambda x: x, lambda x: x * x]
library_function_names = [lambda x: x, lambda x: x + x]

pde_lib = ps.PDELibrary(
    library_functions=library_functions,
    function_names=library_function_names,
    derivative_order=2,
    spatial_grid=x,
    include_bias=False,
    is_uniform=True,
)

# Fit and predict with the non-weak model
opt = ps.SR3(
    threshold=0.05, thresholder="l0", tol=1e-10, normalize_columns=True, max_iter=1000
)
model_for_prediction = ps.SINDy(feature_library=pde_lib, optimizer=opt)
#model_for_prediction.fit(u, t=dt) # default behavior, gives:
## (x0)' = 0.011 x0 + -0.031 x0x0 + -0.001 x0_1 + -0.012 x0x0_1 + -0.028 x0x0x0_1
## I had to add t=dt as an argument to match the result below

#model_for_prediction.fit(u, x_dot=u_dot) # doesn't work!

u_dot = ps.FiniteDifference(axis=1)._differentiate(u, t=dt)
model_for_prediction.fit(u, x_dot=u_dot.flatten(),t =dt) # my guess. Produces 
# (x0)' = 0.108 x0 + -0.313 x0x0 + -0.012 x0_1 + -0.117 x0x0_1 + -0.278 x0x0x0_1 + -0.004 x0x0_11 + 0.005 x0x0x0_11


model_for_prediction.print()

Here flatten u helps. So how can I provide my own x_dot info for PDEs?

My package version:
0.1.dev1618+g694c904.d20231219 3.8.18 | packaged by conda-forge | (default, Oct 10 2023, 15:44:36)
[GCC 12.3.0]