diff --git a/autokoopman/autokoopman.py b/autokoopman/autokoopman.py index 9a210cb..371d3cf 100644 --- a/autokoopman/autokoopman.py +++ b/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"] diff --git a/autokoopman/estimator/koopman.py b/autokoopman/estimator/koopman.py index d642c9a..e047806 100644 --- a/autokoopman/estimator/koopman.py +++ b/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:] diff --git a/autokoopman/observable/__init__.py b/autokoopman/observable/__init__.py new file mode 100644 index 0000000..86f2c5b --- /dev/null +++ b/autokoopman/observable/__init__.py @@ -0,0 +1,7 @@ +""" +Koopman Observables Module +""" + +from .observables import * +from .rff import * +from .reweighted import * \ No newline at end of file diff --git a/autokoopman/core/observables.py b/autokoopman/observable/observables.py similarity index 77% rename from autokoopman/core/observables.py rename to autokoopman/observable/observables.py index ff179a8..cf9a89b 100644 --- a/autokoopman/core/observables.py +++ b/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 diff --git a/autokoopman/observable/reweighted.py b/autokoopman/observable/reweighted.py new file mode 100644 index 0000000..fd10da3 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/autokoopman/observable/rff.py b/autokoopman/observable/rff.py new file mode 100644 index 0000000..19f4613 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/environment.yml b/environment.yml index 5282e13..0c5a2eb 100644 --- a/environment.yml +++ b/environment.yml @@ -14,4 +14,6 @@ dependencies: - gpy==1.13.1 - gpyopt==1.2.6 - torch==2.2.1 - - sphinx_mdinclude==0.5.1 \ No newline at end of file + - sphinx_mdinclude==0.5.1 + - cvxpy==1.4.2 + - sympy==1.12 diff --git a/notebooks/Koopman-Symbolic.ipynb b/notebooks/Koopman-Symbolic.ipynb index 622da2f..68d6627 100644 --- a/notebooks/Koopman-Symbolic.ipynb +++ b/notebooks/Koopman-Symbolic.ipynb @@ -118,7 +118,7 @@ "exprs = fhn._exprs\n", "x0, x1 = variables\n", "basis = generate_monomials(variables, 2)[1:]\n", - "lie_obs = ak.core.observables.SymbolicObservable(variables, obs(exprs, basis, variables))" + "lie_obs = ak.observable.SymbolicObservable(variables, obs(exprs, basis, variables))" ] }, { diff --git a/notebooks/linear-model.ipynb b/notebooks/linear-model.ipynb index e3dfdfe..6d44c19 100644 --- a/notebooks/linear-model.ipynb +++ b/notebooks/linear-model.ipynb @@ -41,7 +41,7 @@ "metadata": {}, "outputs": [], "source": [ - "from autokoopman.core.observables import RFFObservable, PolynomialObservable, IdentityObservable\n", + "from autokoopman.observable import RFFObservable, PolynomialObservable, IdentityObservable\n", "\n", "# augment (combine) observables function\n", "# in this case, combine multiple lengthscales together\n", @@ -192,7 +192,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.7" + "version": "3.9.0" } }, "nbformat": 4, diff --git a/notebooks/weighted-cost-func.ipynb b/notebooks/weighted-cost-func.ipynb index 90b7338..31d945e 100644 --- a/notebooks/weighted-cost-func.ipynb +++ b/notebooks/weighted-cost-func.ipynb @@ -7,12 +7,16 @@ "source": [ "# Weighted Cost Function\n", "\n", - "Shows how to use the cost function requested in [issue #84](https://github.com/EthanJamesLew/AutoKoopman/issues/84)." + "Shows how to use the cost function requested in [issue #84](https://github.com/EthanJamesLew/AutoKoopman/issues/84).\n", + "\n", + "## Note on SW-eDMD\n", + "\n", + "SW-eDMD uses an off-the-shelf optimizer (cvxpy) and so is not very fast yet. **You may need to reduce the number of observables and time points to allow the solve to work.**" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "dfdb47e2-0ea0-4bf8-8279-8500ff3cf21f", "metadata": {}, "outputs": [], @@ -24,12 +28,15 @@ "import sys\n", "sys.path.append(\"..\")\n", "# this is the convenience function\n", - "from autokoopman import auto_koopman" + "from autokoopman import auto_koopman\n", + "import autokoopman as ak\n", + "\n", + "np.random.seed(1234)" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "291d3409-1c8c-44cb-8380-44f08019b57d", "metadata": {}, "outputs": [], @@ -41,12 +48,27 @@ " initial_states=np.random.uniform(low=-2.0, high=2.0, size=(10, 2)),\n", " tspan=[0.0, 6.0],\n", " sampling_period=0.1\n", - ")" + ")\n", + "\n", + "# add garbage states -- we will weight these values to zero\n", + "training_data = ak.TrajectoriesData({\n", + " key: ak.Trajectory(t.times, np.hstack([t.states, np.random.rand() * np.ones((len(t.states), 3))]), t.inputs) for key, t in training_data._trajs.items()\n", + "})" + ] + }, + { + "cell_type": "markdown", + "id": "96e78c86-54a3-4aac-bb55-0cbe51eaee31", + "metadata": {}, + "source": [ + "## Weighting Usage\n", + "\n", + "**Weighting has changed from W-eDMD to our new State Weighted eDMD (SW-eDMD) formulation.** Each trajectory is a sequence of $k$ points $[x_0, x_1, \\cdots, x_k]$ $x_i \\in \\mathbb R^n$. So, you will need weights for each time point *and* state $[w_0, w_1, \\cdots, w_k]$ $w_i \\in \\mathbb R^{+n}$." ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "e2d42e41-46c2-467c-9ce7-9bd6a7c509a1", "metadata": {}, "outputs": [], @@ -64,10 +86,13 @@ " # garbage trajectory\n", " trajectories.append(np.random.rand(*traj.states.shape))\n", " \n", - "\n", " # weight good trajectory by its 1 norm\n", " #w = np.sum(traj.abs().states, axis=1)\n", + " #w = 1/(traj.abs().states+1.0)\n", " w = np.ones(traj.states.shape)\n", + " w[:, -3:] = 0.0\n", + " w[:, :2] = 1.0\n", + " w[:, 0] = 1.0\n", " weights.append(w)\n", "\n", " # weight garbage trajectory to zero\n", @@ -81,12 +106,35 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "280b4bb3-4f7d-4a94-a983-663c6255bc83", + "execution_count": 4, + "id": "ddd76415-6d19-4a38-a2b0-84eb48d0fdfc", "metadata": {}, "outputs": [], "source": [ - "weights[1].shape" + "from autokoopman.observable import ReweightedRFFObservable\n", + "import autokoopman.observable as akobs\n", + "\n", + "X, WX = list(zip(*list((trajectories[i], w) for i, w in enumerate(weights))))\n", + "X, WX = np.vstack(X), np.vstack(WX)\n", + "X, WX = np.tile(X, (3, 1)), np.tile(WX, (3, 1))\n", + "idxs = np.random.permutation(np.arange(len(X)))\n", + "Y, WY = X[idxs], WX[idxs]\n", + "\n", + "reweight_obs = akobs.IdentityObservable() | akobs.ReweightedRFFObservable(5, 40, 1.0, X, Y, WX, WY, 'rff')" + ] + }, + { + "cell_type": "markdown", + "id": "a706f212-36cd-4203-b209-cb7c5ce4ad94", + "metadata": {}, + "source": [ + "## Run SW-eDMD\n", + "\n", + "You can run SW-eDMD by turning on the `auto_koopman` flags\n", + "\n", + "* learning_weights=weights -- weights for SW-eDMD, for *estimator method only.*\n", + "* cost_func=\"weighted\" -- uses weighted cost for *tuning only.*\n", + "* scoring_weights=weights-- weights for *tuning only.*" ] }, { @@ -94,7 +142,19 @@ "execution_count": null, "id": "98510aa7-3416-4181-a493-00500be53f61", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Tuning GridSearchTuner: 0%| | 0/40 [00:00