Merge pull request #112 from decargroup/bugfix/111-predict_trajectory…

…-indexing-is-wrong-when-relift_state=false

`predict_trajectory` indexing is wrong when `relift_state=false`

pykoop/kernel_approximation.py

@@ -96,9 +96,9 @@ class RandomFourierKernelApprox(KernelApproximation):
     ift_ : scipy.stats.rv_continuous
         Probability distribution corresponding to inverse Fourier transform of
         chosen kernel.
-    random_weights_ : np.ndarray, shape (n_features_in_, n_components)
+    random_weights_ : np.ndarray, shape (n_features, n_components)
         Random weights to inner-product with features.
-    random_offsets_ : np.ndarray, shape (n_features_in_, )
+    random_offsets_ : np.ndarray, shape (n_features, )
         Random offsets if ``method`` is ``'weight_offset'``.
 
     Examples
@@ -295,9 +295,9 @@ class RandomBinningKernelApprox(KernelApproximation):
         estimators from :mod:`sklearn.kernel_approximation`.
     ddot_ : scipy.stats.rv_continuous
         Probability distribution corresponding to ``\delta \ddot{k}(\delta)``.
-    pitches_ : np.ndarray, shape (n_features_in_, n_components)
+    pitches_ : np.ndarray, shape (n_features, n_components)
         Grid pitches for each component.
-    shifts_ : np.ndarray, shape (n_features_in_, n_components)
+    shifts_ : np.ndarray, shape (n_features, n_components)
         Grid shifts for each component.
     encoder_ : sklearn.preprocessing.OneHotEncoder
         One-hot encoder used for hashing sample coordinates for each component.
