Fixed errors with PDE and weak libraries with multiple trajectories a…

…nd ensembling on. Added some tests to test_pysindy.py

pysindy/pysindy.py

@@ -312,7 +312,15 @@ def fit(
         if (ensemble or library_ensemble) and n_models is None:
             n_models = 20
         if ensemble and n_subset is None:
-            n_subset = x.shape[0]
+            if multiple_trajectories:
+                if x[0].ndim == 1:
+                    n_subset = x[0].shape[0]
+                else:
+                    n_subset = x[0].shape[-2]
+            elif x.ndim == 1:
+                n_subset = x.shape[0]
+            else:
+                n_subset = x.shape[-2]
         if (
             ensemble
             and isinstance(self.feature_library, WeakPDELibrary)
@@ -330,6 +338,8 @@ def fit(
         pde_libraries = False
         weak_libraries = False
 
+        if isinstance(self.feature_library, WeakPDELibrary):
+            self.feature_library.old_x = np.copy(x)
         if x_dot is None:
             if isinstance(self.feature_library, WeakPDELibrary):
                 if multiple_trajectories:
@@ -339,7 +349,12 @@ def fit(
                 else:
                     x_dot = convert_u_dot_integral(x, self.feature_library)
             elif isinstance(self.feature_library, PDELibrary):
-                if multiple_trajectories:
+                if multiple_trajectories and isinstance(t, Sequence):
+                    x_dot = [
+                        FiniteDifference(d=1, axis=-2)._differentiate(xi, t=ti)
+                        for xi, ti in zip(x, t)
+                    ]
+                elif multiple_trajectories:
                     x_dot = [
                         FiniteDifference(d=1, axis=-2)._differentiate(xi, t=t)
                         for xi in x
@@ -366,7 +381,12 @@ def fit(
                             x, self.feature_library.libraries_[0]
                         )
                 elif pde_libraries:
-                    if multiple_trajectories:
+                    if multiple_trajectories and isinstance(t, Sequence):
+                        x_dot = [
+                            FiniteDifference(d=1, axis=-2)._differentiate(xi, t=ti)
+                            for xi, ti in zip(x, t)
+                        ]
+                    elif multiple_trajectories:
                         x_dot = [
                             FiniteDifference(d=1, axis=-2)._differentiate(xi, t=t)
                             for xi in x
@@ -375,20 +395,16 @@ def fit(
                         x_dot = FiniteDifference(d=1, axis=-2)._differentiate(x, t=t)
 
         if multiple_trajectories:
-            if isinstance(self.feature_library, PDELibrary):
-                self.feature_library.num_trajectories = len(x)
-            if isinstance(self.feature_library, WeakPDELibrary):
-                self.feature_library.num_trajectories = len(x)
+            self.feature_library.num_trajectories = len(x)
             if isinstance(self.feature_library, GeneralizedLibrary):
                 for lib in self.feature_library.libraries_:
                     if isinstance(lib, PDELibrary):
                         lib.num_trajectories = len(x)
                     if isinstance(lib, WeakPDELibrary):
                         lib.num_trajectories = len(x)
-
             x, x_dot = self._process_multiple_trajectories(x, t, x_dot)
-
         else:
+            self.feature_library.num_trajectories = 1
             x = validate_input(x, t)
 
             if self.discrete_time:
@@ -442,7 +458,6 @@ def fit(
                 # save the grid
                 if pde_library_flag == "WeakPDE":
                     old_spatiotemporal_grid = self.feature_library.spatiotemporal_grid
-
                 for i in range(n_models):
                     x_ensemble, x_dot_ensemble = drop_random_rows(
                         x,
@@ -451,6 +466,7 @@ def fit(
                         replace,
                         self.feature_library,
                         pde_library_flag,
+                        multiple_trajectories,
                     )
                     self.model.fit(x_ensemble, x_dot_ensemble)
                     self.coef_list.append(self.model.steps[-1][1].coef_)
@@ -498,6 +514,7 @@ def fit(
                         replace,
                         self.feature_library,
                         pde_library_flag,
+                        multiple_trajectories,
                     )
                     for j in range(n_models):
                         self.feature_library.ensemble_indices = np.sort(
@@ -823,7 +840,6 @@ def _process_multiple_trajectories(self, x, t, x_dot, return_array=True):
         """
         if not isinstance(x, Sequence):
             raise TypeError("Input x must be a list")
-
         if self.discrete_time:
             x = [validate_input(xi) for xi in x]
             if x_dot is None:

pysindy/utils/base.py

@@ -108,7 +108,16 @@ def drop_nan_rows(x, x_dot):
     return x, x_dot
 
 
-def drop_random_rows(x, x_dot, n_subset, replace, feature_library, pde_library_flag):
+def drop_random_rows(
+    x,
+    x_dot,
+    n_subset,
+    replace,
+    feature_library,
+    pde_library_flag,
+    multiple_trajectories,
+):
+    num_trajectories = feature_library.num_trajectories
     # Can't choose random n_subset points if data is from a PDE
     # (and therefore is spatially local).
     # Need to unfold it and just choose n_subset from the temporal slices
@@ -129,28 +138,39 @@ def drop_random_rows(x, x_dot, n_subset, replace, feature_library, pde_library_f
             spatial_grid = feature_library.spatial_grid
             dims = spatial_grid.shape[:-1]
             if len(dims) > 0:
-                num_time = n_samples // np.product(dims)
+                num_time = n_samples // np.product(dims) // num_trajectories
             else:
-                num_time = n_samples
+                num_time = n_samples // num_trajectories
 
         n_features = x.shape[1]
         if n_subset > num_time:
             n_subset = num_time
         rand_inds = np.sort(choice(range(num_time), n_subset, replace=replace))
 
         if len(dims) > 0:
-            x_shaped = np.reshape(x, np.concatenate([dims, [num_time], [n_features]]))
+            x_shaped = np.reshape(
+                x, np.concatenate([dims, [num_time * num_trajectories], [n_features]])
+            )
         else:
-            x_shaped = np.reshape(x, np.concatenate([[num_time], [n_features]]))
+            x_shaped = np.reshape(
+                x, np.concatenate([[num_time * num_trajectories], [n_features]])
+            )
         s0 = [slice(dim) for dim in x_shaped.shape]
-        s0[len(dims)] = rand_inds
+
+        rand_inds_total = []
+        for i in range(num_trajectories):
+            rand_inds_total.append(rand_inds + num_time * i)
+        s0[len(dims)] = rand_inds_total
 
         if len(dims) > 0:
             x_new = np.reshape(
-                x_shaped[tuple(s0)], (np.product(dims) * n_subset, x.shape[1])
+                x_shaped[tuple(s0)],
+                (np.product(dims) * n_subset * num_trajectories, x.shape[1]),
             )
         else:
-            x_new = np.reshape(x_shaped[tuple(s0)], (n_subset, x.shape[1]))
+            x_new = np.reshape(
+                x_shaped[tuple(s0)], (n_subset * num_trajectories, x.shape[1])
+            )
 
         if pde_library_flag == "WeakPDE":
             spatiotemporal_grid = feature_library.spatiotemporal_grid
@@ -159,14 +179,25 @@ def drop_random_rows(x, x_dot, n_subset, replace, feature_library, pde_library_f
             new_spatiotemporal_grid = spatiotemporal_grid[tuple(s1)]
             feature_library.spatiotemporal_grid = new_spatiotemporal_grid
             feature_library._set_up_grids()
-            x_dot_new = convert_u_dot_integral(x_shaped[tuple(s0)], feature_library)
-
+            s0[len(dims)] = rand_inds
+            if multiple_trajectories:
+                x_dot_new = [
+                    convert_u_dot_integral(xi[tuple(s0)], feature_library)
+                    for xi in feature_library.old_x
+                ]
+                x_dot_new = np.vstack(x_dot_new)
+            else:
+                x_dot_new = convert_u_dot_integral(
+                    feature_library.old_x[tuple(s0)], feature_library
+                )
         else:
             x_dot_shaped = np.reshape(
-                x_dot, np.concatenate([dims, [num_time], [n_features]])
+                x_dot,
+                np.concatenate([dims, [num_time * num_trajectories], [n_features]]),
             )
             x_dot_new = np.reshape(
-                x_dot_shaped[tuple(s0)], (np.product(dims) * n_subset, x.shape[1])
+                x_dot_shaped[tuple(s0)],
+                (np.product(dims) * n_subset * num_trajectories, x.shape[1]),
             )
     else:
         # Choose random n_subset points to use

test/conftest.py

@@ -6,6 +6,7 @@
 from scipy.integrate import solve_ivp
 
 from pysindy.differentiation import FiniteDifference
+from pysindy.differentiation import SpectralDerivative
 from pysindy.feature_library import CustomLibrary
 from pysindy.feature_library import FourierLibrary
 from pysindy.feature_library import GeneralizedLibrary
@@ -63,6 +64,42 @@ def data_multiple_trajctories():
     return x_list, t_list
 
 
+@pytest.fixture
+def diffuse_multiple_trajectories():
+    def diffuse(t, u, dx, nx):
+        u = np.reshape(u, nx)
+        du = SpectralDerivative(d=2, axis=0)._differentiate(u, dx)
+        return np.reshape(u * du, nx)
+
+    # Required for accurate solve_ivp results
+    integrator_keywords = {}
+    integrator_keywords["rtol"] = 1e-12
+    integrator_keywords["method"] = "LSODA"
+    integrator_keywords["atol"] = 1e-12
+    N = 200
+    h0 = 1.0
+    L = 5
+    T = 1
+    t = np.linspace(0, T, N)
+    x = np.arange(0, N) * L / N
+    ep = 0.5 * h0
+    y0 = np.reshape(
+        h0 + ep * np.cos(4 * np.pi / L * x) + ep * np.cos(2 * np.pi / L * x), N
+    )
+    y1 = np.reshape(
+        h0 + ep * np.cos(4 * np.pi / L * x) - ep * np.cos(2 * np.pi / L * x), N
+    )
+    dx = x[1] - x[0]
+    sol1 = solve_ivp(
+        diffuse, (t[0], t[-1]), y0=y0, t_eval=t, args=(dx, N), **integrator_keywords
+    )
+    sol2 = solve_ivp(
+        diffuse, (t[0], t[-1]), y0=y1, t_eval=t, args=(dx, N), **integrator_keywords
+    )
+    u = [np.reshape(sol1.y, (N, N, 1)), np.reshape(sol2.y, (N, N, 1))]
+    return t, x, u
+
+
 @pytest.fixture
 def data_discrete_time():
 

test/test_pysindy.py

@@ -658,7 +658,6 @@ def test_ensemble_pdes(data_1d_random_pde):
     model = SINDy(feature_library=pde_lib).fit(
         u, t, ensemble=True, n_models=10, n_subset=len(t) // 2
     )
-    print(np.shape(model.coef_list))
     assert len(model.coef_list) == 10
     model = SINDy(feature_library=pde_lib).fit(
         u, x_dot=u_dot, ensemble=True, n_models=10, n_subset=len(t) // 2
@@ -869,3 +868,62 @@ def test_data_shapes():
     model.fit(x)
     x = np.ones((n, n, n, n, 2))
     model.fit(x)
+
+
+def test_multiple_trajectories_and_ensemble(diffuse_multiple_trajectories):
+    t, x, u = diffuse_multiple_trajectories
+    library_functions = [lambda x: x]
+    library_function_names = [lambda x: x]
+
+    pde_lib = PDELibrary(
+        library_functions=library_functions,
+        function_names=library_function_names,
+        derivative_order=2,
+        spatial_grid=x,
+        is_uniform=True,
+    )
+
+    X, T = np.meshgrid(x, t, indexing="ij")
+    XT = np.transpose([X, T], [1, 2, 0])
+
+    weak_lib = WeakPDELibrary(
+        library_functions=library_functions,
+        function_names=library_function_names,
+        derivative_order=2,
+        spatiotemporal_grid=XT,
+        K=100,
+        is_uniform=False,
+        num_pts_per_domain=30,
+    )
+
+    optimizer = STLSQ(threshold=0.1, alpha=1e-5, normalize_columns=False)
+    model = SINDy(feature_library=pde_lib, optimizer=optimizer, feature_names=["u"])
+    model.fit(u, multiple_trajectories=True, t=t, ensemble=False)
+    print(model.coefficients(), model.coefficients()[0][-1])
+    assert abs(model.coefficients()[0][-1] - 1) < 1e-2
+    assert np.all(model.coefficients()[0][:-1] == 0)
+
+    model = SINDy(feature_library=weak_lib, optimizer=optimizer, feature_names=["u"])
+    model.fit(u, multiple_trajectories=True, t=t, ensemble=False)
+    assert abs(model.coefficients()[0][-1] - 1) < 1e-2
+    assert np.all(model.coefficients()[0][:-1] == 0)
+
+    model = SINDy(feature_library=pde_lib, optimizer=optimizer, feature_names=["u"])
+    model.fit(u, multiple_trajectories=True, t=t, ensemble=True, n_subset=len(t))
+    assert abs(model.coefficients()[0][-1] - 1) < 1e-2
+    assert np.all(model.coefficients()[0][:-1] == 0)
+
+    model = SINDy(feature_library=weak_lib, optimizer=optimizer, feature_names=["u"])
+    model.fit(u, multiple_trajectories=True, t=t, ensemble=True, n_subset=len(t))
+    assert abs(model.coefficients()[0][-1] - 1) < 1e-2
+    assert np.all(model.coefficients()[0][:-1] == 0)
+
+    model = SINDy(feature_library=pde_lib, optimizer=optimizer, feature_names=["u"])
+    model.fit(u, multiple_trajectories=True, t=t, ensemble=True)
+    assert abs(model.coefficients()[0][-1] - 1) < 1e-2
+    assert np.all(model.coefficients()[0][:-1] == 0)
+
+    model = SINDy(feature_library=weak_lib, optimizer=optimizer, feature_names=["u"])
+    model.fit(u, multiple_trajectories=True, t=t, ensemble=True)
+    assert abs(model.coefficients()[0][-1] - 1) < 1e-2
+    assert np.all(model.coefficients()[0][:-1] == 0)