Unreasonable and inconsistent results when running the "2d compressible isothermal flow equations" example in weak formulation SINDy

[nnnyyc]

When I run the example of "2d compressible isothermal flow equations" in weak formulation SINDy, I get a seriously inconsistent results with the document and the true equations even when there is no noise. Here is my result:

Noiseless differential fit:
(u)' = -1.000 uv_2 + 0.998 vp_22 + -1.002 uu_1 + -1.002 p^-1p_1 + 1.000 vp_11
(v)' = -1.002 pv_2 + -1.000 p^-1p_2 + 1.000 u^-1p_22 + -0.997 pu_1 + 0.998 u^-1p_11
(p)' = -0.997 pp_2 + -0.999 v^-1v_2 + -1.000 up_1 + -0.998 v^-1u_1
R2 score:  0.9999402527137371
Noisy differential fit:
(u)' = -1.971 u + -0.899 uv_2 + 0.837 u^-1p_2 + -0.593 uu_1
(v)' = -1.624 v + -4.063 vp_2 + -0.838 pv_2
(p)' = -0.995 pp_2 + -0.997 v^-1v_2 + -0.997 up_1 + -0.987 v^-1u_1
R2 score:  0.8679715519283352
Noiseless weak fit:
(u)' = 0.978 vp_22 + -0.951 uu_1 + -0.994 p^-1p_1 + 0.992 vp_11
(v)' = -0.500 pv_2 + -0.501 p^-1p_2 + 1.002 u^-1p_22 + -0.999 pu_1 + -0.500 pv_12 + -0.501 p^-1p_12 + 1.001 u^-1p_11
(p)' = -1.000 pp_2 + -1.000 v^-1v_2 + -1.000 up_1 + -1.000 v^-1u_1
R2 score:  0.9962515186348716
Noisy weak fit:
(u)' = 0.978 vp_22 + -0.950 uu_1 + -0.993 p^-1p_1 + 0.992 vp_11
(v)' = -41.532 u_1 + 66.057 v_11 + 1.000 u^-1p_22 + 19.705 up_1 + 21.366 vp_1 + -0.604 pu_1 + -30.321 pp_11 + -33.725 u^-1p_11
(p)' = -0.516 pp_2 + -0.883 up_1 + -1.022 v^-1u_1 + -0.516 pp_12
R2 score:  0.9562401139683946

However, the derivative p' should not include v^-1, and the result in the document is similar to (p)' = -0.997 pv_2 + -0.999 vp_2 + -1.000 pu_1 + -0.998 up_1. According to the true equation, the real result should be similar to (p)' = -1.000 pv_2 + -1.000 vp_2 + -1.000 pu_1 + -1.000 up_1.