@@ -339,8 +339,8 @@ def __init__(
             Kernel to approximate. Possible options are
 
                 - ``'laplacian'`` -- Laplacian kernel, with
-                ``\delta \ddot{k}(\delta)`` being :class:`scipy.stats.gamma`
-                with shape parameter ``a=2`` (default).
+                  ``\delta \ddot{k}(\delta)`` being :class:`scipy.stats.gamma`
+                  with shape parameter ``a=2`` (default).
 
             Alternatively, a separable, positive, shift-invariant kernel can be
             implicitly specified by providing ``\delta \ddot{k}(\delta)`` as a

pykoop/koopman_pipeline.py

@@ -2317,6 +2317,12 @@ def predict_multistep(self, X: np.ndarray) -> np.ndarray:
         ------
         ValueError
             If an episode is shorter than ``min_samples_``.
+
+        Warning
+        -------
+        Deprecated in favour of
+        :func:`pykoop.KoopmanPipeline.predict_trajectory`.
+
         """
         sklearn.utils.validation.check_is_fitted(self, 'regressor_fit_')
         X = sklearn.utils.validation.check_array(X, **self._check_array_params)
@@ -2507,26 +2513,56 @@ def predict_trajectory(
                 )
             else:
                 # Number of steps in episode
-                n_steps_i = U_i.shape[0] - self.min_samples_
+                n_steps_i = U_i.shape[0] - self.min_samples_ + 1
                 # Initial conditions
                 Theta_i = np.zeros((n_steps_i, self.n_states_out_))
                 Upsilon_i = np.zeros((n_steps_i, self.n_inputs_out_))
+                X_i = np.zeros((U_i.shape[0], self.n_states_in_))
                 Theta_i[[0], :] = self.lift_state(X0_i, episode_feature=False)
+                X_i[:self.min_samples_, :] = X0_i
                 for k in range(1, n_steps_i + 1):
                     try:
-                        X_ikm1 = self.retract_state(
-                            Theta_i[[k - 1], :],
-                            episode_feature=False,
-                        )
-                        window = np.s_[k:(k + self.min_samples_)]
+                        # Lift input. We need to use the retracted state here
+                        # because some lifting functions are state-dependent.
+                        # ``Theta_i`` should not have any lifting and
+                        # retracting happening, but for ``Upsilon_i``, it is
+                        # unavoidable.
+                        window_km1 = np.s_[(k - 1):(k + self.min_samples_ - 1)]
                         Upsilon_i[[k - 1], :] = self.lift_input(
-                            np.hstack((X_ikm1, U_i[window, :])),
+                            np.hstack((
+                                X_i[window_km1, :],
+                                U_i[window_km1, :],
+                            )),
                             episode_feature=False,
                         )
-                        # Predict
                         if k < n_steps_i:
+                            # Predict next lifted state
                             Theta_i[[k], :] = (Theta_i[[k - 1], :] @ A.T
                                                + Upsilon_i[[k - 1], :] @ B.T)
+                            # Retract and store state. ``k`` is index of
+                            # ``Theta_i``, which is shorter than ``X_i`` when
+                            # there are delays.
+                            #
+                            # If ``min_samples_=3``, ``Theta_i`` is two entries
+                            # shorter than ``X_i``. The first entry of
+                            # ``Theta_i`` occurs at the same "time" as the
+                            # third entry of ``X_i``. Thus ``Theta_i[k]``
+                            # corresponds to
+                            # ``X_i[k + self.min_samples_ - 1]``.
+                            #
+                            # Let ``self.min_samples_=3``. An illustration:
+                            #
+                            #            X_i[0:3]
+                            #            vvvvv
+                            #     X_i: [ | | | . . . . . . . ]
+                            # Theta_i:     [ | . . . . . . . ]
+                            #                ^
+                            #                Theta_i[0]
+                            X_ik = self.retract_state(
+                                Theta_i[[k], :],
+                                episode_feature=False,
+                            )
+                            X_i[[k + self.min_samples_ - 1], :] = X_ik[[-1], :]
                     except ValueError as ve:
                         if (np.all(np.isfinite(X_ikm1))
                                 and np.all(np.isfinite(Theta_i))
@@ -2538,7 +2574,6 @@ def predict_trajectory(
                             Theta_i[crash_index:, :] = 0
                             Upsilon_i[crash_index:, :] = 0
                             break
-                X_i = self.retract_state(Theta_i, episode_feature=False)
             # If prediction crashed, set remaining entries to NaN
             if crash_index is not None:
                 log.warning(f'Prediction diverged at index {crash_index}. '

tests/test_koopman_pipeline.py

@@ -994,6 +994,44 @@ def test_predict_trajectory_no_U(self, kp, mass_spring_damper_sine_input):
             X_sim_exp[[k], :] = Xp[[-1], :]
         np.testing.assert_allclose(X_sim, X_sim_exp)
 
+    def test_predict_trajectory_no_relift_state(self, ndarrays_regression, kp,
+                                                mass_spring_damper_sine_input):
+        """Test :func:`predict_trajectory` without relifting state."""
+        msg = 'Test only works when there is no episode feature.'
+        assert (not mass_spring_damper_sine_input['episode_feature']), msg
+        # Fit estimator
+        kp.fit(
+            mass_spring_damper_sine_input['X_train'],
+            n_inputs=mass_spring_damper_sine_input['n_inputs'],
+            episode_feature=False,
+        )
+        # Extract initial conditions
+        x0 = pykoop.extract_initial_conditions(
+            mass_spring_damper_sine_input['X_train'],
+            kp.min_samples_,
+            n_inputs=mass_spring_damper_sine_input['n_inputs'],
+            episode_feature=False,
+        )
+        # Extract input
+        u = pykoop.extract_input(
+            mass_spring_damper_sine_input['X_train'],
+            n_inputs=mass_spring_damper_sine_input['n_inputs'],
+            episode_feature=False,
+        )
+        # Predict new states
+        X_sim = kp.predict_trajectory(
+            x0,
+            u,
+            episode_feature=False,
+            relift_state=False,
+        )
+        ndarrays_regression.check(
+            {
+                'X_sim': X_sim,
+            },
+            default_tolerance=dict(atol=1e-6, rtol=0),
+        )
+
     def test_predict_multistep(self, kp, mass_spring_damper_sine_input):
         """Test :func:`predict_multistep` (deprecated)."""
         msg = 'Test only works when there is no episode feature.'

tests/test_util.py

@@ -233,6 +233,7 @@ def test_feature_names_out(self, names_in, X, names_out, n_inputs,
     pykoop.example_data_msd,
     pykoop.example_data_vdp,
     pykoop.example_data_pendulum,
+    pykoop.example_data_duffing,
 ])
 class TestExampleData:
     """Test example dynamic model data.