From 953822b12622c2e671aab21e0607c5c5d35d5fe8 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] 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
+    )