fix issues on the pytest version

src/pykoopman/koopman.py

@@ -22,6 +22,7 @@
 from .regression import EDMDc
 from .regression import EnsembleBaseRegressor
 from .regression import PyDMDRegressor
+from .regression import NNDMD
 
 
 class Koopman(BaseEstimator):
@@ -146,13 +147,24 @@ def fit(self, x, y=None, u=None, dt=1):
         if y is None:  # or isinstance(self.regressor, PyDMDRegressor):
             # if there is only 1 trajectory OR regressor is PyDMD
             regressor = self.regressor
+        elif isinstance(self.regressor, NNDMD):
+            regressor = self.regressor
         else:
-            # multiple trajectories
+            # multiple 1-step-trajectories
             regressor = EnsembleBaseRegressor(
                 regressor=self.regressor,
                 func=self.observables.transform,
                 inverse_func=self.observables.inverse,
             )
+            # if x is a list, we need to further change trajectories into 1-step-traj
+            if isinstance(x, list):
+                x_tmp = []
+                y_tmp = []
+                for traj_dat in x:
+                    x_tmp.append(traj_dat[:-1])
+                    y_tmp.append(traj_dat[1:])
+                x = np.hstack(x_tmp)
+                y = np.hstack(y_tmp)
 
         steps = [
             ("observables", self.observables),

src/pykoopman/regression/_base_ensemble.py

@@ -4,6 +4,7 @@
 """
 from __future__ import annotations
 
+import numpy as np
 from sklearn.base import BaseEstimator
 from sklearn.base import clone
 from sklearn.base import TransformerMixin
@@ -77,10 +78,14 @@ def fit(self, X, y, **fit_params):
             functions.
         """
 
-        # store the number of dimension of the target to predict an array of
-        # similar shape at predict
+        # if (
+        #         isinstance(X, np.ndarray)
+        #         and isinstance(y, np.ndarray)
+        #         and X.ndim == 2
+        #         and y.ndim == 2
+        # ):
+        # case 2: x, y are 2D np.ndarray, must be 1-step, no validation
         self._training_dim = y.ndim
-
         # transformers are designed to modify X which is 2d dimensional, we
         # need to modify y accordingly.
         if y.ndim == 1:
@@ -99,12 +104,17 @@ def fit(self, X, y, **fit_params):
 
         if self.regressor is None:
             from sklearn.linear_model import LinearRegression
-
             self.regressor_ = LinearRegression()
         else:
             self.regressor_ = clone(self.regressor)
 
         self.regressor_.fit(X, y_trans, **fit_params)
+        # elif isinstance(X, list) and isinstance(y, list):
+        #     # case 4: x, y are two lists of trajectories, we have validation data
+        #     for
+        #
+        # else:
+        #     raise ValueError("check `x` and `y` for `self.fit`")
 
         if hasattr(self.regressor_, "feature_names_in_"):
             self.feature_names_in_ = self.regressor_.feature_names_in_

test/regression/test_regressors.py

@@ -10,6 +10,7 @@
 from pykoopman.regression import NNDMD
 from pykoopman.regression import PyDMDRegressor
 from sklearn.gaussian_process.kernels import RBF
+import pykoopman as pk
 
 
 class RegressorWithoutFit:
@@ -70,16 +71,26 @@ def test_fit_regressors(data_xy, regressor):
     [
         # case 1,2 only work for pykoopman class
         # case 1: single step single traj, no validation
-        (np.random.rand(200, 3), None),
+        (
+                np.random.rand(200, 3),
+                None
+        ),
         # case 2: single step multiple traj, no validation
-        (np.random.rand(200, 3), np.random.rand(200, 3)),
+        (
+                np.random.rand(200, 3),
+                np.random.rand(200, 3) # because "x" is not a list, so we think this
+                                       # is single step
+        ),
         # case 3,4 works for regressor directly
         # case 3: multiple traj, no validation
-        ([np.random.rand(200, 3), np.random.rand(100, 3)], None),
+        (
+                [np.random.rand(200, 3), np.random.rand(100, 3)],  # this is training
+                None  # no validation
+        ),
         # case 4: multiple traj, with validation
         (
-            [np.random.rand(100, 3), np.random.rand(100, 3)],
-            [np.random.rand(300, 3), np.random.rand(400, 3)],
+            [np.random.rand(100, 3), np.random.rand(100, 3)], # this is training
+            [np.random.rand(300, 3), np.random.rand(400, 3)], # this is validation
         ),
     ],
 )
@@ -107,3 +118,56 @@ def test_fit_nndmd_regressor(data_xy, regressor):
     """test if using nndmd regressor alone will run the fit without error"""
     x, y = data_xy
     regressor.fit(x, y)
+
+@pytest.mark.parametrize(
+    "data_xy",
+    [
+        # # case 1,2 only work for pykoopman class
+        # # case 1: single step single traj, no validation
+        # (
+        #         np.random.rand(200, 3),
+        #         None
+        # ),
+        # # case 2: single step multiple traj, no validation
+        # (
+        #         np.random.rand(200, 3),
+        #         np.random.rand(200, 3) # because "x" is not a list, so we think this
+        #                                # is single step
+        # ),
+        # # case 3,4 works for regressor directly
+        # # case 3: multiple traj, no validation
+        # (
+        #         [np.random.rand(200, 3), np.random.rand(100, 3)],  # this is training
+        #         None  # no validation
+        # ),
+        # case 4: multiple traj, with validation
+        (
+            [np.random.rand(100, 3), np.random.rand(100, 3)], # this is training
+            [np.random.rand(300, 3), np.random.rand(400, 3)], # this is validation
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "regressor",
+    [
+        NNDMD(
+            mode="Dissipative",
+            look_forward=2,
+            config_encoder=dict(
+                input_size=3, hidden_sizes=[32] * 2, output_size=4, activations="swish"
+            ),
+            config_decoder=dict(
+                input_size=4, hidden_sizes=[32] * 2, output_size=3, activations="linear"
+            ),
+            batch_size=512,
+            lbfgs=True,
+            normalize=False,
+            normalize_mode="max",
+            trainer_kwargs=dict(max_epochs=1),
+        )
+    ],
+)
+def test_fit_dlkoopman(data_xy, regressor):
+    """test if using NNDMD regressor work inside pykoopman"""
+    model_d = pk.Koopman(regressor=regressor)
+    model_d.fit(data_xy[0],data_xy[1], dt=1)