Repeated library terms when using the GeneralizedLibrary feature with weak SINDy

I have been trying to reproduce the weak SINDy compressible example by instead creating the function library using the GeneralizedLibrary feature. However, the models produced at the end are not the same and non-unique terms are generated in the feature library. I have included the code to produce this, the first section is the same as the notebook.

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp
from sklearn.linear_model import Lasso
from scipy.io import loadmat
from sklearn.metrics import mean_squared_error
from pysindy.utils.odes import lorenz

import pysindy as ps

# Ignore matplotlib deprecation warnings
import warnings

"""
Solve the 2D system of equations
u_t = -(u.d)u + (-dP + mu d^2 u)/rho
using the ideal gas equation for compressibility.
"""

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

# 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)

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

t = np.linspace(0, T, Nt)
t.shape
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)
patiotemporal_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, :]
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,
)

I then define the feature and function libraries in two different forms, one of which is the same as the Jupyter-notebook, and the other using GeneralizedLibrary.

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)

library_functions_linear = [lambda x: x] 
library_functions_inv = [lambda x:  1 / (1e-6 + abs(x))] 
library_names_linear = [lambda x: x]
library_names_inv = [lambda x:  x + "^-1"]

pde_lib = ps.WeakPDELibrary(
    library_functions=library_functions,
    function_names=library_function_names,
    derivative_order=2,
    spatiotemporal_grid=spatiotemporal_grid,
    is_uniform=True,
    K=1000,
)

pde_lib_lin = ps.WeakPDELibrary(
    library_functions=library_functions_linear,
    function_names=library_names_linear,
    derivative_order=2,
    spatiotemporal_grid=spatiotemporal_grid,
    is_uniform=True,
    K=1000,
)

pde_lib_inv = ps.WeakPDELibrary(
    library_functions=library_functions_inv,
    function_names=library_names_inv,
    derivative_order=2,
    spatiotemporal_grid=spatiotemporal_grid,
    is_uniform=True,
    K=1000,
)

generalized_lib = ps.GeneralizedLibrary(
                    [pde_lib_lin, pde_lib_inv])

I then fit the libraries separately.

optimizer = ps.STLSQ(threshold=5e-1, alpha=1e-12, normalize_columns=False)
model = ps.SINDy(
    feature_library=pde_lib, optimizer=optimizer, feature_names=["u", "v", "p"]
)
model.fit(u_shaped_noiseless, t=t, unbias=False)
print("Noiseless weak fit:")
model.print()
#print(model.get_feature_names())
print(len(model.get_feature_names()))
print("")
print("Generalized lib fit")
optimizer = ps.STLSQ(threshold=5e-1, alpha=1e-12, normalize_columns=False)
model = ps.SINDy(
    feature_library=generalized_lib, optimizer=optimizer, feature_names=["u", "v", "p"]
)
model.fit(u_shaped_noiseless, t=t, unbias=False)
model.print()
print(model.get_feature_names())
print(len(model.get_feature_names()))

which gives the resulting models as

Noiseless weak fit:
(u)' = -0.995 vu_2 + 1.004 p^-1u_22 + -0.999 uu_1 + -1.002 p^-1p_1 + 1.000 p^-1u_11
(v)' = -0.995 vv_2 + -1.002 p^-1p_2 + 1.002 p^-1v_22 + -0.999 uv_1 + 1.003 p^-1v_11
(p)' = -1.000 pv_2 + -1.000 vp_2 + -1.000 pu_1 + -1.000 up_1

Generalized lib fit
(u)' = 1.232 v_2 + 0.821 u_22 + -2.610 v_1 + -0.802 p_1 + 0.906 u_11 + -0.961 vu_2 + -0.964 pv_2 + -0.933 uu_1 + 2.229 pv_1
(v)' = -5.276 u_2 + -1.026 p_2 + 1.003 v_22 + 1.029 v_11 + 5.308 pu_2 + -0.954 vv_2 + -0.916 uv_1
(p)' = -1.000 pv_2 + -1.000 vp_2 + -1.000 pu_1 + -1.000 up_1

So the models are not the same. If I print the feature names using model.get_feature_names(), we can see that the second fit contains 15 more terms. These additional terms are actually terms which have already been included in the library, so are correlated with themselves. The terms which appear twice in the library are

['u_2', 'v_2', 'p_2', 'u_22', 'v_22', 'p_22', 'u_1', 'v_1', 'p_1', 'u_12', 'v_12', 'p_12', 'u_11', 'v_11', 'p_11']

Have I made a mistake forming the generalized library?

dynamicslab / pysindy

Repeated library terms when using the GeneralizedLibrary feature with weak SINDy #395