gegenbauer_expansion

Bases: transformation

The gegenbauer expansion function.

Applies Gegenbauer polynomial expansion to input data.

Notes

The Gegenbauer polynomials, named after mathematician Leopold Gegenbauer, are orthogonal polynomials that generalize both the Legendre and Chebyshev polynomials, and are special cases of Jacobi polynomials.

Formally, the Gegenbauer polynomials are particular solutions of the Gegenbauer differential equation:

\[
    \begin{equation}
    (1 - x^2) y'' - (2 \alpha + 1) x y' + d(d+2 \alpha) y = 0,
    \end{equation}
\]

where \(y = y(x)\) is a function of variable \(x\) and \(d \in N\) is a non-negative integer.

When \(\alpha = \frac{1}{2}\), the Gegenbauer polynomials reduce to the Legendre polynomials introduced earlier; when \(\alpha = 1\), they reduce to the Chebyshev polynomials of the second kind.

The Gegenbauer polynomials can be recursively defined as follows:

Base cases \(n=0\) and \(n=1\):

\[
    \begin{equation}
    P^{(\alpha)}_0(x) = 1 \text{, and } P^{(\alpha)}_1(x) = 2 \alpha x.
    \end{equation}
\]

High-order cases with degree \(n \ge 2\):

\[
    \begin{equation}
    P^{(\alpha)}_n(x) = \frac{2x(n-1+\alpha) P^{(\alpha)}_{n-1}(x) - (n+2\alpha -2) P^{(\alpha)}_{n-2}(x) }{n}
    \end{equation}
\]

Based on the Gegenbauer polynomials, we can define the expansion function as follows:

\[
    \begin{equation}
    \kappa(\mathbf{x} | d, \alpha) = \left[ P^{(\alpha)}_1(\mathbf{x}), P^{(\alpha)}_2(\mathbf{x}), \cdots, P^{(\alpha)}_d(\mathbf{x}) \right] \in R^D,
    \end{equation}
\]

where the output dimension \(D = md\).

Attributes:

Name	Type	Description
d	int	The degree of Gegenbauer polynomial expansion.
alpha	float	Parameter controlling the Gegenbauer polynomial.

Methods:

Name	Description
calculate_D	Calculates the output dimension after expansion.
forward	Performs Gegenbauer polynomial expansion on the input tensor.

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

class gegenbauer_expansion(transformation):
    r"""
        The gegenbauer expansion function.

        Applies Gegenbauer polynomial expansion to input data.

        Notes
        ----------

        The Gegenbauer polynomials, named after mathematician Leopold Gegenbauer, are orthogonal polynomials that generalize both the Legendre and Chebyshev polynomials, and are special cases of Jacobi polynomials.

        Formally, the Gegenbauer polynomials are particular solutions of the Gegenbauer differential equation:

        $$
            \begin{equation}
            (1 - x^2) y'' - (2 \alpha + 1) x y' + d(d+2 \alpha) y = 0,
            \end{equation}
        $$

        where $y = y(x)$ is a function of variable $x$ and $d \in N$ is a non-negative integer.

        When $\alpha = \frac{1}{2}$, the Gegenbauer polynomials reduce to the Legendre polynomials introduced earlier; when $\alpha = 1$, they reduce to the Chebyshev polynomials of the second kind.

        The Gegenbauer polynomials can be recursively defined as follows:

        __Base cases $n=0$ and $n=1$:__

        $$
            \begin{equation}
            P^{(\alpha)}_0(x) = 1 \text{, and } P^{(\alpha)}_1(x) = 2 \alpha x.
            \end{equation}
        $$

        __High-order cases with degree $n \ge 2$:__

        $$
            \begin{equation}
            P^{(\alpha)}_n(x) = \frac{2x(n-1+\alpha) P^{(\alpha)}_{n-1}(x) - (n+2\alpha -2) P^{(\alpha)}_{n-2}(x) }{n}
            \end{equation}
        $$

        Based on the Gegenbauer polynomials, we can define the expansion function as follows:

        $$
            \begin{equation}
            \kappa(\mathbf{x} | d, \alpha) = \left[ P^{(\alpha)}_1(\mathbf{x}), P^{(\alpha)}_2(\mathbf{x}), \cdots, P^{(\alpha)}_d(\mathbf{x}) \right] \in R^D,
            \end{equation}
        $$

        where the output dimension $D = md$.


        Attributes
        ----------
        d : int
            The degree of Gegenbauer polynomial expansion.
        alpha : float
            Parameter controlling the Gegenbauer polynomial.

        Methods
        -------
        calculate_D(m: int)
            Calculates the output dimension after expansion.
        forward(x: torch.Tensor, device='cpu', *args, **kwargs)
            Performs Gegenbauer polynomial expansion on the input tensor.
    """
    def __init__(self, name='gegenbauer_polynomial_expansion', d: int = 2, alpha: float = 1.0, *args, **kwargs):
        """
            Initializes the Gegenbauer polynomial expansion transformation.

            Parameters
            ----------
            name : str, optional
                Name of the transformation. Defaults to 'gegenbauer_polynomial_expansion'.
            d : int, optional
                The maximum order of Gegenbauer polynomials for expansion. Defaults to 2.
            alpha : float, optional
                The alpha parameter for Gegenbauer polynomials. Defaults to 1.0.
            *args : tuple
                Additional positional arguments.
            **kwargs : dict
                Additional keyword arguments.
        """
        super().__init__(name=name, *args, **kwargs)
        self.d = d
        self.alpha = alpha

    def calculate_D(self, m: int):
        """
            Calculates the output dimension after Gegenbauer polynomial expansion.

            Parameters
            ----------
            m : int
                Input dimension.

            Returns
            -------
            int
                Output dimension after expansion.
        """
        return m * self.d

    def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
        """
            Performs Gegenbauer polynomial expansion on the input tensor.

            Parameters
            ----------
            x : torch.Tensor
                Input tensor of shape `(batch_size, input_dim)`.
            device : str, optional
                Device for computation ('cpu', 'cuda'). Defaults to 'cpu'.
            *args : tuple
                Additional positional arguments.
            **kwargs : dict
                Additional keyword arguments.

            Returns
            -------
            torch.Tensor
                Expanded tensor of shape `(batch_size, expanded_dim)`.

            Raises
            ------
            AssertionError
                If the output tensor shape does not match the expected dimensions.
        """
        b, m = x.shape
        x = self.pre_process(x=x, device=device)

        # base case: order 0
        expansion = torch.ones(size=[x.size(0), x.size(1), self.d + 1]).to(device)
        # base case: order 1
        if self.d > 0:
            expansion[:, :, 1] = 2 * self.alpha * x
        # high-order cases
        for n in range(2, self.d + 1):
            expansion[:, :, n] = (n-1+self.alpha)/n * 2*x * expansion[:, :, n-1].clone() - (n-2+2*self.alpha)/n * expansion[:, :, n-2].clone()

        expansion = expansion[:, :, 1:].contiguous().view(x.size(0), -1)

        assert expansion.shape == (b, self.calculate_D(m=m))
        return self.post_process(x=expansion, device=device)

