diff --git a/src/pykoopman/koopman.py b/src/pykoopman/koopman.py index 90c4ac6..b797090 100644 --- a/src/pykoopman/koopman.py +++ b/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), diff --git a/src/pykoopman/regression/_base_ensemble.py b/src/pykoopman/regression/_base_ensemble.py index 970d478..df099fc 100644 --- a/src/pykoopman/regression/_base_ensemble.py +++ b/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_ diff --git a/test/regression/test_regressors.py b/test/regression/test_regressors.py index c5ac0ee..99ea8d8 100644 --- a/test/regression/test_regressors.py +++ b/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) \ No newline at end of file