legendre_expansion

Bases: transformation

The legendre expansion function.

Applies Legendre polynomial expansion to input data.

Notes

The Legendre polynomials, named after mathematician Adrien-Marie Legendre, are defined as an orthogonal system over the interval \([-1, 1]\), where the polynomial term \(P_n(x)\) of degree \(n\) satisfies the following equation:

\[
    \begin{equation}
    \int_{-1}^{+1} P_m(x) P_n(x) dx = 0, \text{ if } m \neq n.
    \end{equation}
\]

Specifically, according to Bonnet's formula, the Legendre polynomials can be recursively represented as follows:

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

\[
    \begin{equation}
    P_0(x) = 1 \text{, and } P_1(x) = x.
    \end{equation}
\]

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

\[
    \begin{equation}
    P_n(x) = \frac{x(2n-1) P_{n-1}(x) - (n-1) P_{n-2}(x) }{n}
    \end{equation}
\]

The Legendre polynomials help define the data expansion function as follows:

\[
    \begin{equation}
    \kappa(\mathbf{x} | d) = \left[ P_1(\mathbf{x}), P_2(\mathbf{x}), \cdots, P_d(\mathbf{x}) \right] \in R^D,
    \end{equation}
\]

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

Attributes:

Name	Type	Description
d	int	The degree of Legendre polynomial expansion.

Methods:

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

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

class legendre_expansion(transformation):
    r"""
        The legendre expansion function.

        Applies Legendre polynomial expansion to input data.

        Notes
        ----------
        The Legendre polynomials, named after mathematician Adrien-Marie Legendre, are defined as an orthogonal system over the interval $[-1, 1]$, where the polynomial term $P_n(x)$ of degree $n$ satisfies the following equation:

        $$
            \begin{equation}
            \int_{-1}^{+1} P_m(x) P_n(x) dx = 0, \text{ if } m \neq n.
            \end{equation}
        $$

        Specifically, according to Bonnet's formula, the Legendre polynomials can be recursively represented as follows:

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

        $$
            \begin{equation}
            P_0(x) = 1 \text{, and } P_1(x) = x.
            \end{equation}
        $$

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

        $$
            \begin{equation}
            P_n(x) = \frac{x(2n-1) P_{n-1}(x) - (n-1) P_{n-2}(x) }{n}
            \end{equation}
        $$

        The Legendre polynomials help define the data expansion function as follows:

        $$
            \begin{equation}
            \kappa(\mathbf{x} | d) = \left[ P_1(\mathbf{x}), P_2(\mathbf{x}), \cdots, P_d(\mathbf{x}) \right] \in R^D,
            \end{equation}
        $$

        where the output dimension $D = md$.

        Attributes
        ----------
        d : int
            The degree of Legendre polynomial expansion.

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

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

    def calculate_D(self, m: int):
        """
            Calculates the output dimension after Legendre 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 Legendre 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] = x
        # high-order cases
        for n in range(2, self.d + 1):
            expansion[:, :, n] = (2*n-1)/n * x * expansion[:, :, n-1].clone() - (n-1)/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='legendre_polynomial_expansion', d=2, *args, **kwargs)

Initializes the Legendre polynomial expansion transformation.

Parameters:

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

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py

def __init__(self, name='legendre_polynomial_expansion', d: int = 2, *args, **kwargs):
    """
        Initializes the Legendre polynomial expansion transformation.

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

calculate_D(m)

Calculates the output dimension after Legendre 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 Legendre polynomial expansion.

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

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

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

Performs Legendre 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 Legendre 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] = x
    # high-order cases
    for n in range(2, self.d + 1):
        expansion[:, :, n] = (2*n-1)/n * x * expansion[:, :, n-1].clone() - (n-1)/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)