From 7d4dd9873624875e52dde64a372340d810bc0b9b Mon Sep 17 00:00:00 2001 From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com> Date: Tue, 5 Dec 2023 17:42:43 +0000 Subject: [PATCH 1/6] ENH: Add function to calculate number of polynomial features This is much faster than fitting the features, and can be used by trapping to replace some constraint index calculations. Also: Add docstring/type annotations Rename a variable so it doesn't shadow newly-imported name --- pysindy/feature_library/polynomial_library.py | 66 ++++++++++++++++--- test/test_feature_library.py | 4 ++ 2 files changed, 62 insertions(+), 8 deletions(-) diff --git a/pysindy/feature_library/polynomial_library.py b/pysindy/feature_library/polynomial_library.py index 75dbf5637..30710041a 100644 --- a/pysindy/feature_library/polynomial_library.py +++ b/pysindy/feature_library/polynomial_library.py @@ -1,7 +1,9 @@ from itertools import chain +from math import comb from typing import Iterator import numpy as np +from numpy.typing import NDArray from scipy import sparse from sklearn.preprocessing import PolynomialFeatures from sklearn.utils.validation import check_is_fitted @@ -73,8 +75,22 @@ def __init__( @staticmethod def _combinations( - n_features, degree, include_interaction, interaction_only, include_bias - ) -> Iterator[tuple]: + n_features: int, + degree: int, + include_interaction: bool, + interaction_only: bool, + include_bias: bool, + ) -> Iterator[tuple[int, ...]]: + """ + Create selection tuples of input indexes for each polynomail term + + Selection tuple iterates the input indexes present in a single term. + For example, (x+y+1)^2 would be in iterator of the tuples: + (), (0,), (1,), (0, 0), (0, 1), (1, 1) + 1 x y x^2 x*y y^2 + + Order of terms is preserved regardless of include_interation. + """ if not include_interaction: return chain( [()] if include_bias else [], @@ -93,7 +109,12 @@ def _combinations( ) @property - def powers_(self): + def powers_(self) -> NDArray[np.dtype("int")]: + """ + The exponents of the polynomial as an array of shape + (n_features_out, n_features_in), where each item is the exponent of the + jth input variable in the ith polynomial term. + """ check_is_fitted(self) combinations = self._combinations( n_features=self.n_features_in_, @@ -208,10 +229,10 @@ def transform(self, x_full): ) if sparse.isspmatrix(x): columns = [] - for comb in combinations: - if comb: + for combo in combinations: + if combo: out_col = 1 - for col_idx in comb: + for col_idx in combo: out_col = x[..., col_idx].multiply(out_col) columns.append(out_col) else: @@ -227,7 +248,36 @@ def transform(self, x_full): ), x.__dict__, ) - for i, comb in enumerate(combinations): - xp[..., i] = x[..., comb].prod(-1) + for i, combo in enumerate(combinations): + xp[..., i] = x[..., combo].prod(-1) xp_full = xp_full + [xp] return xp_full + + +def n_poly_features( + n_in_feat: int, + degree: int, + include_bias: bool = False, + include_interation: bool = True, + interaction_only: bool = False, +) -> int: + """Calculate number of polynomial features + + Args: + n_in_feat: number of input features, e.g. 3 for x, y, z + degree: polynomial degree, e.g. 2 for quadratic + include_bias: whether to include a constant term + include_interaction: whether to include terms mixing multiple inputs + interaction_only: whether to omit terms of x_m * x_n^p for p > 1 + """ + if not include_interation and interaction_only: + raise ValueError("Cannot set interaction only if include_interaction is False") + n_feat = include_bias + if not include_interation: + return n_feat + n_in_feat * degree + for deg in range(1, degree + 1): + if interaction_only: + n_feat += comb(n_in_feat, deg) + else: + n_feat += comb(n_in_feat + deg - 1, deg) + return n_feat diff --git a/test/test_feature_library.py b/test/test_feature_library.py index 8e98b1a0d..8bf27c580 100644 --- a/test/test_feature_library.py +++ b/test/test_feature_library.py @@ -21,6 +21,7 @@ from pysindy.feature_library import TensoredLibrary from pysindy.feature_library import WeakPDELibrary from pysindy.feature_library.base import BaseFeatureLibrary +from pysindy.feature_library.polynomial_library import n_poly_features from pysindy.optimizers import SINDyPI from pysindy.optimizers import STLSQ @@ -880,3 +881,6 @@ def test_polynomial_combinations(include_interaction, interaction_only, bias, ex ) result = tuple(sorted(list(combos))) assert result == expected + assert len(expected) == n_poly_features( + 2, 2, bias, include_interaction, interaction_only + ) From 65304b2b10189c536ca06754b7c9d9f7ebcd48e1 Mon Sep 17 00:00:00 2001 From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com> Date: Fri, 8 Dec 2023 16:47:07 +0000 Subject: [PATCH 2/6] ENH: add _make_constraints to TrappingSR3 There are several major differences from the version in example 8: * Coupled with PolynomialLibrary logic using PolynomialLibrary.powers_. Previously, this coupling was implicit and actually tied to the specific order in Example 8, not the order users would expect. * Instead of handling pre-flattened 2D array, constraints are built with target and feature axes separately, letting the caller handle formatting the constraints as "target" or "feature". This gives us flexibility to later clean up the constraint logic outside this function without needing to return here and removes the need for striding calculations like: n_tgts * (n_tgts + k - 1) + i + n_tgts * int(j * (2 * n_tgts - j - 3) / 2.0) Since constraints are built directly from the terms and exponents in the library, its much easier to read and debug the code. * Renamed "r" as "n_tgts", etc. Single-letter index variables are removed. When they occur in homogenous iterators, they are concatenated with what they're iterating (e.g. tgt_i, tgt_j) * Broke out the different kinds of constraints into helper functions to make it easier to test. This allows follow on work to auto-attach constraints to a problem. --- pysindy/optimizers/trapping_sr3.py | 94 ++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py index dad0aa8c9..def9d71ee 100644 --- a/pysindy/optimizers/trapping_sr3.py +++ b/pysindy/optimizers/trapping_sr3.py @@ -1,12 +1,19 @@ import warnings +from itertools import combinations as combo_nr +from itertools import combinations_with_replacement as combo_wr +from math import comb from typing import Tuple import cvxpy as cp import numpy as np +from numpy import intp +from numpy.typing import NDArray from scipy.linalg import cho_factor from scipy.linalg import cho_solve from sklearn.exceptions import ConvergenceWarning +from ..feature_library.polynomial_library import n_poly_features +from ..feature_library.polynomial_library import PolynomialLibrary from ..utils import reorder_constraints from .sr3 import SR3 @@ -813,3 +820,90 @@ def _reduce(self, x, y): self.coef_ = coef_sparse.T self.objective_history = objective_history + + +def _make_constraints(n_tgts: int): + """Create constraints for the Quadratic terms in TrappingSR3. + + These are the constraints from equation 5 of the Trapping SINDy paper. + + Args: + n_tgts: number of coordinates or modes for which you're fitting an ODE. + + Returns: + A tuple of the constraint zeros, and a constraint matrix to multiply + by the coefficient matrix of Polynomial terms. Number of constraints is + ``n_tgts + 2 * math.comb(n_tgts, 2) + math.comb(n_tgts, 3)``. + Constraint matrix is of shape ``(n_constraint, n_feature, n_tgt)``. + To get "feature" order constraints, use + ``np.reshape(constraint_matrix, (n_constraints, -1))``. + To get "target" order constraints, transpose axis 1 and 2 before + reshaping. + """ + n_terms = n_poly_features(n_tgts, degree=2, include_bias=False) + lib = PolynomialLibrary(2, include_bias=False).fit(np.zeros((1, n_tgts))) + terms = [(t_ind, exps) for t_ind, exps in enumerate(lib.powers_)] + + # index of tgt -> index of its pure quadratic term + pure_terms = {np.argmax(exps): t_ind for t_ind, exps in terms if max(exps) == 2} + # two indexes of tgts -> index of their mixed quadratic term + mixed_terms = { + frozenset(np.argwhere(exponent == 1).flatten()): t_ind + for t_ind, exponent in terms + if max(exponent) == 1 and sum(exponent) == 2 + } + constraint_mat = np.vstack( + ( + _pure_constraints(n_tgts, n_terms, pure_terms), + _antisymm_double_constraint(n_tgts, n_terms, pure_terms, mixed_terms), + _antisymm_triple_constraints(n_tgts, n_terms, mixed_terms), + ) + ) + + return np.zeros(len(constraint_mat)), constraint_mat + + +def _pure_constraints( + n_tgts: int, n_terms: int, pure_terms: dict[intp, int] +) -> NDArray[np.dtype("float")]: + """Set constraints for coefficients adorning terms like a_i^3 = 0""" + constraint_mat = np.zeros((n_tgts, n_terms, n_tgts)) + for constr_ind, (tgt_ind, term_ind) in zip(range(n_tgts), pure_terms.items()): + constraint_mat[constr_ind, term_ind, tgt_ind] = 1.0 + return constraint_mat + + +def _antisymm_double_constraint( + n_tgts: int, + n_terms: int, + pure_terms: dict[intp, int], + mixed_terms: dict[frozenset[intp] : int], +) -> NDArray[np.dtype("float")]: + """Set constraints for coefficients adorning terms like a_i^2 * a_j=0""" + constraint_mat_1 = np.zeros((comb(n_tgts, 2), n_terms, n_tgts)) # a_i^2 * a_j + constraint_mat_2 = np.zeros((comb(n_tgts, 2), n_terms, n_tgts)) # a_i * a_j^2 + for constr_ind, ((tgt_i, tgt_j), mix_term) in zip( + range(n_tgts), mixed_terms.items() + ): + constraint_mat_1[constr_ind, mix_term, tgt_i] = 1.0 + constraint_mat_1[constr_ind, pure_terms[tgt_i], tgt_j] = 1.0 + constraint_mat_2[constr_ind, mix_term, tgt_j] = 1.0 + constraint_mat_2[constr_ind, pure_terms[tgt_j], tgt_i] = 1.0 + + return np.concatenate((constraint_mat_1, constraint_mat_2), axis=0) + + +def _antisymm_triple_constraints( + n_tgts: int, n_terms: int, mixed_terms: dict[frozenset[intp] : int] +) -> NDArray[np.dtype("float")]: + constraint_mat = np.zeros((comb(n_tgts, 3), n_terms, n_tgts)) # a_ia_ja_k + + def find_symm_term(a, b): + return mixed_terms[frozenset({a, b})] + + for constr_ind, (tgt_i, tgt_j, tgt_k) in enumerate(combo_nr(range(n_tgts), 3)): + constraint_mat[constr_ind, find_symm_term(tgt_j, tgt_k), tgt_i] = 1 + constraint_mat[constr_ind, find_symm_term(tgt_k, tgt_i), tgt_j] = 1 + constraint_mat[constr_ind, find_symm_term(tgt_i, tgt_j), tgt_k] = 1 + + return constraint_mat From 018eb9a80defc1518f57b73d54789eed17d5199b Mon Sep 17 00:00:00 2001 From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com> Date: Sun, 10 Dec 2023 17:53:58 +0000 Subject: [PATCH 3/6] TST: Add Trapping _make_constraints test --- test/test_optimizers.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/test/test_optimizers.py b/test/test_optimizers.py index c69ce9823..385d809c7 100644 --- a/test/test_optimizers.py +++ b/test/test_optimizers.py @@ -30,6 +30,8 @@ from pysindy.optimizers import TrappingSR3 from pysindy.optimizers import WrappedOptimizer from pysindy.optimizers.stlsq import _remove_and_decrement +from pysindy.optimizers.trapping_sr3 import _antisymm_triple_constraints +from pysindy.optimizers.trapping_sr3 import _make_constraints from pysindy.utils import supports_multiple_targets from pysindy.utils.odes import enzyme @@ -1129,3 +1131,18 @@ def test_remove_and_decrement(): existing_vals=existing_vals, vals_to_remove=vals_to_remove ) np.testing.assert_array_equal(expected, result) + + +def test_trapping_constraints(): + # x, y, x^2, xy, y^2 + constraint_rhs, constraint_lhs = _make_constraints(2) + stable_coefs = np.array([[0, 0, 0, 1, -1], [0, 0, -1, 1, 0]]) + result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0))) + np.testing.assert_array_equal(constraint_rhs, result) + + # xy, xz, yz + stable_coefs = np.array([[0, 0, -1], [0, 0.5, 0], [0.5, 0, 0]]) + mixed_terms = {frozenset((0, 1)): 0, frozenset((0, 2)): 1, frozenset((1, 2)): 2} + constraint_lhs = _antisymm_triple_constraints(3, 3, mixed_terms) + result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0))) + assert result[0] == 0 From 0870f57ab40d7d2ca65f3e24fd85d91ee240ff82 Mon Sep 17 00:00:00 2001 From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com> Date: Sun, 10 Dec 2023 18:41:00 +0000 Subject: [PATCH 4/6] ENH: Allow all polynomial library library options in trapping (other than degree, which has to be 2) --- pysindy/optimizers/trapping_sr3.py | 6 ++++-- test/test_optimizers.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py index def9d71ee..1d8410bcb 100644 --- a/pysindy/optimizers/trapping_sr3.py +++ b/pysindy/optimizers/trapping_sr3.py @@ -822,13 +822,15 @@ def _reduce(self, x, y): self.objective_history = objective_history -def _make_constraints(n_tgts: int): +def _make_constraints(n_tgts: int, **kwargs): """Create constraints for the Quadratic terms in TrappingSR3. These are the constraints from equation 5 of the Trapping SINDy paper. Args: n_tgts: number of coordinates or modes for which you're fitting an ODE. + kwargs: Keyword arguments to PolynomialLibrary such as + ``include_bias``. Returns: A tuple of the constraint zeros, and a constraint matrix to multiply @@ -841,7 +843,7 @@ def _make_constraints(n_tgts: int): reshaping. """ n_terms = n_poly_features(n_tgts, degree=2, include_bias=False) - lib = PolynomialLibrary(2, include_bias=False).fit(np.zeros((1, n_tgts))) + lib = PolynomialLibrary(2, **kwargs).fit(np.zeros((1, n_tgts))) terms = [(t_ind, exps) for t_ind, exps in enumerate(lib.powers_)] # index of tgt -> index of its pure quadratic term diff --git a/test/test_optimizers.py b/test/test_optimizers.py index 385d809c7..ac65049e2 100644 --- a/test/test_optimizers.py +++ b/test/test_optimizers.py @@ -1135,7 +1135,7 @@ def test_remove_and_decrement(): def test_trapping_constraints(): # x, y, x^2, xy, y^2 - constraint_rhs, constraint_lhs = _make_constraints(2) + constraint_rhs, constraint_lhs = _make_constraints(2, include_bias=False) stable_coefs = np.array([[0, 0, 0, 1, -1], [0, 0, -1, 1, 0]]) result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0))) np.testing.assert_array_equal(constraint_rhs, result) From e00724773c491ae382b8371bfc5b247735addd7b Mon Sep 17 00:00:00 2001 From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com> Date: Tue, 2 Jan 2024 16:45:04 +0000 Subject: [PATCH 5/6] CLN: remove unused import --- pysindy/optimizers/trapping_sr3.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py index 1d8410bcb..e9816481e 100644 --- a/pysindy/optimizers/trapping_sr3.py +++ b/pysindy/optimizers/trapping_sr3.py @@ -1,6 +1,5 @@ import warnings from itertools import combinations as combo_nr -from itertools import combinations_with_replacement as combo_wr from math import comb from typing import Tuple From f53151a20749a1664e36dba272b4ff3ec47217c4 Mon Sep 17 00:00:00 2001 From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com> Date: Mon, 8 Jan 2024 13:19:21 +0000 Subject: [PATCH 6/6] BUG(_make_constraints): Request correct num of terms if include_bias --- pysindy/optimizers/trapping_sr3.py | 3 ++- test/test_optimizers.py | 9 +++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py index e9816481e..6e445875e 100644 --- a/pysindy/optimizers/trapping_sr3.py +++ b/pysindy/optimizers/trapping_sr3.py @@ -841,7 +841,8 @@ def _make_constraints(n_tgts: int, **kwargs): To get "target" order constraints, transpose axis 1 and 2 before reshaping. """ - n_terms = n_poly_features(n_tgts, degree=2, include_bias=False) + include_bias = kwargs.get("include_bias", False) + n_terms = n_poly_features(n_tgts, degree=2, include_bias=include_bias) lib = PolynomialLibrary(2, **kwargs).fit(np.zeros((1, n_tgts))) terms = [(t_ind, exps) for t_ind, exps in enumerate(lib.powers_)] diff --git a/test/test_optimizers.py b/test/test_optimizers.py index ac65049e2..462f93a45 100644 --- a/test/test_optimizers.py +++ b/test/test_optimizers.py @@ -1133,13 +1133,18 @@ def test_remove_and_decrement(): np.testing.assert_array_equal(expected, result) -def test_trapping_constraints(): +@pytest.mark.parametrize("include_bias", (True, False)) +def test_trapping_constraints(include_bias): # x, y, x^2, xy, y^2 - constraint_rhs, constraint_lhs = _make_constraints(2, include_bias=False) + constraint_rhs, constraint_lhs = _make_constraints(2, include_bias=include_bias) stable_coefs = np.array([[0, 0, 0, 1, -1], [0, 0, -1, 1, 0]]) + if include_bias: + stable_coefs = np.concatenate(([[0], [0]], stable_coefs), axis=1) result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0))) np.testing.assert_array_equal(constraint_rhs, result) + +def test_trapping_mixed_only(): # xy, xz, yz stable_coefs = np.array([[0, 0, -1], [0, 0.5, 0], [0.5, 0, 0]]) mixed_terms = {frozenset((0, 1)): 0, frozenset((0, 2)): 1, frozenset((1, 2)): 2}