"""Assemble sparse univariate and multivariate basis polynomials."""
import math
from numbers import Real
from typing import Callable, Optional, Sequence, TypeAlias
import numpy as np
import scipy.integrate
import scipy.special
from numpy.polynomial import Polynomial
from scipy.stats import norm as std_normal
from . import sampler as pt_sampler
from .parameter import Parameter
Function1D: TypeAlias = Callable[[np.ndarray], np.ndarray]
FunctionND: TypeAlias = Callable[[np.ndarray], np.ndarray]
# Function1D: TypeAlias = Callable[[Float[np.ndarray, "*shape"]],
# Float[np.ndarray, " *shape"]]
# FunctionND: TypeAlias = Callable[[Float[np.ndarray, "*shape dimension"]],
# Float[np.ndarray, " *shape"]]
[docs]def univariate_basis(
params: Sequence[Parameter],
degs: Sequence[int] | np.ndarray,
) -> list[list[Polynomial]]:
"""Assemble a univariate polynomial basis.
Set polynomial basis up to deg for each parameter in `params` according to
the parameter distribution and area of definition.
Parameters
----------
params : list of `pythia.parameter.Parameter`
Parameters to compute univariate basis function for.
degs : array_like
Max. degrees of univariate polynomials for each parameter.
Returns
-------
:
List of normalized univariate polynomials w.r.t. parameter domain and
distribution up to specified degree for each parameter in `params`.
"""
assert len(params) == len(degs)
def get_basis(param: Parameter, deg: int) -> list[Polynomial]:
pdf = pt_sampler.assign_sampler(param).pdf
if param.distribution == "uniform":
return normalize_polynomial(pdf, set_legendre_basis(param, deg), param)
elif param.distribution == "normal":
return normalize_polynomial(pdf, set_hermite_basis(param, deg), param)
elif param.distribution == "gamma":
return normalize_polynomial(pdf, set_laguerre_basis(param, deg), param)
elif param.distribution == "beta":
return normalize_polynomial(pdf, set_jacobi_basis(param, deg), param)
else:
raise ValueError(
f'Unsupported distribution "{param.distribution}"' f" for {param.name}"
)
return [get_basis(param, deg) for param, deg in zip(params, degs)]
[docs]def multivariate_basis(
univariate_bases: Sequence[Sequence[Function1D]],
indices: np.ndarray,
partial: Optional[list[int]] = None,
) -> list[Callable]:
"""Assemble multivariate polynomial basis.
Set the (partial derivative of the) multivariate (product) polynomial basis
functions.
Parameters
----------
univariate_bases : list of list of callable
Univariate basis functions for parameters. Is called by
`univariate_bases[paramIdx][deg]()`.
indices : array_like
Array of multiindices for multivariate basis functions.
partial : list of int
Number of partial derivatives for each dimension. Length is same as
`univariate_bases`.
Returns
-------
:
List of multivariate product polynomials with univariate degrees as
specified in `indices`.
"""
assert len(univariate_bases) == indices.shape[1]
if partial is not None:
assert len(partial) == indices.shape[1]
basis_list = []
for index in indices:
def fun(x: np.ndarray, index: np.ndarray = index) -> np.ndarray:
if not 1 <= x.ndim <= 2:
raise ValueError(f"Wrong ndim '{x.ndim}'")
if x.ndim == 1:
x.shape = 1, -1
if partial is None:
basis = [univariate_bases[k][mu_k] for k, mu_k in enumerate(index)]
else:
basis = [
univariate_bases[k][mu_k].deriv(partial[k])
for k, mu_k in enumerate(index)
]
return np.prod([basis[k](x[:, k]) for k, _ in enumerate(index)], axis=0)
basis_list += [fun]
return basis_list
[docs]def normalize_polynomial(
weight: Function1D,
basis: Sequence[Polynomial],
param: Parameter,
) -> list[Polynomial]:
r"""Normalize orthogonal polynomials.
Normalize a polynomial of an orthogonal system with respect to the scalar
product
.. math::
a(u,v)_\mathrm{pdf} = \int u(p) v(p) \mathrm{pdf}(p) \mathrm{d}p.
The normalized polynomial :math:`\phi_j` for any given polynomial
:math:`P_j` is given by :math:`\phi_j = P_j / \sqrt{c_j}`
for the constant
:math:`c_j = \int \mathrm{pdf}(p) * P_j(p)^2 \mathrm{d}p`.
Parameters
----------
weight : callable
Probability density function.
basis : list of `numpy.polynomial.Polynomial`
Polynomials to normalize w.r.t. weight.
param : `pythia.parameter.Parameter`
Parameter used for distribution and domain information.
Returns
-------
:
List of normalized univariate polynomials.
"""
if param.distribution == "normal":
# TODO: Only works for certain weights!
assert np.allclose(
weight(np.linspace(-5, 5, 100)), std_normal.pdf(np.linspace(-5, 5, 100))
)
def norm(polynomial: Polynomial, degree: int) -> float:
return np.sqrt(math.factorial(degree))
else:
def norm(polynomial: Polynomial, degree: int) -> float:
def integrand(x: np.ndarray) -> np.ndarray:
return weight(x) * polynomial(x) ** 2
return np.sqrt(
scipy.integrate.quad(integrand, param.domain[0], param.domain[1])[0]
)
return [p / norm(p, j) for j, p in enumerate(basis)]
[docs]def set_legendre_basis(param: Parameter, deg: int) -> list[Polynomial]:
"""Generate list of the Legendre Polynomials.
Generate the Legendre Polynomials up to certain degree on the interval
specified by the parameter.
Parameters
----------
param : `pythia.parameters.Parameter`
Parameter for basis function. Needs to be uniformly distributed.
deg : int
Maximum degree for polynomials.
Returns
-------
:
List of Legendre polynomials up to (including) degree specified in
`deg`.
"""
return [
np.polynomial.legendre.Legendre([0] * j + [1], param.domain).convert(
kind=Polynomial
)
for j in range(deg + 1)
]
[docs]def set_hermite_basis(param: Parameter, deg: int) -> list[Polynomial]:
"""Generate list of probabilists Hermite polynomials.
Generate the Hermite Polynomials up to certain degree according to the
mean and variance of the specified parameter.
Parameters
----------
param : `pythia.parameters.Parameter`
Parameter for basis function. Needs to be normal distributed.
deg : int
Maximum degree for polynomials.
Returns
-------
:
List of probabilists Hermite polynomials up to (including) degree
specified in `deg`.
"""
assert isinstance(param.mean, Real)
assert isinstance(param.var, Real)
std = np.sqrt(param.var)
a = -param.mean / (std * np.sqrt(2))
b = 1 / (np.sqrt(2) * std)
shift_and_scale = Polynomial([a, b])
def scaled_hermite(degree: int) -> Polynomial:
p = np.polynomial.hermite.Hermite([0] * degree + [1])
return 2 ** (-degree / 2) * p(shift_and_scale)
return [scaled_hermite(j) for j in range(deg + 1)]
[docs]def set_jacobi_basis(param: Parameter, deg: int) -> list[Polynomial]:
"""Generate list of Jacobi polynomials.
Generate the Jacobi Polynomials up to certain degree on the interval
and DoFs specified by the parameter.
.. note::
The Jacobi polynomials have leading coefficient 1.
Parameters
----------
param : `pythia.parameters.Parameter`
Parameter for basis function. Needs to be Beta-distributed.
deg : int
Maximum degree for polynomials.
Returns
-------
:
List of Jacobi polynomials up to (including) degree specified in `deg`.
"""
domain = tuple(param.domain)
assert (
len(domain) == 2 and isinstance(domain[0], Real) and isinstance(domain[1], Real)
)
a = transform_interval(0.0, [-1, 1], domain)
b = transform_interval(1.0, [-1, 1], domain) - a
shift_and_scale = Polynomial([a, b])
def scaled_jacobi(degree: int) -> Polynomial:
if degree == 0:
return Polynomial([1])
assert isinstance(param.alpha, Real)
assert isinstance(param.beta, Real)
roots, _ = scipy.special.roots_jacobi(degree, param.beta - 1, param.alpha - 1)
coeff = np.polynomial.polynomial.polyfromroots(shift_and_scale(roots))
return Polynomial(coeff)
return [scaled_jacobi(j) for j in range(deg + 1)]
[docs]def set_laguerre_basis(param: Parameter, deg: int) -> list[Polynomial]:
"""Generate list of Leguerre polynomials.
Generate the generalized Laguerre polynomials up to certain degree on
the interval and DoFs specified by the parameter.
Parameters
----------
param : `pythia.parameters.Parameter`
Parameter for basis function. Needs to be Gamma-distributed.
deg : int
Maximum degree for polynomials.
Returns
-------
:
List of Laguerre polynomials up to (including) degree specified in
`deg`.
"""
assert isinstance(param.beta, Real)
shift_and_scale = Polynomial([param.domain[0], 1 / param.beta])
def scaled_laguerre(degree: int) -> Polynomial:
if degree == 0:
return Polynomial([1])
assert isinstance(param.alpha, Real)
roots, _ = scipy.special.roots_genlaguerre(degree, param.alpha - 1)
coeff = np.polynomial.polynomial.polyfromroots(shift_and_scale(roots))
return Polynomial(coeff)
return [scaled_laguerre(j) for j in range(deg + 1)]