Merge pull request #43 from dynamicslab/fix-predict-simulate-bug

fix bug in prediction and upgrade nndmd

.github/workflows/run-tests.yml

@@ -39,8 +39,8 @@ jobs:
         py.test test
     - name: Generate coverage report
       run: |
-        pip install pytest
-        pip install pytest-cov
+        pip install pytest==7.4.4
+        pip install pytest-cov==4.1.0
         pytest --cov=./ --cov-report=xml
     - name: Upload coverage reports to Codecov
       uses: codecov/codecov-action@v3

.gitignore~

@@ -0,0 +1,39 @@
+build
+dist
+*.egg-info
+# automatically generated by setuptools-scm
+pysindy/version.py
+# sphinx gallery files
+docs/examples
+# virtual environment
+venv
+
+# eggs
+.eggs
+
+dev
+
+**/*cache*
+
+.coverage
+coverage.xml
+
+.idea
+
+docs/api
+docs/_build
+
+.ipynb_checkpoints
+*/.ipynb_checkpoints/*
+*/lightning_logs/*
+lightning_logs
+.DS_Store
+.vscode
+
+*.pyc
+
+*.sublime*
+
+Pipfile
+Pipfile.lock
+/htmlcov/

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pykoopman"
-version = "1.0.8"
+version = "1.0.9"
 authors = [
     { name = "Shaowu Pan", email = "[email protected]" },
     { name = "Eurika Kaiser", email = "[email protected]" },
@@ -42,6 +42,7 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = [
+    "pytest <= 7.4.4",
     "pytest-cov ~= 4.1.0",
     "pytest-lazy-fixture ~= 0.6.3",
     "flake8-builtins-unleashed ~= 1.3.1",

requirements-dev.txt

@@ -0,0 +1,24 @@
+-e .
+-r requirements.txt
+# -r requirements-torch.txt --extra-index-url https://download.pytorch.org/whl/cu121/
+-f https://download.pytorch.org/whl/cu121/torch_stable.html
+-r requirements-torch.txt
+
+pytest <= 7.4.4
+pytest-cov ~= 4.1.0
+pytest-lazy-fixture ~= 0.6.3
+flake8-builtins-unleashed ~= 1.3.1
+setuptools_scm ~= 8.0.2
+setuptools_scm_git_archive
+jupyter >= 1.0.0
+notebook > 7.0.0, <= 7.0.4
+nbsphinx
+sphinx-codeautolink
+sphinx >= 3, <= 7.0.0
+sphinxcontrib-apidoc
+sphinx_rtd_theme
+pre-commit
+sphinx-nbexamples
+jupyter_contrib_nbextensions
+PyQt5
+osqp

requirements-torch.txt

@@ -0,0 +1,3 @@
+torch == 2.1.0+cu121
+torchvision
+lightning

requirements-torch.txt~

@@ -0,0 +1,3 @@
+torch
+torchvision
+lightning

requirements.txt

@@ -0,0 +1,8 @@
+matplotlib == 3.5.0
+derivative ~= 0.6.0
+scikit-learn == 1.1.3
+numpy >= 1.20, <= 1.26
+scipy > 1.6.0, <= 1.11.2
+pydmd > 0.4, <= 0.4.1
+optht ~= 0.2.0
+prettytable > 3.0.0, <= 3.9.0

src/pykoopman/common/validation.py

@@ -7,12 +7,18 @@
 
 
 def validate_input(x, t=T_DEFAULT):
-    if not isinstance(x, np.ndarray):
-        raise ValueError("x must be array-like")
+    if not isinstance(x, np.ndarray) and not isinstance(x, list):
+        raise ValueError("x must be array-like OR a list of array-like")
+    elif isinstance(x, list):
+        for i in range(len(x)):
+            x[i] = validate_input(x[i], t)
+        return x
     elif x.ndim == 1:
         x = x.reshape(-1, 1)
     x = check_array(x)
 
+    # add another case if x is a list of trajectory
+
     if t is not T_DEFAULT:
         if t is None:
             raise ValueError("t must be a scalar or array-like.")

src/pykoopman/koopman.py

@@ -7,7 +7,6 @@
 
 import numpy as np
 from numpy import empty
-from numpy import vstack
 from pydmd import DMD
 from pydmd import DMDBase
 from sklearn.base import BaseEstimator
@@ -186,7 +185,9 @@ def fit(self, x, y=None, u=None, dt=1):
             self.n_control_features_ = self._pipeline.steps[1][1].n_control_features_
 
         # compute amplitudes
-        if y is None:
+        if isinstance(x, list):
+            self._amplitudes = None
+        elif y is None:
             if hasattr(self.observables, "n_consumed_samples"):
                 # g0 = self.observables.transform(
                 #     x[0 : 1 + self.observables.n_consumed_samples]
@@ -196,6 +197,7 @@ def fit(self, x, y=None, u=None, dt=1):
                 )
             else:
                 # g0 = self.observables.transform(x[0:1])
+
                 self._amplitudes = np.abs(self.psi(x[0:1].T))
         else:
             self._amplitudes = None
@@ -221,12 +223,13 @@ def predict(self, x, u=None):
                 Time series of external actuation/control.
 
         Returns:
-            y: numpy.ndarray, shape (n_samples, n_input_features)
+            x_next: numpy.ndarray, shape (n_samples, n_input_features)
                 Predicted state one timestep in the future.
         """
 
         check_is_fitted(self, "n_output_features_")
-        return self.observables.inverse(self._step(x, u))
+        x_next = self.observables.inverse(self._step(x, u))
+        return x_next
 
     def simulate(self, x0, u=None, n_steps=1):
         """Simulate an initial state forward in time with the learned Koopman model.
@@ -247,44 +250,46 @@ def simulate(self, x0, u=None, n_steps=1):
                 Number of forward steps to be simulated.
 
         Returns:
-            y: numpy.ndarray, shape (n_steps, n_input_features)
+            x: numpy.ndarray, shape (n_steps, n_input_features)
                 Simulated states.
-                Note that `y[0, :]` is one timestep ahead of `x0`.
+                Note that `x[0, :]` is one timestep ahead of `x0`.
         """
         check_is_fitted(self, "n_output_features_")
         # Could have an option to only return the end state and not all
         # intermediate states to save memory.
 
-        y = empty((n_steps, self.n_input_features_), dtype=self.A.dtype)
-        if u is None:
-            y[0] = self.predict(x0)
+        if x0.ndim == 1:  # handle non-time delay input but 1D accidently
+            x0 = x0.reshape(-1, 1)
+        elif x0.ndim == 2 and x0.shape[0] > 1:  # handle time delay input
+            x0 = x0.T
         else:
-            y[0] = self.predict(x0, u[0])
+            raise TypeError("Check your initial condition shape!")
 
-        if isinstance(self.observables, TimeDelay):
-            n_consumed_samples = self.observables.n_consumed_samples
-            if u is None:
-                for k in range(n_consumed_samples):
-                    y[k + 1] = self.predict(vstack((x0[k + 1 :], y[: k + 1])))
-                for k in range(n_consumed_samples, n_steps - 1):
-                    y[k + 1] = self.predict(y[k - n_consumed_samples : k + 1])
-            else:
-                for k in range(n_consumed_samples):
-                    y[k + 1] = self.predict(vstack((x0[k + 1 :], y[: k + 1])), u[k + 1])
-                for k in range(n_consumed_samples, n_steps - 1):
-                    y[k + 1] = self.predict(y[k - n_consumed_samples : k + 1], u[k + 1])
+        # x = empty((n_steps, self.n_input_features_), dtype=self.A.dtype)
+        y = empty((n_steps, self.A.shape[0]), dtype=self.W.dtype)
 
+        if u is None:
+            # lifted eigen space and move 1 step forward
+            y[0] = self.lamda @ self.psi(x0).flatten()
+
+            # iterate in the lifted space
+            for k in range(n_steps - 1):
+                # tmp = self.W @ self.lamda**(k+1) @ y[0].reshape(-1,1)
+                y[k + 1] = self.lamda @ y[k]
+            x = np.transpose(self.W @ y.T)
+            x = x.astype(self.A.dtype)
         else:
-            if u is None:
-                for k in range(n_steps - 1):
-                    y[k + 1] = self.predict(y[k].reshape(1, -1))
-            else:
-                for k in range(n_steps - 1):
-                    y[k + 1] = self.predict(
-                        y[k].reshape(1, -1), u[k + 1].reshape(1, -1)
-                    )
+            # lifted space (not eigen)
+            y[0] = self.A @ self.phi(x0).flatten() + self.B @ u[0]
+
+            # iterate in the lifted space
+            for k in range(n_steps - 1):
+                tmp = self.A @ y[k].reshape(-1, 1) + self.B @ u[k + 1].reshape(-1, 1)
+                y[k + 1] = tmp.flatten()
+            x = np.transpose(self.C @ y.T)
+            x = x.astype(self.A.dtype)
 
-        return y
+        return x
 
     def get_feature_names(self, input_features=None):
         """Get the names of the individual features constituting the observables.

src/pykoopman/observables/_identity.py

@@ -34,12 +34,20 @@ def fit(self, x, y=None):
 
         Returns:
             self: Returns a fit instance of the class `pykoopman.observables.Identity`.
+
+        Note:
+            only identity mapping is supported for list of arb trajectories
         """
         x = validate_input(x)
-        self.n_input_features_ = self.n_output_features_ = x.shape[1]
+        if not isinstance(x, list):
+            self.n_input_features_ = self.n_output_features_ = x.shape[1]
+            self.measurement_matrix_ = np.eye(x.shape[1]).T
+        else:
+            self.n_input_features_ = self.n_output_features_ = x[0].shape[1]
+            self.measurement_matrix_ = np.eye(x[0].shape[1]).T
+
         self.n_consumed_samples = 0
 
-        self.measurement_matrix_ = np.eye(x.shape[1]).T
         return self
 
     def transform(self, x):

src/pykoopman/observables/_time_delay.py

@@ -62,7 +62,8 @@ def __init__(self, delay=1, n_delays=2):
         Initialize the TimeDelay class with given parameters.
 
         Args:
-            delay (int, optional): The length of each delay. Defaults to 1.
+            delay (int, optional): The length of each delay. Defaults to 1. Or
+                we say this is the "stride of delay".
             n_delays (int, optional): The number of delays to compute for each
                 variable. Defaults to 2.
 

src/pykoopman/regression/_nndmd.py

@@ -91,7 +91,9 @@ class FFNN(nn.Module):
         layers (nn.ModuleList): A list of the neural network layers.
     """
 
-    def __init__(self, input_size, hidden_sizes, output_size, activations):
+    def __init__(
+        self, input_size, hidden_sizes, output_size, activations, include_state=False
+    ):
         super(FFNN, self).__init__()
 
         activations_dict = {
@@ -103,6 +105,12 @@ def __init__(self, input_size, hidden_sizes, output_size, activations):
             "mish": nn.Mish(),
             "linear": nn.Identity(),
         }
+        # whether to directly encode state in observables
+        self.include_state = include_state
+        if self.include_state:
+            output_size_ = output_size - input_size
+        else:
+            output_size_ = output_size
 
         # Define the activation
         act = activations_dict[activations]
@@ -117,7 +125,8 @@ def __init__(self, input_size, hidden_sizes, output_size, activations):
             bias = True
 
         if len(hidden_sizes) == 0:
-            self.layers.append(nn.Linear(input_size, output_size, bias))
+            # if no hidden layer, then entire NN is just a linear one
+            self.layers.append(nn.Linear(input_size, output_size_, bias))
         else:
             self.layers.append(nn.Linear(input_size, hidden_sizes[0], bias))
             if activations != "linear":
@@ -132,8 +141,7 @@ def __init__(self, input_size, hidden_sizes, output_size, activations):
                     self.layers.append(act)
 
             # Define the last output layer
-            bias_last = False  # True  # last layer with bias
-            self.layers.append(nn.Linear(hidden_sizes[-1], output_size, bias_last))
+            self.layers.append(nn.Linear(hidden_sizes[-1], output_size_, bias=False))
 
     def forward(self, x):
         """Performs a forward pass through the neural network.
@@ -144,11 +152,24 @@ def forward(self, x):
         Returns:
             torch.Tensor: The output tensor of the neural network.
         """
+        in_x = x
         for layer in self.layers:
             x = layer(x)
+
+        if self.include_state:
+            x = torch.cat((in_x, x), 1)
         return x
 
 
+class HardCodedLinearLayer(nn.Module):
+    def __init__(self, input_size, output_size):
+
+        pass
+
+    def forward(self, x):
+        pass
+
+
 class BaseKoopmanOperator(nn.Module):
     """Base class for Koopman operator models.
 
@@ -186,6 +207,8 @@ def forward(self, x):
         """
         Computes the forward pass of the `BaseKoopmanOperator`.
 
+        Given `x` as a row vector, return `x @ K.T`
+
         Args:
             x (torch.Tensor): The input tensor.
 
@@ -362,6 +385,7 @@ def __init__(
         config_decoder=dict(),
         lbfgs=False,
         std_koopman=1e-1,
+        include_state=False,
     ):
         super(DLKoopmanRegressor, self).__init__()
 
@@ -373,6 +397,7 @@ def __init__(
             hidden_sizes=config_encoder["hidden_sizes"],
             output_size=config_encoder["output_size"],
             activations=config_encoder["activations"],
+            include_state=include_state,
         )
 
         self._decoder = FFNN(
@@ -1076,6 +1101,7 @@ def __init__(
         normalize_mode="equal",
         normalize_std_factor=2.0,
         std_koopman=1e-1,
+        include_state=False,
         trainer_kwargs={},
     ):
         """Initializes the NNDMD model."""
@@ -1091,6 +1117,7 @@ def __init__(
         self.normalize_std_factor = normalize_std_factor
         self.batch_size = batch_size
         self.std_koopman = std_koopman
+        self.include_state = include_state
 
         # build DLK regressor
         self._regressor = DLKoopmanRegressor(