added handling of 3d numpy arrays and lists of uneven numpy arrays

README.rst

@@ -204,7 +204,7 @@ you can run the following to automatically reformat your staged code
 
 .. code-block:: bash
 
-    pre-commit -a -v
+    pre-commit run -a -v
 
 Note that you will then need to re-stage any changes ``pre-commit`` made to your code.
 

src/pykoopman/common/validation.py

@@ -40,11 +40,27 @@ def validate_input(x, t=T_DEFAULT):
 
 def check_array(x, **kwargs):
     if np.iscomplexobj(x):
-        return skl_check_array(x.real, **kwargs) + 1j * skl_check_array(
-            x.imag, **kwargs
-        )
+        if x.ndim == 3:
+            # Handle 3D arrays by processing each 2D slice
+            result = np.zeros_like(x)
+            for i in range(x.shape[0]):
+                result[i] = skl_check_array(x[i].real, **kwargs) + 1j * skl_check_array(
+                    x[i].imag, **kwargs
+                )
+            return result
+        else:
+            return skl_check_array(x.real, **kwargs) + 1j * skl_check_array(
+                x.imag, **kwargs
+            )
     else:
-        return skl_check_array(x, **kwargs)
+        if x.ndim == 3:
+            # Handle 3D arrays by processing each 2D slice
+            result = np.zeros_like(x)
+            for i in range(x.shape[0]):
+                result[i] = skl_check_array(x[i], **kwargs)
+            return result
+        else:
+            return skl_check_array(x, **kwargs)
 
 
 def drop_nan_rows(arr, *args):

src/pykoopman/differentiation/_derivative.py

@@ -62,7 +62,7 @@ def get_params(self, deep=True):
 
         return params
 
-    def __call__(self, x, t):
+    def __call__(self, x, t, axis=0):
         """
         Perform numerical differentiation by calling the ``dxdt`` method.
 
@@ -86,4 +86,4 @@ def __call__(self, x, t):
                 raise ValueError("t must be a positive constant or an array")
             t = arange(x.shape[0]) * t
 
-        return dxdt(x, t, axis=0, **self.kwargs)
+        return dxdt(x, t, axis=axis, **self.kwargs)

src/pykoopman/koopman.py

@@ -21,6 +21,7 @@
 from .regression import DMDc
 from .regression import EDMDc
 from .regression import EnsembleBaseRegressor
+from .regression import HAVOK
 from .regression import NNDMD
 from .regression import PyDMDRegressor
 
@@ -146,9 +147,23 @@ 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
+            y_flag = True
+            # regressor = self.regressor
+            x, y = self._detect_reshape(x, offset=True)
+            if isinstance(self.regressor, HAVOK):
+                regressor = self.regressor
+                y_flag = False
+            else:
+                regressor = EnsembleBaseRegressor(
+                    regressor=self.regressor,
+                    func=self.observables.transform,
+                    inverse_func=self.observables.inverse,
+                )
+            # regressor = self.regressor
         elif isinstance(self.regressor, NNDMD):
             regressor = self.regressor
+            y_flag = False
+
         else:
             # multiple 1-step-trajectories
             regressor = EnsembleBaseRegressor(
@@ -157,14 +172,17 @@ def fit(self, x, y=None, u=None, dt=1):
                 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)
+            x, _ = self._detect_reshape(x, offset=False)
+            y, _ = self._detect_reshape(y, offset=False)
+            y_flag = False
+            # 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),
@@ -199,7 +217,7 @@ def fit(self, x, y=None, u=None, dt=1):
         # compute amplitudes
         if isinstance(x, list):
             self._amplitudes = None
-        elif y is None:
+        elif y_flag:
             if hasattr(self.observables, "n_consumed_samples"):
                 # g0 = self.observables.transform(
                 #     x[0 : 1 + self.observables.n_consumed_samples]
@@ -239,6 +257,8 @@ def predict(self, x, u=None):
                 Predicted state one timestep in the future.
         """
 
+        x = validate_input(x)
+
         check_is_fitted(self, "n_output_features_")
         x_next = self.observables.inverse(self._step(x, u))
         return x_next
@@ -276,7 +296,6 @@ def simulate(self, x0, u=None, n_steps=1):
             x0 = x0.T
         else:
             raise TypeError("Check your initial condition shape!")
-
         # x = empty((n_steps, self.n_input_features_), dtype=self.A.dtype)
         y = empty((n_steps, self.A.shape[0]), dtype=self.W.dtype)
 
@@ -428,7 +447,6 @@ def C(self):
         # if not isinstance(self.observables, RadialBasisFunction):
         #     raise ValueError("this type of self.observable has no C")
         # return self._pipeline.steps[0][1].measurement_matrix_
-
         measure_mat = self._pipeline.steps[0][1].measurement_matrix_
         ur = self._pipeline.steps[-1][1].ur
         C = measure_mat @ ur
