jacobi_expansion

Bases: transformation

The jacobi data expansion function.

It performs the jacobi expansion of the input vector, and returns the expansion result. The class inherits from the base expansion class (i.e., the transformation class in the module directory).

...

Notes

For input vector \(\mathbf{x} \in R^m\), its jacobi expansion up to degree \(d\) can be represented as $$ \begin{equation} \kappa(\mathbf{x} | d) = \left[ P_1^{(\alpha, \beta)}(\mathbf{x}), P_2^{(\alpha, \beta)}(\mathbf{x}), \cdots, P_d^{(\alpha, \beta)}(\mathbf{x})\right] \in {R}^D, \end{equation} $$ where \(P_d^{(\alpha, \beta)}(\mathbf{x})\) denotes the jacobi expansion polynomial of \(\mathbf{x}\) degree \(d\). The output dimension of jacobi expansion will be \(D = m d\).

As to the specific representations of jacobi polynomials, they can be defined recursively based on the lower-degree terms according to the following equations:

(1) Base jacobi polynomials with degree \(d=0\), \(d=1\) and \(d=2\): $$ \begin{align} P_0^{(\alpha, \beta)}(x) &= 1,\\ P_1^{(\alpha, \beta)}(x) &= (\alpha + 1) + (\alpha + \beta + 2) \frac{(x-1)}{2},\\ P_2^{(\alpha, \beta)}(x) &= \frac{(\alpha+1)(\alpha+2)}{2} + (\alpha+2)(\alpha+\beta+3) \frac{x-1}{2} + \frac{(\alpha + \beta + 3)(\alpha + \beta + 4)}{2} \left( \frac{x-1}{2} \right)^2. \end{align} $$

(2) Higher-degree jacobi polynomials with \(d \ge 2\): $$ \begin{align} P_d^{(\alpha, \beta)}(x) &= \frac{(2d + \alpha + \beta -1) \left[ (2d + \alpha + \beta)(2d + \alpha + \beta -2) x + (\alpha^2 - \beta^2) \right]}{2d(d + \alpha + \beta)(2d + \alpha + \beta - 2) } P_{d-1}^{(\alpha, \beta)}(x)\\ & - \frac{2(d+\alpha-1)(d+\beta-1)(2d+\alpha+\beta)}{2d(d + \alpha + \beta)(2d + \alpha + \beta - 2) }P_{d-2}^{(\alpha, \beta)}(x). \end{align} $$

By default, the input and output can also be processed with the optional pre- or post-processing functions in the jacobi expansion function.

Attributes:

Name	Type	Description
name	str, default = 'jacobi_polynomial_expansion'	Name of the expansion function.
d	int, default = 2	Degree of jacobi expansion.
alpha	float, default = 1.0	Parameter of jacobi polynomial representation.
beta	float, default = 1.0	Parameter of jacobi polynomial representation.

Methods:

Name	Description
__init__	It performs the initialization of the expansion function.
calculate_D	It calculates the expansion space dimension D based on the input dimension parameter m.
forward	It implements the abstract forward method declared in the base expansion class.

Source code in tinybig/expansion/recursive_expansion.py

