# Path: formulaic/transforms/poly.py
# @stateful_transform
# def poly(
#     x: numpy.typing.ArrayLike, degree: int = 1, raw: bool = False, _state=None
# ) -> numpy.ndarray:
#     """
#     Generate a basis for a polynomial vector-space representation of `x`.
# 
#     The basis vectors returned by this transform can be used, for example, to
#     capture non-linear dependence on `x` in a linear regression.
# 
#     Args:
#         x: The vector for which a polynomial vector space should be generated.
#         degree: The degree of the polynomial vector space.
#         raw: Whether to return "raw" basis vectors (e.g. `[x, x**2, x**3]`). If
#             `False`, an orthonormal set of basis vectors is returned instead
#             (see notes below for more information).
# 
#     Returns:
#         A two-dimensional numpy array with `len(x)` rows, and `degree` columns.
#         The columns represent the basis vectors of the polynomial vector-space.
# 
#     Notes:
#         This transform is an implementation of the "three-term recurrence
#         relation" for monic orthogonal polynomials. There are many good
#         introductions to these recurrence relations, including:
#             https://dec41.user.srcf.net/h/IB_L/numerical_analysis/2_3
#         Another common approach is QR factorisation, where the columns of Q are
#         the orthogonal basis vectors. However, our implementation outperforms
#         numpy's QR decomposition, and does not require needless computation of
#         the R matrix. It should also be noted that orthogonal polynomial bases
#         are unique up to the choice of inner-product and scaling, and so all
#         methods will result in the same set of polynomials.
# 
#         When used as a stateful transform, we retain the coefficients that
#         uniquely define the polynomials; and so new data will be evaluated
#         against the same polynomial bases as the original dataset. However,
#         the polynomial basis will almost certainly *not* be orthogonal for the
#         new data. This is because changing the incoming dataset is equivalent to
#         changing your choice of inner product.
# 
#         Using orthogonal basis vectors (as compared to the "raw" vectors) allows
#         you to increase the degree of the polynomial vector space without
#         affecting the coefficients of lower-order components in a linear
#         regression. This stability is often attractive during exploratory data
#         analysis, but does not otherwise change the results of a linear
#         regression.
# 
#         `nan` values in `x` will be ignored and progagated through to generated
#         polynomials.
# 
#         The signature of this transform is intentionally chosen to be compatible
#         with R.
#     """
# 
#     if raw:
#         return numpy.stack([numpy.power(x, k) for k in range(1, degree + 1)], axis=1)
# 
#     x = numpy.array(x)
# 
#     # Check if we already have already generated the alpha and beta coefficients.
#     # If not, we enter "training" mode.
#     training = False
#     alpha = _state.get("alpha")
#     norms2 = _state.get("norms2")
# 
#     if alpha is None:
#         training = True
#         alpha = {}
#         norms2 = {}
# 
#     # Build polynomials iteratively using the monic three-term recurrence relation
#     # Note that alpha and beta are fixed if not in "training" mode.
#     P = numpy.empty((x.shape[0], degree + 1))
#     P[:, 0] = 1
# 
#     def get_alpha(k):
#         if training and k not in alpha:
#             alpha[k] = numpy.sum(x * P[:, k] ** 2) / numpy.sum(P[:, k] ** 2)
#         return alpha[k]
# 
#     def get_norm(k):
#         if training and k not in norms2:
#             norms2[k] = numpy.sum(P[:, k] ** 2)
#         return norms2[k]
# 
#     def get_beta(k):
#         return get_norm(k) / get_norm(k - 1)
# 
#     for i in range(1, degree + 1):
#         P[:, i] = (x - get_alpha(i - 1)) * P[:, i - 1]
#         if i >= 2:
#             P[:, i] -= get_beta(i - 1) * P[:, i - 2]
# 
#     # Renormalize so we provide an orthonormal basis.
#     P /= numpy.array([numpy.sqrt(get_norm(k)) for k in range(0, degree + 1)])
# 
#     if training:
#         _state["alpha"] = alpha
#         _state["norms2"] = norms2
# 
#     # Return basis dropping the first (constant) column
#     return FactorValues(
#         P[:, 1:], column_names=tuple(str(i) for i in range(1, degree + 1))
#     )