@@ -571,3 +589,49 @@ def _regressor(self):
         # the _regressor.fit to update the model coefficients.
         # call this function with _regressor()
         return self._pipeline.steps[1][1]
+
+    def _detect_reshape(self, X, offset=True):
+        """
+        Detect the shape of the input data and reshape it accordingly to return
+        both X and Y in the correct shape.
+        """
+        s1 = -1 if offset else None
+        s2 = 1 if offset else None
+        if isinstance(X, np.ndarray):
+            if X.ndim == 1:
+                X = X.reshape(-1, 1)
+
+            if X.ndim == 2:
+                self.n_samples_, self.n_input_features_ = X.shape
+                self.n_trials_ = 1
+                return X[:s1], X[s2:]
+            elif X.ndim == 3:
+                self.n_trials_, self.n_samples_, self.n_input_features_ = X.shape
+                X, Y = X[:, :s1, :], X[:, s2:, :]
+                return X.reshape(-1, X.shape[2]), Y.reshape(
+                    -1, Y.shape[2]
+                )  # time*trials, features
+
+        elif isinstance(X, list):
+            assert all(isinstance(x, np.ndarray) for x in X)
+            self.n_trials_tot, self.n_samples_tot, self.n_input_features_tot = (
+                [],
+                [],
+                [],
+            )
+            X_tot, Y_tot = [], []
+            for x in X:
+                x, y = self._detect_reshape(x)
+                X_tot.append(x)
+                Y_tot.append(y)
+                self.n_trials_tot.append(self.n_trials_)
+                self.n_samples_tot.append(self.n_samples_)
+                self.n_input_features_tot.append(self.n_input_features_)
+            X = np.concatenate(X_tot, axis=0)
+            Y = np.concatenate(Y_tot, axis=0)
+
+            self.n_trials_ = sum(self.n_trials_tot)
+            self.n_samples_ = sum(self.n_samples_tot)
+            self.n_input_features_ = sum(self.n_input_features_tot)
+
+            return X, Y

src/pykoopman/observables/_base.py

@@ -79,9 +79,15 @@ def inverse(self, y):
             ValueError: If the shape of the input does not match the expected shape.
         """
         check_is_fitted(self, ["n_input_features_", "measurement_matrix_"])
+        if isinstance(y, list):
+            return [self.inverse(y_trial) for y_trial in y]
+        if y.ndim == 3:
+            return np.array([self.inverse(y_trial) for y_trial in y])
+
         if y.ndim == 1:
             y = y.reshape(1, -1)
-        if y.shape[1] != self.n_output_features_:
+
+        if y.shape[-1] != self.n_output_features_:
             raise ValueError(
                 "Wrong number of input features."
                 f"Expected y.shape[1] = {self.n_output_features_}; "

src/pykoopman/observables/_custom_observables.py

@@ -139,6 +139,12 @@ def transform(self, x):
         check_is_fitted(self, "n_output_features_")
         x = validate_input(x)
 
+        if isinstance(x, list):
+            return [self.transform(x_trial) for x_trial in x]
+
+        if x.ndim == 3:
+            return np.array([self.transform(x_trial) for x_trial in x])
+
         n_samples, n_features = x.shape
 
         if n_features != self.n_input_features_:

src/pykoopman/observables/_polynomial.py

@@ -141,6 +141,12 @@ def transform(self, x):
             ValueError: If the input data is not valid or the shape of `x` does not
             match training shape.
         """
+        if isinstance(x, list):
+            return [self.transform(x_trial) for x_trial in x]
+
+        if x.ndim == 3:
+            return np.array([self.transform(x_trial) for x_trial in x])
+
         check_is_fitted(self, "n_input_features_")
 
         x = check_array(x, order="F", dtype=FLOAT_DTYPES, accept_sparse=("csr", "csc"))

src/pykoopman/observables/_radial_basis_functions.py

@@ -181,6 +181,12 @@ def transform(self, x):
                 input features expected by the transformer.
         """
         check_is_fitted(self, ["n_input_features_", "centers"])
+        if isinstance(x, list):
+            return [self.transform(x_trial) for x_trial in x]
+
+        if x.ndim == 3:
+            return np.array([self.transform(x_trial) for x_trial in x])
+
         x = validate_input(x)
 
         if x.shape[1] != self.n_input_features_:

src/pykoopman/observables/_random_fourier_features.py

@@ -120,6 +120,12 @@ def transform(self, x):
         """
 
         check_is_fitted(self, "n_input_features_")
+        if isinstance(x, list):
+            return [self.transform(x_trial) for x_trial in x]
+
+        if x.ndim == 3:
+            return np.array([self.transform(x_trial) for x_trial in x])
+
         z = np.zeros((x.shape[0], self.n_output_features_))
         z_rff = self._rff_lifting(x)
         if self.include_state:

