fix bug in test generation of traj lengths, update deepkoopman call of

next step traj

autokoopman/estimator/deepkoopman.py

@@ -329,7 +329,7 @@ def _get_loss(mseloss, l1loss, x, xn, y, yn, xr, Xn):
             for step_size in range(1, self.rollout_steps + 1):
                 # upload to GPU and normalize the data
                 # NOTE: this is expensive :(
-                X, Xn, U = trajs.n_step_matrices(step_size)
+                X, Xn, U, _ = trajs.n_step_matrices(step_size)
                 X = torch.tensor(X.T, dtype=torch.float32).to(self.device)
                 Xn = torch.tensor(Xn.T, dtype=torch.float32).to(self.device)
                 U = (

test/unit_test/test_autokoopman.py

@@ -30,7 +30,6 @@ def test_autokoopman(obs, opt, cost, normalize):
     sp = 0.1
     ivs = np.random.uniform(low=-2.0, high=2.0, size=(20, 2))
     t_ends = [np.random.random() + 1.0 for idx in range(len(ivs))]
-    lengths = [int(t_end // sp) for t_end in t_ends]
     training_data = UniformTimeTrajectoriesData(
         {
             idx: fhn.solve_ivp(
@@ -45,7 +44,7 @@ def test_autokoopman(obs, opt, cost, normalize):
     # produce scoring weights if the cost is weighted
     if cost == "weighted":
         scoring_weights = {
-            idx: np.ones((length,)) * 0.01 for idx, length in enumerate(lengths)
+            idx: np.ones(len(traj.states)) * 0.01 for idx, traj in training_data._trajs.items() 
         }
     else:
         scoring_weights = None
@@ -86,8 +85,7 @@ def test_autokoopman_np(obs, opt, cost, normalize):
     # given issue #29, let's make these differently sized
     sp = 0.1
     ivs = np.random.uniform(low=-2.0, high=2.0, size=(20, 2))
-    t_ends = [np.random.random() + 1.0 for idx in range(len(ivs))]
-    lengths = [int(t_end // sp) for t_end in t_ends]
+    t_ends = [1.5 for idx in range(len(ivs))]
     training_data = UniformTimeTrajectoriesData(
         {
             idx: fhn.solve_ivp(
@@ -102,15 +100,15 @@ def test_autokoopman_np(obs, opt, cost, normalize):
     # produce scoring weights if the cost is weighted
     if cost == "weighted":
         scoring_weights = {
-            idx: np.ones((length,)) * 0.01 for idx, length in enumerate(lengths)
+            idx: np.ones((len(t.states),)) * 0.01 for idx, t in training_data._trajs.items()
         }
     else:
         scoring_weights = None
 
     # put into numpy arrays for testing
     training_data_np = []
     scoring_weights_np = []
-    for name in training_data.traj_names:
+    for name in range(len(ivs)):
         training_data_np.append(training_data[name].states)
         if scoring_weights is not None:
             scoring_weights_np.append(scoring_weights[name])