import numpy
import pytest
from formulaic.transforms.poly import poly
            -2.883000e-01,
            -2.522625e-01,
            -2.162250e-01,
            -1.801875e-01,
            -1.441500e-01,
            -1.081125e-01,
            -7.207500e-02,
            -3.603750e-02,
            -3.000725e-17,
            3.603750e-02,
            7.207500e-02,
            1.081125e-01,
            1.441500e-01,
            1.801875e-01,
            2.162250e-01,
            2.522625e-01,
            2.883000e-01,
            3.243375e-01,
            3.603750e-01,
        ]

        assert numpy.allclose(
            V[:, 0],
            r_reference,
        )

        assert pytest.approx(state["alpha"], {0: 0.5})
        assert pytest.approx(state["norms2"], {0: 21.0, 2: 1.925})

        assert numpy.allclose(

# Path: formulaic/transforms/poly.py
# @stateful_transform
# def poly(
#     x: numpy.typing.ArrayLike, degree: int = 1, raw: bool = False, _state=None
# ) -> numpy.ndarray:
#     """
#     Generate a basis for a polynomial vector-space representation of `x`.
# 
#     The basis vectors returned by this transform can be used, for example, to
#     capture non-linear dependence on `x` in a linear regression.
# 
#     Args:
#         x: The vector for which a polynomial vector space should be generated.
#         degree: The degree of the polynomial vector space.
#         raw: Whether to return "raw" basis vectors (e.g. `[x, x**2, x**3]`). If
#             `False`, an orthonormal set of basis vectors is returned instead
#             (see notes below for more information).
# 
#     Returns:
#         A two-dimensional numpy array with `len(x)` rows, and `degree` columns.
#         The columns represent the basis vectors of the polynomial vector-space.
# 
#     Notes:
#         This transform is an implementation of the "three-term recurrence
#         relation" for monic orthogonal polynomials. There are many good
#         introductions to these recurrence relations, including:
#             https://dec41.user.srcf.net/h/IB_L/numerical_analysis/2_3
#         Another common approach is QR factorisation, where the columns of Q are
#         the orthogonal basis vectors. However, our implementation outperforms
#         numpy's QR decomposition, and does not require needless computation of
#         the R matrix. It should also be noted that orthogonal polynomial bases
#         are unique up to the choice of inner-product and scaling, and so all
#         methods will result in the same set of polynomials.
# 
#         When used as a stateful transform, we retain the coefficients that
#         uniquely define the polynomials; and so new data will be evaluated
#         against the same polynomial bases as the original dataset. However,
#         the polynomial basis will almost certainly *not* be orthogonal for the
#         new data. This is because changing the incoming dataset is equivalent to
#         changing your choice of inner product.
# 
#         Using orthogonal basis vectors (as compared to the "raw" vectors) allows
#         you to increase the degree of the polynomial vector space without
#         affecting the coefficients of lower-order components in a linear
#         regression. This stability is often attractive during exploratory data
#         analysis, but does not otherwise change the results of a linear
#         regression.
# 
#         `nan` values in `x` will be ignored and progagated through to generated
#         polynomials.
# 
#         The signature of this transform is intentionally chosen to be compatible
#         with R.
#     """
# 
#     if raw:
#         return numpy.stack([numpy.power(x, k) for k in range(1, degree + 1)], axis=1)
# 
#     x = numpy.array(x)
# 
#     # Check if we already have already generated the alpha and beta coefficients.
#     # If not, we enter "training" mode.
#     training = False
#     alpha = _state.get("alpha")
#     norms2 = _state.get("norms2")
# 
#     if alpha is None:
#         training = True
#         alpha = {}
#         norms2 = {}
# 
#     # Build polynomials iteratively using the monic three-term recurrence relation
#     # Note that alpha and beta are fixed if not in "training" mode.
#     P = numpy.empty((x.shape[0], degree + 1))
#     P[:, 0] = 1
# 
#     def get_alpha(k):
#         if training and k not in alpha:
#             alpha[k] = numpy.sum(x * P[:, k] ** 2) / numpy.sum(P[:, k] ** 2)
#         return alpha[k]
# 
#     def get_norm(k):
#         if training and k not in norms2:
#             norms2[k] = numpy.sum(P[:, k] ** 2)
#         return norms2[k]
# 
#     def get_beta(k):
#         return get_norm(k) / get_norm(k - 1)
# 
#     for i in range(1, degree + 1):
#         P[:, i] = (x - get_alpha(i - 1)) * P[:, i - 1]
#         if i >= 2:
#             P[:, i] -= get_beta(i - 1) * P[:, i - 2]
# 
#     # Renormalize so we provide an orthonormal basis.
#     P /= numpy.array([numpy.sqrt(get_norm(k)) for k in range(0, degree + 1)])
# 
#     if training:
#         _state["alpha"] = alpha
#         _state["norms2"] = norms2
# 
#     # Return basis dropping the first (constant) column
#     return FactorValues(
#         P[:, 1:], column_names=tuple(str(i) for i in range(1, degree + 1))
#     )


# Path: tests/transforms/test_poly.py
import numpy
import pytest
from formulaic.transforms.poly import poly

            -2.883000e-01,
            -2.522625e-01,
            -2.162250e-01,
            -1.801875e-01,
            -1.441500e-01,
            -1.081125e-01,
            -7.207500e-02,
            -3.603750e-02,
            -3.000725e-17,
            3.603750e-02,
            7.207500e-02,
            1.081125e-01,
            1.441500e-01,
            1.801875e-01,
            2.162250e-01,
            2.522625e-01,
            2.883000e-01,
            3.243375e-01,
            3.603750e-01,
        ]

        assert numpy.allclose(
            V[:, 0],
            r_reference,
        )

        assert pytest.approx(state["alpha"], {0: 0.5})
        assert pytest.approx(state["norms2"], {0: 21.0, 2: 1.925})

        assert numpy.allclose(

            poly(data, _state=state)[:, 0],