add sw-edmd to continuous koopman

EthanJamesLew · Apr 16, 2024 · 4bec0ad · 4bec0ad
1 parent 18f6dd9
commit 4bec0ad
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/autokoopman/estimator/koopman.py b/autokoopman/estimator/koopman.py
@@ -210,9 +210,20 @@ def fit_gradient(
         if weights is None:
             self._A, self._B = dmdc(G, Gp, U.T if U is not None else U, self.rank)
         else:
-            self._A, self._B = wdmdc(
-                G, Gp, U.T if U is not None else U, self.rank, weights
-            )
+            # TODO: change this condition to be more accurate
+            if False:  # len(weights[0].shape) == 1:
+                self._A, self._B = wdmdc(
+                    G, Gp, U.T if U is not None else U, self.rank, weights
+                )
+            else:
+                self._A, self._B = swdmdc(
+                    G,
+                    Gp,
+                    U.T if U is not None else U,
+                    self.rank,
+                    [self.obs.obs_grad(xi) for xi in X.T],
+                    weights,
+                )
         self._has_input = U is not None
 
     @property