__init__(name='gegenbauer_polynomial_expansion', d=2, alpha=1.0, *args, **kwargs)

Initializes the Gegenbauer polynomial expansion transformation.

Parameters:

Name	Type	Description	Default
name	str	Name of the transformation. Defaults to 'gegenbauer_polynomial_expansion'.	'gegenbauer_polynomial_expansion'
d	int	The maximum order of Gegenbauer polynomials for expansion. Defaults to 2.	2
alpha	float	The alpha parameter for Gegenbauer polynomials. Defaults to 1.0.	1.0
*args	tuple	Additional positional arguments.	()
**kwargs	dict	Additional keyword arguments.	{}

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

def __init__(self, name='gegenbauer_polynomial_expansion', d: int = 2, alpha: float = 1.0, *args, **kwargs):
    """
        Initializes the Gegenbauer polynomial expansion transformation.

        Parameters
        ----------
        name : str, optional
            Name of the transformation. Defaults to 'gegenbauer_polynomial_expansion'.
        d : int, optional
            The maximum order of Gegenbauer polynomials for expansion. Defaults to 2.
        alpha : float, optional
            The alpha parameter for Gegenbauer polynomials. Defaults to 1.0.
        *args : tuple
            Additional positional arguments.
        **kwargs : dict
            Additional keyword arguments.
    """
    super().__init__(name=name, *args, **kwargs)
    self.d = d
    self.alpha = alpha

calculate_D(m)

Calculates the output dimension after Gegenbauer polynomial expansion.

Parameters:

Name	Type	Description	Default
m	int	Input dimension.	required

Returns:

Type	Description
int	Output dimension after expansion.

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

def calculate_D(self, m: int):
    """
        Calculates the output dimension after Gegenbauer polynomial expansion.

        Parameters
        ----------
        m : int
            Input dimension.

        Returns
        -------
        int
            Output dimension after expansion.
    """
    return m * self.d

forward(x, device='cpu', *args, **kwargs)

Performs Gegenbauer polynomial expansion on the input tensor.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor of shape (batch_size, input_dim).	required
device	str	Device for computation ('cpu', 'cuda'). Defaults to 'cpu'.	'cpu'
*args	tuple	Additional positional arguments.	()
**kwargs	dict	Additional keyword arguments.	{}

Returns:

Type	Description
Tensor	Expanded tensor of shape (batch_size, expanded_dim).

Raises:

Type	Description
AssertionError	If the output tensor shape does not match the expected dimensions.

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
    """
        Performs Gegenbauer polynomial expansion on the input tensor.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape `(batch_size, input_dim)`.
        device : str, optional
            Device for computation ('cpu', 'cuda'). Defaults to 'cpu'.
        *args : tuple
            Additional positional arguments.
        **kwargs : dict
            Additional keyword arguments.

        Returns
        -------
        torch.Tensor
            Expanded tensor of shape `(batch_size, expanded_dim)`.

        Raises
        ------
        AssertionError
            If the output tensor shape does not match the expected dimensions.
    """
    b, m = x.shape
    x = self.pre_process(x=x, device=device)

    # base case: order 0
    expansion = torch.ones(size=[x.size(0), x.size(1), self.d + 1]).to(device)
    # base case: order 1
    if self.d > 0:
        expansion[:, :, 1] = 2 * self.alpha * x
    # high-order cases
    for n in range(2, self.d + 1):
        expansion[:, :, n] = (n-1+self.alpha)/n * 2*x * expansion[:, :, n-1].clone() - (n-2+2*self.alpha)/n * expansion[:, :, n-2].clone()

    expansion = expansion[:, :, 1:].contiguous().view(x.size(0), -1)

    assert expansion.shape == (b, self.calculate_D(m=m))
    return self.post_process(x=expansion, device=device)