Source code for skpoly._bernstein

import numpy as np
from scipy.stats import binom

from ._base import _BasePolynomialBasisTransformer




class BernsteinFeatures(_BasePolynomialBasisTransformer):
    """Generate Bernstein polynomial features.

    Parameters
    ----------
    degree : int, default=3
        Highest degree of the basis.
    feature_range : tuple of float, default=(0.0, 1.0)
        Interval ``[a, b]`` that bounds the input features.
    include_bias : bool, default=False
        If ``True``, include the first basis function of each feature.
    tensor_product : bool, default=False
        If ``True``, append tensor product features for each pair of columns.

    Attributes
    ----------
    n_features_in_ : int
        Number of features in the input passed to :meth:`fit`.
    feature_range_ : tuple of float
        Validated feature range.
    n_output_features_ : int
        Total number of output features generated by :meth:`transform`.
    """

    def _setup(self) -> None:
        return None

    def _scale_to_basis(self, X: np.ndarray) -> np.ndarray:
        lower, upper = self.feature_range_
        scaled = (X - lower) / (upper - lower)
        return np.clip(scaled, 0.0, 1.0, out=scaled)

    def _evaluate_basis(self, X: np.ndarray) -> np.ndarray:
        if X.shape[1] == 0:
            return np.empty((X.shape[0], 0, self.degree + 1), dtype=np.float64)

        orders = np.arange(self.degree + 1)
        vanders = []
        for column in range(X.shape[1]):
            probabilities = binom.pmf(orders, self.degree, X[:, [column]])
            vanders.append(probabilities.astype(np.float64, copy=False))

        return np.stack(vanders, axis=1)



__all__ = ["BernsteinFeatures"]