class jacobi_expansion(transformation):
    r"""
    The jacobi data expansion function.

    It performs the jacobi expansion of the input vector, and returns the expansion result.
    The class inherits from the base expansion class (i.e., the transformation class in the module directory).

    ...

    Notes
    ----------
    For input vector $\mathbf{x} \in R^m$, its jacobi expansion up to degree $d$ can be represented as
    $$
        \begin{equation}
            \kappa(\mathbf{x} | d) = \left[ P_1^{(\alpha, \beta)}(\mathbf{x}), P_2^{(\alpha, \beta)}(\mathbf{x}),  \cdots, P_d^{(\alpha, \beta)}(\mathbf{x})\right] \in {R}^D,
        \end{equation}
    $$
    where $P_d^{(\alpha, \beta)}(\mathbf{x})$ denotes the jacobi expansion polynomial of $\mathbf{x}$ degree $d$.
    The output dimension of jacobi expansion will be $D = m d$.

    As to the specific representations of jacobi polynomials, they can be defined recursively based on the
    lower-degree terms according to the following equations:

    (1) **Base jacobi polynomials with degree $d=0$, $d=1$ and $d=2$:**
    $$
        \begin{align}
            P\_0^{(\alpha, \beta)}(x) &= 1,\\\\
            P\_1^{(\alpha, \beta)}(x) &= (\alpha + 1) + (\alpha + \beta + 2) \frac{(x-1)}{2},\\\\
            P\_2^{(\alpha, \beta)}(x) &= \frac{(\alpha+1)(\alpha+2)}{2} + (\alpha+2)(\alpha+\beta+3) \frac{x-1}{2} + \frac{(\alpha + \beta + 3)(\alpha + \beta + 4)}{2} \left( \frac{x-1}{2} \right)^2.
        \end{align}
    $$

    (2) **Higher-degree jacobi polynomials with $d \ge 2$:**
    $$
        \begin{align}
            P\_d^{(\alpha, \beta)}(x) &= \frac{(2d + \alpha + \beta -1) \left[ (2d + \alpha + \beta)(2d + \alpha + \beta -2) x + (\alpha^2 - \beta^2) \right]}{2d(d + \alpha + \beta)(2d + \alpha + \beta - 2) } P\_{d-1}^{(\alpha, \beta)}(x)\\\\
            & - \frac{2(d+\alpha-1)(d+\beta-1)(2d+\alpha+\beta)}{2d(d + \alpha + \beta)(2d + \alpha + \beta - 2) }P\_{d-2}^{(\alpha, \beta)}(x).
        \end{align}
    $$

    By default, the input and output can also be processed with the optional pre- or post-processing functions
    in the jacobi expansion function.

    Attributes
    ----------
    name: str, default = 'jacobi_polynomial_expansion'
        Name of the expansion function.
    d: int, default = 2
        Degree of jacobi expansion.
    alpha: float, default = 1.0
        Parameter of jacobi polynomial representation.
    beta: float, default = 1.0
        Parameter of jacobi polynomial representation.
    Methods
    ----------
    __init__
        It performs the initialization of the expansion function.

    calculate_D
        It calculates the expansion space dimension D based on the input dimension parameter m.

    forward
        It implements the abstract forward method declared in the base expansion class.

    """
    def __init__(self, name: str = 'jacobi_polynomial_expansion', d: int = 2, alpha: float = 1.0, beta: float = 1.0, *args, **kwargs):
        r"""
        The initialization method of jacobi expansion function.

        It initializes a jacobi expansion object based on the input function name.
        This method will also call the initialization method of the base class as well.

        Parameters
        ----------
        name: str, default = 'jacobi_polynomial_expansion'
            The name of the jacobi expansion function.
        d: int, default = 5
            The degree of the jacobi expansion function.
        alpha: float, default = 1.0
            Parameter of jacobi polynomial representation.
        beta: float, default = 1.0
            Parameter of jacobi polynomial representation.

        Returns
        ----------
        transformation
            The jacobi expansion function.
        """
        super().__init__(name=name, *args, **kwargs)
        self.d = d
        self.alpha = alpha
        self.beta = beta

    def calculate_D(self, m: int):
        r"""
        The expansion dimension calculation method.

        It calculates the intermediate expansion space dimension based on the input dimension parameter m.
        For the jacobi expansion function, the expansion space dimension is determined by m and d,
        which can be represented as:

        $$ D = m d. $$

        Parameters
        ----------
        m: int
            The dimension of the input space.

        Returns
        -------
        int
            The dimension of the expansion space.
        """
        return m * self.d

    def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
        r"""
        The forward method of the data expansion function.

        It performs the jacobi data expansion of the input data and returns the expansion result
        according to the following equation:
        $$
            \begin{equation}
                \kappa(\mathbf{x} | d) = \left[ P_1^{(\alpha, \beta)}(\mathbf{x}), P_2^{(\alpha, \beta)}(\mathbf{x}),  \cdots, P_d^{(\alpha, \beta)}(\mathbf{x})\right] \in {R}^D.
            \end{equation}
        $$


        Parameters
        ----------
        x: torch.Tensor
            The input data vector.
        device: str, default = 'cpu'
            The device to perform the data expansion.

        Returns
        ----------
        torch.Tensor
            The expanded data vector of the input.
        """
        b, m = x.shape
        x = self.pre_process(x=x, device=device)

        expansion = torch.ones(size=[x.size(0), x.size(1), self.d+1]).to(device)
        if self.d > 0:
            expansion[:,:,1] = ((self.alpha-self.beta) + (self.alpha+self.beta+2) * x) / 2
        for n in range(2, self.d+1):
            coeff_1 = 2*n*(n+self.alpha+self.beta)*(2*n+self.alpha+self.beta-2)
            coeff_2 = (2*n+self.alpha+self.beta-1)*(2*n+self.alpha+self.beta)*(2*n+self.alpha+self.beta-2)
            coeff_3 = (2*n+self.alpha+self.beta-1)*(self.alpha**2-self.beta**2)
            coeff_4 = 2*(n+self.alpha-1)*(n+self.beta-1)*(2*n+self.alpha+self.beta)
            expansion[:,:,n] = ((coeff_2/coeff_1)*x + coeff_3/coeff_1)*expansion[:,:,n-1].clone() - (coeff_4/coeff_1)*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='jacobi_polynomial_expansion', d=2, alpha=1.0, beta=1.0, *args, **kwargs)

The initialization method of jacobi expansion function.

