Merge pull request #88 from EthanJamesLew/feature/sw-edmd

Feature: State Weighted eDMD (SW-eDMD)

Author: EthanJamesLew

autokoopman/autokoopman.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-import autokoopman.core.observables as kobs
+import autokoopman.observable as kobs
 from autokoopman.core.trajectory import (
     TrajectoriesData,
     UniformTimeTrajectoriesData,
@@ -24,7 +24,7 @@
 from autokoopman.tuner.gridsearch import GridSearchTuner
 from autokoopman.tuner.montecarlo import MonteCarloTuner
 from autokoopman.tuner.bayesianopt import BayesianOptTuner
-from autokoopman.core.observables import KoopmanObservable
+from autokoopman.observable.observables import KoopmanObservable
 from autokoopman.core.format import hide_prints
 
 __all__ = ["auto_koopman"]

autokoopman/estimator/koopman.py

@@ -26,7 +26,11 @@ def dmdc(X, Xp, U, r):
     U, Sigma, V = U[:, :r], np.diag(Sigma)[:r, :r], V.conj().T[:, :r]
 
     # get the transformation
-    Atilde = Yp @ V @ np.linalg.inv(Sigma) @ U.conj().T
+    try:
+        Atilde = Yp @ V @ np.linalg.inv(Sigma) @ U.conj().T
+    except:
+        Atilde = Yp @ V @ np.linalg.pinv(Sigma) @ U.conj().T
+
     return Atilde[:, :state_size], Atilde[:, state_size:]
 
 
@@ -48,33 +52,69 @@ def wdmdc(X, Xp, U, r, W):
     U, Sigma, V = U[:, :r], np.diag(Sigma)[:r, :r], V.conj().T[:, :r]
 
     # get the transformation
-    Atilde = Yp @ V @ np.linalg.inv(Sigma) @ U.conj().T
+    # get the transformation
+    try:
+        Atilde = Yp @ V @ np.linalg.inv(Sigma) @ U.conj().T
+    except:
+        Atilde = Yp @ V @ np.linalg.pinv(Sigma) @ U.conj().T
     return Atilde[:, :state_size], Atilde[:, state_size:]
 
 
 def swdmdc(X, Xp, U, r, Js, W):
     """State Weighted Dynamic Mode Decomposition with Control (wDMDC)"""
+    import cvxpy as cp
+    import warnings
+    
     assert len(W.shape) == 2, "weights must be 2D for snapshot x state"
 
     if U is not None:
-        Y = np.hstack((X, U)).T
+        Y = np.hstack((X, U))
     else:
-        Y = X.T
-    Yp = Xp.T
-
-    # compute observables weights from state weights
-    Wy = np.vstack([(np.abs(J) @ np.atleast_2d(w).T).T for J, w in zip(Js, W)])
-
-    # apply weights element-wise
-    Y, Yp = Wy.T * Y, Wy.T * Yp
-    state_size = Yp.shape[0]
-
-    # compute Atilde
-    U, Sigma, V = np.linalg.svd(Y, False)
-    U, Sigma, V = U[:, :r], np.diag(Sigma)[:r, :r], V.conj().T[:, :r]
-
+        Y = X
+    Yp = Xp
+    state_size = Yp.shape[1]
+
+    n_snap, n_obs = Yp.shape
+    n_inps = U.shape[1] if U is not None else 0
+    _, n_states = Js[0].shape
+
+    # so the objective isn't numerically unstable
+    sf = (1.0 / n_snap)
+
+    # check that rank is less than the number of observations
+    if r > n_obs:
+        warnings.warn("Rank must be less than the number of observations. Reducing rank to n_obs.")
+        r = n_obs
+    
+    # koopman operator
+    # Variables for the low-rank approximation
+    K_U = cp.Variable((n_obs, r))
+    K_V = cp.Variable((r, n_obs + n_inps))
+
+    # SW-eDMD objective
+    weights_obj = np.vstack([(np.abs(J) @ w) for J, w in zip(Js, W)]).T 
+    P = sf * cp.multiply(weights_obj, Yp.T - (K_U @ K_V) @ Y.T)
+    # add regularization 
+    objective = cp.Minimize(cp.sum_squares(P) + 1E-4 * 1.0 / (n_obs**2) * cp.norm(K_U @ K_V, "fro"))
+
+    # unconstrained problem
+    constraints = None
+
+    # SW-eDMD problem
+    prob = cp.Problem(objective, constraints)
+
+    # solve for the SW-eDMD Koopman operator
+    try:
+        _ = prob.solve(solver=cp.CLARABEL)
+        #_ = prob.solve(solver=cp.ECOS)
+        if K_U.value is None or K_V.value is None:
+            raise Exception("SW-eDMD (cvxpy) Optimization failed to converge.")
+    except:
+        warnings.warn("SW-eDMD (cvxpy) Optimization failed to converge. Switching to unweighted DMDc.")
+        return dmdc(X, Xp, U, r)
+    
     # get the transformation
-    Atilde = Yp @ V @ np.linalg.inv(Sigma) @ U.conj().T
+    Atilde = K_U.value @ K_V.value 
     return Atilde[:, :state_size], Atilde[:, state_size:]
 
 

autokoopman/observable/__init__.py

@@ -0,0 +1,7 @@
+"""
+Koopman Observables Module
+"""
+
+from .observables import *
+from .rff import *
+from .reweighted import *

autokoopman/core/observables.py renamed to autokoopman/observable/observables.py

@@ -4,8 +4,6 @@
 import numpy as np
 import sympy as sp  # type: ignore
 
-from scipy.stats import cauchy, laplace
-
 
 class KoopmanObservable(abc.ABC):
     """
@@ -163,51 +161,3 @@ def __init__(self, dimension, degree) -> None:
 
     def obs_fcn(self, X: np.ndarray) -> np.ndarray:
         return self.poly.transform(np.atleast_2d(X)).T
-
-
-class RFFObservable(KoopmanObservable):
-    def __init__(self, dimension, num_features, gamma, metric="rbf"):
-        super(RFFObservable, self).__init__()
-        self.gamma = gamma
-        self.dimension = dimension
-        self.metric = metric
-        self.D = num_features
-        # Generate D iid samples from p(w)
-        if self.metric == "rbf":
-            self.w = np.sqrt(2 * self.gamma) * np.random.normal(
-                size=(self.D, self.dimension)
-            )
-        elif self.metric == "laplace":
-            self.w = cauchy.rvs(scale=self.gamma, size=(self.D, self.dimension))
-        # Generate D iid samples from Uniform(0,2*pi)
-        self.u = 2 * np.pi * np.random.rand(self.D)
-
-    def obs_fcn(self, X: np.ndarray) -> np.ndarray:
-        # modification...
-        if len(X.shape) == 1:
-            x = np.atleast_2d(X.flatten()).T
-        else:
-            x = X.T
-        w = self.w.T
-        u = self.u[np.newaxis, :].T
-        s = np.sqrt(2 / self.D)
-        Z = s * np.cos(x.T @ w + u.T)
-        return Z.T
-
-    def obs_grad(self, X: np.ndarray) -> np.ndarray:
-        if len(X.shape) == 1:
-            x = np.atleast_2d(X.flatten()).T
-        else:
-            x = X.T
-        x = np.atleast_2d(X.flatten()).T
-        w = self.w.T
-        u = self.u[np.newaxis, :].T
-        s = np.sqrt(2 / self.D)
-        # TODO: make this sparse?
-        Z = -s * np.diag(np.sin(u + w.T @ x).flatten()) @ w.T
-        return Z
-
-    def compute_kernel(self, X: np.ndarray) -> np.ndarray:
-        Z = self.transform(X)
-        K = Z.dot(Z.T)
-        return K

autokoopman/observable/reweighted.py

@@ -0,0 +1,151 @@
+from autokoopman.observable import SymbolicObservable, KoopmanObservable
+import math
+import numpy as np
+import sympy as sp
+from scipy.stats import cauchy, laplace
+from scipy.optimize import nnls
+
+
+def make_gaussian_kernel(sigma=1.0):
+  """the baseline"""
+  def kernel(x, xp):
+    return np.exp(-np.linalg.norm(x-xp)**2.0 / (2.0*sigma**2.0))
+  return kernel
+
+
+def rff_reweighted_map(n, X, Y, wx, wy, sigma=1.0):
+  """build a RFF explicit feature map"""
+  assert len(X) == len(Y)
+  R = n
+  D = X.shape[1]
+  N = len(X)
+  W    = np.random.normal(loc=0, scale=(1.0/np.sqrt(sigma)), size=(R, D))
+  b    = np.random.uniform(-np.pi, np.pi, size = R)
+
+  # get ground truth kernel for training
+  kernel = make_gaussian_kernel(sigma)
+
+  # solve NNLS problem
+  M = np.zeros((N, R))
+  bo = np.zeros((N,))
+  for jdx, (xi, yi, wxi, wyi) in enumerate(zip(X, Y, wx, wy)):
+    wi = np.sum(wxi) * np.sum(wyi)
+    k = wi * kernel(xi, yi)
+    if np.isclose(np.abs(wi), 0.0):
+      continue
+    bo[jdx] = k
+    for idx, (omegai, bi) in enumerate(zip(W, b)):
+      M[jdx, idx] = wi * np.cos(np.dot(omegai, xi) + bi) * np.cos(np.dot(omegai, yi) + bi)
+
+  a, _ = nnls(M, bo, maxiter=1000)
+
+  # get new values
+  new_idxs = (np.abs(a) > 3E-5)
+  W = W[new_idxs]
+  b = b[new_idxs]
+  a = a[new_idxs]
+
+  return a, W, b
+
+
+def subsampled_dense_grid(d, D, gamma, deg=8):
+  """sparse grid gaussian quadrature"""
+  points, weights = np.polynomial.hermite.hermgauss(deg)
+  points *= 2 * math.sqrt(gamma)
+  weights /= math.sqrt(math.pi)
+  # Now @weights must sum to 1.0, as the kernel value at 0 is 1.0
+  subsampled_points = np.random.choice(points, size=(D, d), replace=True, p=weights)
+  subsampled_weights = np.ones(D) / math.sqrt(D)
+  return subsampled_points, subsampled_weights
+
+
+def dense_grid(d, D, gamma, deg=8):
+  """dense grid gaussian quadrature"""
+  points, weights = np.polynomial.hermite.hermgauss(deg)
+  points *= 2 * math.sqrt(gamma)
+  weights /= math.sqrt(math.pi)
+  points = np.atleast_2d(points).T
+  return points, weights
+
+
+def sparse_grid_map(n, d, sigma=1.0):
+  """sparse GQ explicit map"""
+  d, D = d, n
+  gamma = 1/(2*sigma*sigma)
+  points, weights = subsampled_dense_grid(d, D, gamma=gamma)
+  def inner(x):
+    prod = x @ points.T
+    return np.hstack((weights * np.cos(prod), weights * np.sin(prod)))
+  return inner
+
+
+def dense_grid_map(n, d, sigma=1.0):
+  """dense GQ explicit map"""
+  d, D = d, n
+  gamma = 1/(2*sigma*sigma)
+  points, weights = dense_grid(d, D, gamma=gamma)
+  def inner(x):
+    prod = x @ points.T
+    return np.hstack((weights * np.cos(prod), weights * np.sin(prod)))
+  return inner
+
+
+def quad_reweighted_map(n, X, Y, wx, wy, sigma=1.0):
+  """build a RFF explicit feature map"""
+  assert len(X) == len(Y)
+  R = int(n / 2.0)
+  D = X.shape[1]
+  N = len(X)
+  W, a = subsampled_dense_grid(D, R, gamma=1/(2*sigma*sigma))
+  #W    = np.random.normal(loc=0, scale=(1.0/np.sqrt(sigma)), size=(R, D))
+  b    = np.random.uniform(-np.pi, np.pi, size = R)
+  #print(X.shape, W.shape)
+  #b = np.arccos(-np.sqrt(np.cos(2*X.T.dot(W)) + 1)/np.sqrt(2.0))
+  #print(b)
+
+  # get ground truth kernel for training
+  kernel = make_gaussian_kernel(sigma)
+
+  # solve NNLS problem
+  M = np.zeros((N, R))
+  bo = np.zeros((N,))
+  for jdx, (xi, yi, wxi, wyi) in enumerate(zip(X, Y, wx, wy)):
+      k = kernel(xi, yi)
+      bo[jdx] = k
+      for idx, (omegai, ai, bi) in enumerate(zip(W, a, b)):
+        M[jdx, idx] = np.cos(np.dot(omegai, xi - yi)) 
+
+  a, _ = nnls(M, bo, maxiter=1000)
+
+  # get new values
+  new_idxs = (np.abs(a) > 1E-7)
+  W = W[new_idxs]
+  b = b[new_idxs]
+  a = a[new_idxs]
+
+  return a, W, b
+
+
+class ReweightedRFFObservable(SymbolicObservable):
+    def __init__(self, dimension, num_features, gamma, X, Y, wx, wy, feat_type='rff'):
+        self.dimension, self.num_features = dimension, num_features
+        n = num_features
+        if feat_type == 'quad':
+          self.a, self.W, self.b = quad_reweighted_map(n, X, Y, wx, wy, np.sqrt(1/(2*gamma)))
+        elif feat_type == 'rff':
+          self.a, self.W, self.b = rff_reweighted_map(n, X, Y, wx, wy, np.sqrt(1/(2*gamma)))
+        else:
+          raise ValueError('feat_type must be quad or rff')
+        
+        # make sympy variables for each of the state dimensions
+        self.variables = [f'x{i}' for i in range(self.dimension)]
+        self._observables = sp.symbols(self.variables)
+        X = sp.Matrix(self._observables)
+        
+        # build observables sympy expressions from self.weights from self.weights, x.T @ points
+        self.expressions = []
+        for ai, wi, bi in zip(self.a, self.W, self.b):
+            self.expressions.append(np.sqrt(ai) * sp.cos(X.dot(wi)))
+            self.expressions.append(np.sqrt(ai) * sp.sin(X.dot(wi)))
+
+        super(ReweightedRFFObservable, self).__init__(self.variables, self.expressions) 

autokoopman/observable/rff.py

@@ -0,0 +1,52 @@
+import numpy as np
+from scipy.stats import cauchy, laplace
+
+from .observables import KoopmanObservable
+
+
+class RFFObservable(KoopmanObservable):
+    def __init__(self, dimension, num_features, gamma, metric="rbf"):
+        super(RFFObservable, self).__init__()
+        self.gamma = gamma
+        self.dimension = dimension
+        self.metric = metric
+        self.D = num_features
+        # Generate D iid samples from p(w)
+        if self.metric == "rbf":
+            self.w = np.sqrt(2 * self.gamma) * np.random.normal(
+                size=(self.D, self.dimension)
+            )
+        elif self.metric == "laplace":
+            self.w = cauchy.rvs(scale=self.gamma, size=(self.D, self.dimension))
+        # Generate D iid samples from Uniform(0,2*pi)
+        self.u = 2 * np.pi * np.random.rand(self.D)
+
+    def obs_fcn(self, X: np.ndarray) -> np.ndarray:
+        # modification...
+        if len(X.shape) == 1:
+            x = np.atleast_2d(X.flatten()).T
+        else:
+            x = X.T
+        w = self.w.T
+        u = self.u[np.newaxis, :].T
+        s = np.sqrt(2 / self.D)
+        Z = s * np.cos(x.T @ w + u.T)
+        return Z.T
+
+    def obs_grad(self, X: np.ndarray) -> np.ndarray:
+        if len(X.shape) == 1:
+            x = np.atleast_2d(X.flatten()).T
+        else:
+            x = X.T
+        x = np.atleast_2d(X.flatten()).T
+        w = self.w.T
+        u = self.u[np.newaxis, :].T
+        s = np.sqrt(2 / self.D)
+        # TODO: make this sparse?
+        Z = -s * np.diag(np.sin(u + w.T @ x).flatten()) @ w.T
+        return Z
+
+    def compute_kernel(self, X: np.ndarray) -> np.ndarray:
+        Z = self.transform(X)
+        K = Z.dot(Z.T)
+        return K