Here is my code, which is almost consistent with document, except the function_library is commented and the inputs library_functions and function_names are used when building weak_pde_lib( refer to #478 ). Besides, my python version is 3.8.10 and numpy version is 1.24.4.

At first, we ipmort necessary package and define the function used in solve_ivp.

import numpy as np
from scipy.integrate import solve_ivp

import pysindy as ps

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

# integration keywords for solve_ivp, typically needed for chaotic systems
integrator_keywords = {}
integrator_keywords["rtol"] = 1e-12
integrator_keywords["method"] = "LSODA"
integrator_keywords["atol"] = 1e-12


def compressible(t, U, dx, N, mu, RT):
    u = U.reshape(N, N, 3)[:, :, 0]
    v = U.reshape(N, N, 3)[:, :, 1]
    rho = U.reshape(N, N, 3)[:, :, 2]
    ux = ps.differentiation.FiniteDifference(
        d=1,
        axis=0,
        periodic=True,
    )._differentiate(u, dx)
    uy = ps.differentiation.FiniteDifference(
        d=1,
        axis=1,
        periodic=True,
    )._differentiate(u, dx)
    uxx = ps.differentiation.FiniteDifference(
        d=2,
        axis=0,
        periodic=True,
    )._differentiate(u, dx)
    uyy = ps.differentiation.FiniteDifference(
        d=2,
        axis=1,
        periodic=True,
    )._differentiate(u, dx)
    vx = ps.differentiation.FiniteDifference(
        d=1,
        axis=0,
        periodic=True,
    )._differentiate(v, dx)
    vy = ps.differentiation.FiniteDifference(
        d=1,
        axis=1,
        periodic=True,
    )._differentiate(v, dx)
    vxx = ps.differentiation.FiniteDifference(
        d=2,
        axis=0,
        periodic=True,
    )._differentiate(v, dx)
    vyy = ps.differentiation.FiniteDifference(
        d=2,
        axis=1,
        periodic=True,
    )._differentiate(v, dx)
    px = ps.differentiation.FiniteDifference(
        d=1,
        axis=0,
        periodic=True,
    )._differentiate(rho * RT, dx)
    py = ps.differentiation.FiniteDifference(
        d=1,
        axis=1,
        periodic=True,
    )._differentiate(rho * RT, dx)
    ret = np.zeros((N, N, 3))
    ret[:, :, 0] = -(u * ux + v * uy) - (px - mu * (uxx + uyy)) / rho
    ret[:, :, 1] = -(u * vx + v * vy) - (py - mu * (vxx + vyy)) / rho
    ret[:, :, 2] = -(u * px / RT + v * py / RT + rho * ux + rho * vy)
    return ret.reshape(3 * N * N)

Then generate the data and add noise。

N = 64
Nt = 100
L = 5
T = 0.5
mu = 1
RT = 1

t = np.linspace(0, T, Nt)
x = np.arange(0, N) * L / N
y = np.arange(0, N) * L / N
dx = x[1] - x[0]

# some arbitrary initial conditions
y0 = np.zeros((N, N, 3))
y0[:, :, 0] = (
    -np.sin(2 * np.pi / L * x)[:, np.newaxis]
    + 0.5 * np.cos(2 * 2 * np.pi / L * y)[np.newaxis, :]
)
y0[:, :, 1] = (
    0.5 * np.cos(2 * np.pi / L * x)[:, np.newaxis]
    - np.sin(2 * 2 * np.pi / L * y)[np.newaxis, :]
)
y0[:, :, 2] = (
    1
    + 0.5
    * np.cos(2 * np.pi / L * x)[:, np.newaxis]
    * np.cos(2 * 2 * np.pi / L * y)[np.newaxis, :]
)

sol = solve_ivp(
    compressible,
    (t[0], t[-1]),
    y0=y0.reshape(3 * N * N),
    t_eval=t,
    args=(dx, N, mu, RT),
    method="RK45",
    rtol=1e-8,
    atol=1e-8,
)

u_shaped_noiseless = sol.y.reshape(N, N, 3, -1).transpose(0, 1, 3, 2)
u_dot_noiseless = ps.FiniteDifference(d=1, axis=2)._differentiate(u_shaped_noiseless, t)

ep = 5e-3
np.random.seed(100)
u_shaped_noisy = u_shaped_noiseless + 2 * ep * (
    0.5 - np.random.random(size=(N, N, Nt, 3))
)
u_dot_noisy = ps.FiniteDifference(d=1, axis=2)._differentiate(u_shaped_noisy, t)

Finally, run SINDy and weak SINDy, and print the results.

spatial_grid = np.zeros((N, N, 2))
spatial_grid[:, :, 0] = x[:, np.newaxis]
spatial_grid[:, :, 1] = y[np.newaxis, :]

library_functions = [lambda x: x, lambda x: 1 / (1e-6 + abs(x))]
library_function_names = [lambda x: x, lambda x: x + "^-1"]

pde_lib = ps.PDELibrary(
    library_functions=library_functions,
    function_names=library_function_names,
    derivative_order=2,
    spatial_grid=spatial_grid,
)

# Fit the clean data (prints the correct model)
optimizer = ps.STLSQ(threshold=5e-1, alpha=1e-10, normalize_columns=False)
model = ps.SINDy(
    feature_library=pde_lib, optimizer=optimizer, feature_names=["u", "v", "p"]
)
model.fit(u_shaped_noiseless, t=t)
print("Noiseless differential fit:")
model.print()
print("R2 score: ", model.score(u_shaped_noiseless, t=t))

# Fit the noisy data (prints a bad model)
optimizer = ps.STLSQ(threshold=5e-1, alpha=1e-10, normalize_columns=False)
model = ps.SINDy(
    feature_library=pde_lib, optimizer=optimizer, feature_names=["u", "v", "p"]
)
model.fit(u_shaped_noisy, t=t)
print("Noisy differential fit:")
model.print()
print("R2 score: ", model.score(u_shaped_noisy, t=t))

spatiotemporal_grid = np.zeros((N, N, Nt, 3))
spatiotemporal_grid[:, :, :, 0] = x[:, np.newaxis, np.newaxis]
spatiotemporal_grid[:, :, :, 1] = y[np.newaxis, :, np.newaxis]
spatiotemporal_grid[:, :, :, 2] = t[np.newaxis, np.newaxis, :]

library_functions = [lambda x: x, lambda x: 1 / (1e-6 + abs(x))]
library_function_names = [lambda x: x, lambda x: x + "^-1"]
np.random.seed(100)

weak_pde_lib = ps.WeakPDELibrary(
    library_functions=library_functions,
    function_names=library_function_names,
    # function_library=ps.CustomLibrary(library_functions=library_functions,function_names=library_function_names),
    derivative_order=2,
    spatiotemporal_grid=spatiotemporal_grid,
    K=2000,
    H_xt=[L / 10, L / 10, T / 10],
)

# Fit the clean data
optimizer = ps.STLSQ(threshold=5e-1, alpha=1e-12, normalize_columns=False)
model = ps.SINDy(
    feature_library=weak_pde_lib, optimizer=optimizer, feature_names=["u", "v", "p"]
)
model.fit(u_shaped_noiseless, t=t)
print("Noiseless weak fit:")
model.print()
print("R2 score: ", model.score(u_shaped_noiseless, t=t))

# Fit the noisy data
optimizer = ps.STLSQ(threshold=5e-1, alpha=1e-12, normalize_columns=False)
model = ps.SINDy(
    feature_library=weak_pde_lib, optimizer=optimizer, feature_names=["u", "v", "p"]
)
model.fit(u_shaped_noisy, t=t)
print("Noisy weak fit:")
model.print()
print("R2 score: ", model.score(u_shaped_noisy, t=t))

Thank you very much for reading my question, and if there is any important information not mentioned above, please let me know.