From 870726de9bdcfb7aef826791e5ab719ffb87f2a6 Mon Sep 17 00:00:00 2001
From: Steven Dahdah <steven.dahdah@mail.mcgill.ca>
Date: Mon, 14 Aug 2023 16:45:57 -0400
Subject: [PATCH] Add more failing unit tests

---
 tests/test_koopman_pipeline.py | 117 +++++++++++++++++++++++++++++++--
 1 file changed, 110 insertions(+), 7 deletions(-)

diff --git a/tests/test_koopman_pipeline.py b/tests/test_koopman_pipeline.py
index 9207975..b7df2f9 100644
--- a/tests/test_koopman_pipeline.py
+++ b/tests/test_koopman_pipeline.py
@@ -1672,34 +1672,137 @@ def test_combine_episodes(self, X, episodes, episode_feature):
                 'min_samples_': 1,
             },
         ),
-        # With delays but no episode feature
+        # State delay, with episode feature
         (
             pykoop.SplitPipeline(
                 lifting_functions_state=pykoop.DelayLiftingFn(1, 0),
                 lifting_functions_input=None,
             ),
-            np.array(['x0', 'x1', 'u0', 'u1']),
-            # TODO
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
             np.array([
+                [0, 0, 0, 0, 1, 1],
                 [0, 1, 2, 3, 4, 5],
                 [5, 4, 3, 2, 1, 0],
                 [4, 5, 6, 7, 8, 9],
                 [0, 8, 7, 6, 5, 4],
             ]).T,
-            np.array(['x0', 'x1', 'u0', 'u1']),
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
             np.array([
+                [0, 0, 0, 0, 1, 1],
                 [0, 1, 2, 3, 4, 5],
                 [5, 4, 3, 2, 1, 0],
                 [4, 5, 6, 7, 8, 9],
                 [0, 8, 7, 6, 5, 4],
             ]).T,
             2,
-            False,
+            True,
             {
-                'n_features_in_': 4,
+                'n_features_in_': 5,
                 'n_states_in_': 2,
                 'n_inputs_in_': 2,
-                'n_features_out_': 4,
+                'n_features_out_': 5,
+                'n_states_out_': 2,
+                'n_inputs_out_': 2,
+                'min_samples_': 1,
+            },
+        ),
+        # Input delay, with episode feature
+        (
+            pykoop.SplitPipeline(
+                lifting_functions_state=None,
+                lifting_functions_input=pykoop.DelayLiftingFn(0, 1),
+            ),
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
+            np.array([
+                [0, 0, 0, 0, 1, 1],
+                [0, 1, 2, 3, 4, 5],
+                [5, 4, 3, 2, 1, 0],
+                [4, 5, 6, 7, 8, 9],
+                [0, 8, 7, 6, 5, 4],
+            ]).T,
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
+            np.array([
+                [0, 0, 0, 0, 1, 1],
+                [0, 1, 2, 3, 4, 5],
+                [5, 4, 3, 2, 1, 0],
+                [4, 5, 6, 7, 8, 9],
+                [0, 8, 7, 6, 5, 4],
+            ]).T,
+            2,
+            True,
+            {
+                'n_features_in_': 5,
+                'n_states_in_': 2,
+                'n_inputs_in_': 2,
+                'n_features_out_': 5,
+                'n_states_out_': 2,
+                'n_inputs_out_': 2,
+                'min_samples_': 1,
+            },
+        ),
+        # State delay with nonzero ``n_delays_input``, with episode feature
+        (
+            pykoop.SplitPipeline(
+                lifting_functions_state=pykoop.DelayLiftingFn(1, 1),
+                lifting_functions_input=None,
+            ),
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
+            np.array([
+                [0, 0, 0, 0, 1, 1],
+                [0, 1, 2, 3, 4, 5],
+                [5, 4, 3, 2, 1, 0],
+                [4, 5, 6, 7, 8, 9],
+                [0, 8, 7, 6, 5, 4],
+            ]).T,
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
+            np.array([
+                [0, 0, 0, 0, 1, 1],
+                [0, 1, 2, 3, 4, 5],
+                [5, 4, 3, 2, 1, 0],
+                [4, 5, 6, 7, 8, 9],
+                [0, 8, 7, 6, 5, 4],
+            ]).T,
+            2,
+            True,
+            {
+                'n_features_in_': 5,
+                'n_states_in_': 2,
+                'n_inputs_in_': 2,
+                'n_features_out_': 5,
+                'n_states_out_': 2,
+                'n_inputs_out_': 2,
+                'min_samples_': 1,
+            },
+        ),
+        # Input delay with nonzero ``n_delays_state``, with episode feature
+        (
+            pykoop.SplitPipeline(
+                lifting_functions_state=None,
+                lifting_functions_input=pykoop.DelayLiftingFn(1, 1),
+            ),
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
+            np.array([
+                [0, 0, 0, 0, 1, 1],
+                [0, 1, 2, 3, 4, 5],
+                [5, 4, 3, 2, 1, 0],
+                [4, 5, 6, 7, 8, 9],
+                [0, 8, 7, 6, 5, 4],
+            ]).T,
+            np.array(['ep', 'x0', 'x1', 'u0', 'u1']),
+            np.array([
+                [0, 0, 0, 0, 1, 1],
+                [0, 1, 2, 3, 4, 5],
+                [5, 4, 3, 2, 1, 0],
+                [4, 5, 6, 7, 8, 9],
+                [0, 8, 7, 6, 5, 4],
+            ]).T,
+            2,
+            True,
+            {
+                'n_features_in_': 5,
+                'n_states_in_': 2,
+                'n_inputs_in_': 2,
+                'n_features_out_': 5,
                 'n_states_out_': 2,
                 'n_inputs_out_': 2,
                 'min_samples_': 1,