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/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py index dad0aa8c9..6e445875e 100644 --- a/pysindy/optimizers/trapping_sr3.py +++ b/pysindy/optimizers/trapping_sr3.py @@ -1,12 +1,18 @@ import warnings +from itertools import combinations as combo_nr +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 +819,93 @@ def _reduce(self, x, y): self.coef_ = coef_sparse.T self.objective_history = objective_history + + +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 + 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. + """ + 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_)] + + # 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 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 + ) diff --git a/test/test_optimizers.py b/test/test_optimizers.py index 171825c4b..c9417a1f1 100644 --- a/test/test_optimizers.py +++ b/test/test_optimizers.py @@ -32,6 +32,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 @@ -1133,6 +1135,26 @@ def test_remove_and_decrement(): np.testing.assert_array_equal(expected, result) +@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=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} + constraint_lhs = _antisymm_triple_constraints(3, 3, mixed_terms) + result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0))) + assert result[0] == 0 + + def test_pickle(data_lorenz): x, t = data_lorenz y = PolynomialLibrary(degree=2).fit_transform(x)