It initializes a jacobi expansion object based on the input function name. This method will also call the initialization method of the base class as well.

Parameters:

Name	Type	Description	Default
name	str	The name of the jacobi expansion function.	'jacobi_polynomial_expansion'
d	int	The degree of the jacobi expansion function.	2
alpha	float	Parameter of jacobi polynomial representation.	1.0
beta	float	Parameter of jacobi polynomial representation.	1.0

Returns:

Type	Description
transformation	The jacobi expansion function.

Source code in tinybig/expansion/recursive_expansion.py

def __init__(self, name: str = 'jacobi_polynomial_expansion', d: int = 2, alpha: float = 1.0, beta: float = 1.0, *args, **kwargs):
    r"""
    The initialization method of jacobi expansion function.

    It initializes a jacobi expansion object based on the input function name.
    This method will also call the initialization method of the base class as well.

    Parameters
    ----------
    name: str, default = 'jacobi_polynomial_expansion'
        The name of the jacobi expansion function.
    d: int, default = 5
        The degree of the jacobi expansion function.
    alpha: float, default = 1.0
        Parameter of jacobi polynomial representation.
    beta: float, default = 1.0
        Parameter of jacobi polynomial representation.

    Returns
    ----------
    transformation
        The jacobi expansion function.
    """
    super().__init__(name=name, *args, **kwargs)
    self.d = d
    self.alpha = alpha
    self.beta = beta

calculate_D(m)

The expansion dimension calculation method.

It calculates the intermediate expansion space dimension based on the input dimension parameter m. For the jacobi expansion function, the expansion space dimension is determined by m and d, which can be represented as:

\[ D = m d. \]

Parameters:

Name	Type	Description	Default
m	int	The dimension of the input space.	required

Returns:

Type	Description
int	The dimension of the expansion space.

Source code in tinybig/expansion/recursive_expansion.py

def calculate_D(self, m: int):
    r"""
    The expansion dimension calculation method.

    It calculates the intermediate expansion space dimension based on the input dimension parameter m.
    For the jacobi expansion function, the expansion space dimension is determined by m and d,
    which can be represented as:

    $$ D = m d. $$

    Parameters
    ----------
    m: int
        The dimension of the input space.

    Returns
    -------
    int
        The dimension of the expansion space.
    """
    return m * self.d

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

The forward method of the data expansion function.

It performs the jacobi data expansion of the input data and returns the expansion result according to the following equation: $$ \begin{equation} \kappa(\mathbf{x} | d) = \left[ P_1^{(\alpha, \beta)}(\mathbf{x}), P_2^{(\alpha, \beta)}(\mathbf{x}), \cdots, P_d^{(\alpha, \beta)}(\mathbf{x})\right] \in {R}^D. \end{equation} $$

Parameters:

Name	Type	Description	Default
x	Tensor	The input data vector.	required
device		The device to perform the data expansion.	'cpu'

Returns:

Type	Description
Tensor	The expanded data vector of the input.

Source code in tinybig/expansion/recursive_expansion.py

def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
    r"""
    The forward method of the data expansion function.

    It performs the jacobi data expansion of the input data and returns the expansion result
    according to the following equation:
    $$
        \begin{equation}
            \kappa(\mathbf{x} | d) = \left[ P_1^{(\alpha, \beta)}(\mathbf{x}), P_2^{(\alpha, \beta)}(\mathbf{x}),  \cdots, P_d^{(\alpha, \beta)}(\mathbf{x})\right] \in {R}^D.
        \end{equation}
    $$


    Parameters
    ----------
    x: torch.Tensor
        The input data vector.
    device: str, default = 'cpu'
        The device to perform the data expansion.

    Returns
    ----------
    torch.Tensor
        The expanded data vector of the input.
    """
    b, m = x.shape
    x = self.pre_process(x=x, device=device)

    expansion = torch.ones(size=[x.size(0), x.size(1), self.d+1]).to(device)
    if self.d > 0:
        expansion[:,:,1] = ((self.alpha-self.beta) + (self.alpha+self.beta+2) * x) / 2
    for n in range(2, self.d+1):
        coeff_1 = 2*n*(n+self.alpha+self.beta)*(2*n+self.alpha+self.beta-2)
        coeff_2 = (2*n+self.alpha+self.beta-1)*(2*n+self.alpha+self.beta)*(2*n+self.alpha+self.beta-2)
        coeff_3 = (2*n+self.alpha+self.beta-1)*(self.alpha**2-self.beta**2)
        coeff_4 = 2*(n+self.alpha-1)*(n+self.beta-1)*(2*n+self.alpha+self.beta)
        expansion[:,:,n] = ((coeff_2/coeff_1)*x + coeff_3/coeff_1)*expansion[:,:,n-1].clone() - (coeff_4/coeff_1)*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)