src/pykoopman/observables/_time_delay.py

@@ -122,11 +122,17 @@ def transform(self, x):
             y (array-like): The transformed data, shape (n_samples - delay * n_delays,
             n_output_features).
         """
-
         check_is_fitted(self, "n_input_features_")
+
+        if isinstance(x, list):
+            return [self.transform(x_trial) for x_trial in x]
+
+        if x.ndim == 3:
+            return np.array([self.transform(x_trial) for x_trial in x])
+
         x = validate_input(x)
 
-        if x.shape[1] != self.n_input_features_:
+        if x.shape[-1] != self.n_input_features_:
             raise ValueError(
                 "Wrong number of input features. "
                 f"Expected x.shape[1] = {self.n_input_features_}; "
@@ -151,7 +157,6 @@ def transform(self, x):
             y[i - self._n_consumed_samples, self.n_input_features_ :] = x[
                 self._delay_inds(i), :
             ].flatten()
-
         return y
 
     def get_feature_names(self, input_features=None):

src/pykoopman/regression/_base.py

@@ -4,6 +4,7 @@
 from abc import ABC
 from abc import abstractmethod
 
+import numpy as np
 from sklearn.base import BaseEstimator
 
 
@@ -54,6 +55,75 @@ def __init__(self, regressor):
             raise AttributeError("regressor does not have a callable predict method")
         self.regressor = regressor
 
+    def _detect_reshape(self, X, offset=True):
+        """
+        Detect the shape of the input data and reshape it accordingly to return
+        both X and Y in the correct shape.
+        """
+        s1 = -1 if offset else None
+        s2 = 1 if offset else None
+        if isinstance(X, np.ndarray):
+            if X.ndim == 1:
+                X = X.reshape(-1, 1)
+
+            if X.ndim == 2:
+                self.n_samples_, self.n_input_features_ = X.shape
+                self.n_trials_ = 1
+                return X[:s1], X[s2:]
+            elif X.ndim == 3:
+                self.n_trials_, self.n_samples_, self.n_input_features_ = X.shape
+                X, Y = X[:, :s1, :], X[:, s2:, :]
+                return X.reshape(-1, X.shape[2]), Y.reshape(
+                    -1, Y.shape[2]
+                )  # time*trials, features
+
+        elif isinstance(X, list):
+            assert all(isinstance(x, np.ndarray) for x in X)
+            self.n_trials_tot, self.n_samples_tot, self.n_input_features_tot = (
+                [],
+                [],
+                [],
+            )
+            X_tot, Y_tot = [], []
+            for x in X:
+                x, y = self._detect_reshape(x)
+                X_tot.append(x)
+                Y_tot.append(y)
+                self.n_trials_tot.append(self.n_trials_)
+                self.n_samples_tot.append(self.n_samples_)
+                self.n_input_features_tot.append(self.n_input_features_)
+            X = np.concatenate(X_tot, axis=0)
+            Y = np.concatenate(Y_tot, axis=0)
+
+            self.n_trials_ = sum(self.n_trials_tot)
+            self.n_samples_ = sum(self.n_samples_tot)
+            self.n_input_features_ = sum(self.n_input_features_tot)
+
+            return X, Y
+
+    def _return_orig_shape(self, X):
+        """
+        X will be a 2d array of shape (n_samples * n_trials, n_features).
+        This function will return the original shape of X.
+        """
+        if not hasattr(self, "n_trials_tot"):
+            X = X.reshape(self.n_trials_, -1, self.n_input_features_)
+            if X.shape[0] == 1:
+                X = X[0]
+            return X
+
+        else:
+            X_tot = []
+            prev_t = 0
+            for i in range(len(self.n_trials_tot)):
+                X_i = X[prev_t : prev_t + self.n_trials_tot[i] * self.n_samples_tot[i]]
+                X_i = X_i.reshape(
+                    self.n_trials_tot[i], -1, self.n_input_features_tot[i]
+                )
+                X_tot.append(X_i)
+                prev_t += self.n_trials_tot[i] * self.n_samples_tot[i]
+            return X_tot
+
     def fit(self, x, y=None):
         raise NotImplementedError
 

src/pykoopman/regression/_base_ensemble.py

@@ -92,8 +92,8 @@ def fit(self, X, y, **fit_params):
         # FIXME: a FunctionTransformer can return a 1D array even when validate
         # is set to True. Therefore, we need to check the number of dimension
         # first.
-        if y_trans.ndim == 2 and y_trans.shape[1] == 1:
-            y_trans = y_trans.squeeze(axis=1)
+        # if y_trans.ndim == 2 and y_trans.shape[1] == 1:
+        #     y_trans = y_trans.squeeze(axis=1)
 
         if self.regressor is None:
             from sklearn.linear_model import LinearRegression