"""
Class that stores and perfom operation on the coefficients of a polynomial series.
"""
from __future__ import annotations
import typing
import logging
import numpy as np
import numpy.typing as npt
from scipy.special import eval_chebyt, eval_chebyu
from .grid import Grid
[docs]
class Polynomial:
"""
Class that stores and perfoms operations on the coefficients of a polynomial series.
"""
# Static tuples of allowed basis and direction strings
ALLOWED_BASES: typing.Final[tuple[str, ...]] = ("Cardinal", "Chebyshev", "Array")
ALLOWED_DIRECTIONS: typing.Final[tuple[str, ...]] = ("z", "pz", "pp", "Array")
[docs]
def __init__(
self,
coefficients: np.ndarray,
grid: Grid, basis: str | tuple[str,...]="Cardinal",
direction: str | tuple[str,...]="z",
endpoints: bool | tuple[bool,...]=False,
):
"""
Initialization of Polynomial object.
Parameters
----------
coefficients : array-like
Array of rank N containing the coefficients of a polynomial defined
by the object grid.
grid : Grid
An object of the Grid class defining the polynomial.
basis : string or tuple of strings, optional
Tuple of length N specifying in what basis each dimension of
coefficients is defined. Each component can either be 'Cardinal',
'Chebyshev' or 'Array'. The latter is to be used if the axis does
not corresponds to a polynomial series. Can also be a single
string, in which case all the dimensions are assumed to be in that
basis. The default is 'Cardinal'.
direction : string or tuple of strings, optional
Tuple of length N specifying what direction each dimension of
coefficients represents. Each component can either be 'z', 'pz',
'pp' or 'Array'. The latter can only be used if basis is 'Array'
for that axis. Can also be a single string, in which case all the
dimensions are assumed to be in that direction. The default is 'z'.
endpoints : bool or tuple of bool, optional
Tuple of length N specifying wheither each dimension includes the
endpoints. Can also be a single bool, in which case all the
dimensions are assumed to be the same. If False, the polynomial is
assumed to be 0 at the endpoints. The default is False.
Returns
-------
None.
"""
self.coefficients = np.asanyarray(coefficients)
self.rank = len(self.coefficients.shape)
self.grid = grid
if isinstance(basis, str):
basis = self.rank * (basis,)
self._checkBasis(basis)
self.basis = basis
if isinstance(direction, str):
direction = self.rank * (direction,)
self._checkDirection(direction)
if isinstance(endpoints, bool):
endpoints = self.rank * (endpoints,)
self._checkEndpoints(endpoints)
self.direction = direction
self.endpoints = endpoints
self._checkCoefficients(self.coefficients)
def __getitem__(self, key: int | tuple[int,...]) -> Polynomial:
basisList: list[str] = []
endpointsList: list[bool] = []
directionList: list[str] = []
if not isinstance(key, tuple):
key = (key,)
n = 0 # pylint: disable=invalid-name
for i, k in enumerate(key):
if isinstance(k, int):
n += 1 # pylint: disable=invalid-name
elif isinstance(k, slice):
basisList.append(self.basis[i])
directionList.append(self.direction[i])
endpointsList.append(self.endpoints[i])
n += 1 # pylint: disable=invalid-name
elif k is None:
basisList.append("Array")
directionList.append("z")
endpointsList.append(False)
else:
raise ValueError("Polynomial error: invalid key.")
basis = tuple(basisList) + self.basis[n:]
direction = tuple(directionList) + self.direction[n:]
endpoints = tuple(endpointsList) + self.endpoints[n:]
coefficients = np.array(self.coefficients[key])
return Polynomial(coefficients, self.grid, basis, direction, endpoints)
def __mul__(self, poly: Polynomial | np.ndarray | float) -> Polynomial:
if isinstance(poly, Polynomial):
assert self._isBroadcastable(
self.coefficients, poly.coefficients
), "Polynomial error: the two Polynomial objects are not broadcastable."
basis, direction, endpoints = self._findContraction(poly)
return Polynomial(
self.coefficients * poly.coefficients,
self.grid,
basis,
direction,
endpoints,
)
newCoeff = np.array(poly * self.coefficients)
assert (
len(newCoeff.shape) == self.rank
), """Polynomial error: the rank of the resulting Polynomial object must be
the same as the original one."""
return Polynomial(
newCoeff, self.grid, self.basis, self.direction, self.endpoints
)
def __add__(self, poly: Polynomial | np.ndarray | float) -> Polynomial:
if isinstance(poly, Polynomial):
assert self._isBroadcastable(
self.coefficients, poly.coefficients
), "Polynomial error: the two Polynomial objects are not broadcastable."
basis, direction, endpoints = self._findContraction(poly)
return Polynomial(
self.coefficients + poly.coefficients,
self.grid,
basis,
direction,
endpoints,
)
newCoeff = poly + self.coefficients
newCoeff = np.asanyarray(newCoeff)
assert (
newCoeff.ndim == self.rank
), """Polynomial error: the rank of the resulting Polynomial object must be
the same as the original one."""
return Polynomial(
newCoeff, self.grid, self.basis, self.direction, self.endpoints
)
def __sub__(self, poly: Polynomial | np.ndarray | float) -> Polynomial:
return self.__add__((-1.0) * poly)
def __rmul__(self, poly: Polynomial | np.ndarray | float) -> Polynomial:
return self.__mul__(poly)
def __radd__(self, poly: Polynomial | np.ndarray | float) -> Polynomial:
return self.__add__(poly)
def __rsub__(self, poly: Polynomial | np.ndarray | float) -> Polynomial:
return (-1) * self.__sub__(poly)
def _findContraction(
self,
poly: Polynomial,
) -> tuple[tuple[str,...], tuple[str,...], tuple[bool,...]]:
"""
Find the tuples basis, direction and endpoints resulting from the
contraction of self and poly
Parameters
----------
poly : Polynomial
Polynomial object.
Returns
-------
basis : tuple
basis tuple of the contracted polynomial.
direction : tuple
direction tuple of the contracted polynomial.
endpoints : tuple
endpoints tuple of the contracted polynomial.
"""
assert (
self.rank == poly.rank
), """Polynomial error: you can only combine two Polynomial objects with the
same rank."""
basis, endpoints, direction = [], [], []
for i in range(self.rank):
assert (
self.coefficients.shape[i] == 1
or poly.coefficients.shape[i] == 1
or (
self.basis[i] == poly.basis[i]
and self.direction[i] == poly.direction[i]
and self.endpoints[i] == poly.endpoints[i]
)
), "Polynomial error: the two Polynomial objects are not broadcastable."
if self.coefficients.shape[i] > 1:
basis.append(self.basis[i])
direction.append(self.direction[i])
endpoints.append(self.endpoints[i])
else:
basis.append(poly.basis[i])
direction.append(poly.direction[i])
endpoints.append(poly.endpoints[i])
return tuple(basis), tuple(direction), tuple(endpoints)
[docs]
def changeBasis(
self,
newBasis: str | tuple[str,...],
inverseTranspose: bool=False,
) -> None:
"""
Change the basis of the polynomial. Will change self.coefficients
accordingly.
Parameters
----------
newBasis : string or tuple of strings, optional
Tuple of length N specifying in what basis each dimension of
self.coefficients is defined. Each component can either be
'Cardinal' or 'Chebyshev'. Can also be a single string, in which
case all the dimensions are assumed to be in that basis.
inverseTranspose : bool, optional
If True, take the inverse-transpose of the transformation matrix.
This is useful, for example, when changing the basis of the
collision array.
Returns
-------
None.
"""
if isinstance(newBasis, str):
newBasis = self.rank * (newBasis,)
self._checkBasis(newBasis)
for i in range(self.rank):
if (
newBasis[i] != self.basis[i]
and newBasis[i] != "Array"
and self.basis[i] != "Array"
):
# Choosing the appropriate x, n and restriction
x = self.grid.getCompactCoordinates( # pylint: disable=invalid-name
self.endpoints[i], self.direction[i]
)
n, restriction = None, None # pylint: disable=invalid-name
if self.endpoints[i]:
if self.direction[i] == "z":
n = np.arange(self.grid.M + 1) # pylint: disable=invalid-name
elif self.direction[i] == "pz":
n = np.arange(self.grid.N + 1) # pylint: disable=invalid-name
else:
n = np.arange(self.grid.N) # pylint: disable=invalid-name
else:
if self.direction[i] == "z":
n = np.arange(2, self.grid.M + 1) # pylint: disable=invalid-name
restriction = "full"
elif self.direction[i] == "pz":
n = np.arange(2, self.grid.N + 1) # pylint: disable=invalid-name
restriction = "full"
else:
n = np.arange(1, self.grid.N) # pylint: disable=invalid-name
restriction = "partial"
# Computing the Tn matrix
tnMatrix = np.array(self.chebyshev(x[:, None], n[None, :], restriction))
if newBasis[i] == "Chebyshev":
tnMatrix = np.linalg.inv(tnMatrix)
if inverseTranspose:
tnMatrix = np.transpose(np.linalg.inv(tnMatrix))
tnMatrix = np.expand_dims(
tnMatrix,
tuple(np.arange(i)) + tuple(np.arange(i + 2, self.rank + 1))
)
# Contracting M with self.coefficient
self.coefficients = np.sum(
tnMatrix * np.expand_dims(self.coefficients, i), axis=i + 1
)
self.basis = newBasis
[docs]
def evaluate(
self,
compactCoord: np.ndarray,
axes: tuple[int,...] | None=None,
) -> np.ndarray | float:
"""
Evaluates the polynomial at the compact coordinates x.
Parameters
----------
compactCoord : array-like
Compact coordinates at which to evaluate the polynomial. Must have
a shape (len(axes),:) or (len(axes),).
axes : tuple or None, optional
Axes along which to be evaluated. If None, evaluate the polynomial
along all the axes. Default is None.
Returns
-------
array-like
Values of the polynomial at x.
"""
compactCoord = np.asarray(compactCoord)
if axes is None:
axes = tuple(np.arange(self.rank))
assert (
compactCoord.shape[0] == len(axes) and 1 <= len(compactCoord.shape) <= 2
), f"""Polynomial error: compactCoord has shape {compactCoord.shape} but must
be ({self.rank},:) or ({self.rank},)."""
singlePoint = False
if len(compactCoord.shape) == 1:
compactCoord = compactCoord.reshape((len(axes), 1))
singlePoint = True
polynomials = np.ones((compactCoord.shape[1],) + self.coefficients.shape)
for j, i in enumerate(axes):
assert (
self.basis[i] != "Array"
), "Polynomial error: cannot evaluate along an 'Array' axis."
# Choosing the appropriate n
n: np.ndarray # pylint: disable=invalid-name
if self.endpoints[i]:
if self.direction[i] == "z":
n = np.arange(self.grid.M + 1) # pylint: disable=invalid-name
elif self.direction[i] == "pz":
n = np.arange(self.grid.N + 1) # pylint: disable=invalid-name
else:
n = np.arange(self.grid.N) # pylint: disable=invalid-name
else:
if self.direction[i] == "z":
n = np.arange(1, self.grid.M) # pylint: disable=invalid-name
elif self.direction[i] == "pz":
n = np.arange(1, self.grid.N) # pylint: disable=invalid-name
else:
n = np.arange(self.grid.N - 1) # pylint: disable=invalid-name
# Computing the polynomial basis in the i direction
pn: np.ndarray # pylint: disable=invalid-name
if self.basis[i] == "Cardinal":
pn = np.array( # pylint: disable=invalid-name
self.cardinal(compactCoord[j, :, None], n[None, :],
self.direction[i]))
elif self.basis[i] == "Chebyshev":
restriction = None
if not self.endpoints[i]:
n += 1 # pylint: disable=invalid-name
if self.direction[i] == "z":
restriction = "full"
elif self.direction[i] == "pz":
restriction = "full"
else:
restriction = "partial"
pn = np.array( # pylint: disable=invalid-name
self.chebyshev(compactCoord[j, :, None], n[None, :], restriction))
polynomials *= np.expand_dims(
pn, tuple(np.arange(1, i + 1)) + tuple(np.arange(i + 2, self.rank + 1))
)
result = np.sum(
self.coefficients[None, ...] * polynomials, axis=tuple(np.array(axes) + 1)
)
if singlePoint:
return float(result[0])
return np.array(result)
[docs]
def cardinal(
self,
compactCoord: npt.ArrayLike,
n: npt.ArrayLike, # pylint: disable=invalid-name
direction: str,
) -> np.ndarray:
r"""
Computes the cardinal polynomials :math:`C_n(x)` defined by grid.
Parameters
----------
compactCoord : array_like
Compact coordinate at which to evaluate the Chebyshev polynomial. Must be
broadcastable with n.
n : array_like
Order of the cardinal polynomial to evaluate. Must be
broadcastable with x.
direction : string
Select the direction in which to compute the matrix.
Can either be 'z', 'pz' or 'pp'.
Returns
-------
cn : array_like
Values of the cardinal functions.
"""
compactCoord = np.asarray(compactCoord)
n = np.asarray(n)
assert self._isBroadcastable(
compactCoord, n
), "Polynomial error: x and n are not broadcastable."
assert direction in self.ALLOWED_DIRECTIONS, (
f"Polynomial error: unkown direction {direction}"
)
# Selecting the appropriate grid and resizing it
grid = self.grid.getCompactCoordinates(True, direction)
completeGrid = np.expand_dims(
grid, tuple(np.arange(1, len((n * compactCoord).shape) + 1)))
nGrid = grid[n]
# Computing all the factor in the product defining the cardinal functions
cardinalPartial = np.divide(
compactCoord - completeGrid,
nGrid - completeGrid,
where=nGrid - completeGrid != 0
)
# Multiplying all the factors to get the cardinal functions
cardinal = np.prod(np.where(nGrid - completeGrid == 0, 1, cardinalPartial),
axis=0)
return np.array(cardinal)
[docs]
def chebyshev(
self,
compactCoord: npt.ArrayLike,
n: npt.ArrayLike, # pylint: disable=invalid-name
restriction: str | None=None
) -> np.ndarray:
r"""
Computes the Chebyshev polynomial :math:`T_n(x)`.
Parameters
----------
compactCoord : array_like
Compact coordinate at which to evaluate the Chebyshev polynomial. Must be
broadcastable with n.
n : array_like
Order of the Chebyshev polynomial to evaluate. Must be
broadcastable with x.
restriction : None or string, optional
Select the restriction on the Chebyshev basis. If None, evaluates
the unrestricted basis. If 'full', the polynomials are 0 at
:math:`x=\pm 1`. If 'partial', the polynomials are 0 at :math:`x=+1`.
Returns
-------
chen : array_like
Values of the polynomial
"""
compactCoord = np.asarray(compactCoord)
n = np.asarray(n)
assert self._isBroadcastable(
compactCoord, n
), "Polynomial error: compactCoord and n are not broadcastable."
# Computing the unrestricted basis
cheb = eval_chebyt(n, compactCoord)
# Applying the restriction
if restriction == "partial":
cheb -= 1
elif restriction == "full":
cheb -= np.where(n % 2 == 0, 1, compactCoord)
return np.array(cheb)
[docs]
def integrate(
self,
axis: int | tuple[int,...] | None=None,
weight: npt.ArrayLike=1,
) -> Polynomial | float:
r"""
Computes the integral of the polynomial :math:`\int_{-1}^1 dx P(x)w(x)`
along some axis using Gauss-Chebyshev-Lobatto quadrature.
Parameters
----------
axis : None, int or tuple
axis along which the integral is taken. Can either be None, a int or a
tuple of int. If None, integrate along all the axes.
weight : array-like, optional
Integration weight. Must be an object broadcastable with
self.coefficients. Default is 1.
Returns
-------
Polynomial or float
If axis=None, returns a float. Otherwise, returns an object of the
class Polynomial containing the coefficients of the
integrated polynomial along the remaining axes.
"""
if weight is None:
weight = 1
if axis is None:
axis = tuple(np.arange(self.rank))
if isinstance(axis, int):
axis = (axis,)
self._checkAxis(axis)
# Express the integrated axes in the cardinal basis
basis = []
for i in range(self.rank):
if i in axis:
assert (
self.basis[i] != "Array"
), "Polynomial error: cannot integrate along an 'Array' axis."
basis.append("Cardinal")
else:
basis.append(self.basis[i])
self.changeBasis(tuple(basis))
integrand = weight * self.coefficients
newBasis, newDirection, newEndpoints = [], [], []
for i in range(self.rank):
if i in axis:
compactCoord = self.grid.getCompactCoordinates(
self.endpoints[i], self.direction[i]
)
weights = np.pi * np.ones(compactCoord.size)
if self.direction[i] == "z":
weights /= self.grid.M
elif self.direction[i] == "pz":
weights /= self.grid.N
elif self.direction[i] == "pp":
weights /= self.grid.N - 1
if not self.endpoints[i]:
weights[0] /= 2
if self.endpoints[i]:
weights[0] /= 2
weights[-1] /= 2
integrand *= np.expand_dims(
np.sqrt(1 - compactCoord**2) * weights,
tuple(np.arange(i)) + tuple(np.arange(i + 1, self.rank)),
)
else:
newBasis.append(self.basis[i])
newDirection.append(self.direction[i])
newEndpoints.append(self.endpoints[i])
result = np.sum(integrand, axis)
## This check fails too easily with extended types
"""
if isinstance(result, float):
return result
"""
if np.asanyarray(result).ndim == 0:
return float(result)
return Polynomial(
result,
self.grid,
tuple(newBasis),
tuple(newDirection),
tuple(newEndpoints),
)
[docs]
def derivative(self, axis: int | tuple[int,...]) -> Polynomial:
"""
Computes the derivative of the polynomial and returns it in a
Polynomial object.
Parameters
----------
axis : int or tuple
axis along which the derivative is taken. Can either be a int or a
tuple of int.
Returns
-------
Polynomial
Object of the class Polynomial containing the coefficients of the
derivative polynomial (in the compact coordinates). The axis along
which the derivative is taken is always returned with the endpoints
in the cardinal basis.
"""
if isinstance(axis, int):
axis = (axis,)
self._checkAxis(axis)
coeffDeriv = np.array(self.coefficients)
basis, endpoints = [], []
for i in range(self.rank):
if i in axis:
assert (
self.basis[i] != "Array"
), "Polynomial error: cannot differentiate along an 'Array' axis."
derivMatrix = self.derivMatrix(
self.basis[i], self.direction[i], self.endpoints[i]
)
derivMatrix = np.expand_dims(
derivMatrix,
tuple(np.arange(i)) + tuple(np.arange(i + 2, self.rank + 1))
)
coeffDeriv = np.sum(derivMatrix*np.expand_dims(coeffDeriv, i), axis=i+1)
basis.append("Cardinal")
endpoints.append(True)
else:
basis.append(self.basis[i])
endpoints.append(self.endpoints[i])
return Polynomial(
coeffDeriv, self.grid, tuple(basis), self.direction, tuple(endpoints)
)
[docs]
def matrix(self, basis: str, direction: str, endpoints: bool=False) -> np.ndarray:
r"""
Returns the matrix :math:`M_{ij}=T_j(x_i)` or :math:`M_{ij}=C_j(x_i)` computed
in a specific direction.
Parameters
----------
basis : string
Select the basis of polynomials. Can be 'Cardinal' or 'Chebyshev'
direction : string
Select the direction in which to compute the matrix. Can either be 'z', 'pz'
or 'pp'
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
"""
if basis == "Cardinal":
return self._cardinalMatrix(direction, endpoints)
if basis == "Chebyshev":
return self._chebyshevMatrix(direction, endpoints)
raise ValueError("basis must be either 'Cardinal' or 'Chebyshev'.")
[docs]
def derivMatrix(
self,
basis: str,
direction: str,
endpoints: bool=False,
) -> np.ndarray:
"""
Computes the derivative matrix of either the Chebyshev or cardinal polynomials
in some direction.
Parameters
----------
basis : string
Select the basis of polynomials. Can be 'Cardinal' or 'Chebyshev'
direction : string
Select the direction in which to compute the matrix. Can be 'z', 'pz' or
'pp'.
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
Returns
-------
deriv : array_like
Derivative matrix.
"""
assert basis in ['Cardinal', 'Chebyshev'], """basis must be either
'Cardinal' or 'Chebyshev'."""
if basis == "Cardinal":
return self._cardinalDeriv(direction, endpoints)
return self._chebyshevDeriv(direction, endpoints)
def _cardinalMatrix(self, direction: str, endpoints: bool=False) -> np.ndarray:
r"""
Returns the matrix :math:`M_{ij}=C_j(x_i)` computed in a specific direction.
Parameters
----------
direction : string
Select the direction in which to compute the matrix. Can either be 'z', 'pz'
or 'pp'.
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
"""
assert direction in ['z','pz','pp'], """direction must be either 'z',
'pz' or 'pp'."""
if direction == "z":
return np.identity(self.grid.M - 1 + 2 * endpoints)
if direction == "pz":
return np.identity(self.grid.N - 1 + 2 * endpoints)
return np.identity(self.grid.N - 1 + endpoints)
def _chebyshevMatrix(self, direction: str, endpoints: bool=False) -> np.ndarray:
r"""
Returns the matrix :math:`M_{ij}=T_j(x_i)` computed in a specific direction.
Parameters
----------
direction : string
Select the direction in which to compute the matrix. Can either be 'z', 'pz'
or 'pp'
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
"""
grid: np.ndarray
n: np.ndarray # pylint: disable=invalid-name
restriction = None
if direction == "z":
grid = self.grid.getCompactCoordinates(endpoints)[0]
n = np.arange(grid.size) + 2 - 2 * endpoints # pylint: disable=invalid-name
restriction = "full"
elif direction == "pz":
grid = self.grid.getCompactCoordinates(endpoints)[1]
n = np.arange(grid.size) + 2 - 2 * endpoints # pylint: disable=invalid-name
restriction = "full"
elif direction == "pp":
grid = self.grid.getCompactCoordinates(endpoints)[2]
n = np.arange(grid.size) + 1 - endpoints # pylint: disable=invalid-name
restriction = "partial"
if endpoints:
restriction = None
return self.chebyshev(grid[:, None], n[None, :], restriction)
def _cardinalDeriv(self, direction: str, endpoints: bool=False) -> np.ndarray:
"""
Computes the derivative matrix of the cardinal functions in some direction.
Parameters
----------
direction : string
Select the direction in which to compute the matrix. Can either be 'z', 'pz'
or 'pp'.
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
Returns
-------
deriv : array_like
Derivative matrix.
"""
grid = self.grid.getCompactCoordinates(True, direction)
# Computing the diagonal part
diagonal = np.sum(
np.where(
grid[:, None] - grid[None, :] == 0,
0,
np.divide(
1,
grid[:, None] - grid[None, :],
where=grid[:, None] - grid[None, :] != 0,
),
),
axis=1,
)
# Computing the off-diagonal part
offDiagonal = np.prod(
np.where(
(grid[:, None, None] - grid[None, None, :])
* (grid[None, :, None] - grid[None, None, :])
== 0,
1,
np.divide(
grid[None, :, None] - grid[None, None, :],
grid[:, None, None] - grid[None, None, :],
where=grid[:, None, None] - grid[None, None, :] != 0,
),
),
axis=-1,
)
# Putting all together
derivWithEndpoints = np.where(
grid[:, None] - grid[None, :] == 0,
diagonal[:, None],
np.divide(
offDiagonal,
grid[:, None] - grid[None, :],
where=grid[:, None] - grid[None, :] != 0,
),
)
deriv: np.ndarray
if not endpoints:
if direction in ["z", "pz"]:
deriv = derivWithEndpoints[1:-1, :]
elif direction == "pp":
deriv = derivWithEndpoints[:-1, :]
else:
deriv = derivWithEndpoints
return np.transpose(deriv)
def _chebyshevDeriv(self, direction: str, endpoints: bool=False) -> np.ndarray:
"""
Computes the derivative matrix of the Chebyshev polynomials in some direction.
Parameters
----------
direction : string
Select the direction in which to compute the matrix. Can either be 'z', 'pz'
or 'pp'.
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
Returns
-------
deriv : array_like
Derivative matrix.
"""
grid = self.grid.getCompactCoordinates(True, direction)
n: np.ndarray # pylint: disable=invalid-name
restriction = None
if direction == "z":
n = np.arange(2 - 2 * endpoints, grid.size) # pylint: disable=invalid-name
restriction = "full"
elif direction == "pz":
n = np.arange(2 - 2 * endpoints, grid.size) # pylint: disable=invalid-name
restriction = "full"
elif direction == "pp":
n = np.arange(1 - endpoints, grid.size) # pylint: disable=invalid-name
restriction = "partial"
deriv = n[None, :] * eval_chebyu(n[None, :] - 1, grid[:, None])
if restriction == "full" and not endpoints:
deriv -= np.where(n[None, :] % 2 == 0, 0, 1)
return np.array(deriv)
def _checkBasis(self, basis: tuple[str,...]) -> None:
assert isinstance(
basis, tuple
), "Polynomial error: basis must be a tuple or a string."
assert (
len(basis) == self.rank
), 'Polynomial error: basis must be a tuple of length "rank".'
for x in basis: # pylint: disable=invalid-name
assert x in self.ALLOWED_BASES, f"Polynomial error: unkown basis {x}"
def _checkDirection(self, direction: tuple[str,...]) -> None:
assert isinstance(
direction, tuple
), "Polynomial error: direction must be a tuple or a string."
assert (
len(direction) == self.rank
), 'Polynomial error: direction must be a tuple of length "rank".'
for i, x in enumerate(direction): # pylint: disable=invalid-name
assert x in self.ALLOWED_DIRECTIONS, (
f"Polynomial error: unkown direction {x}"
)
if x == "Array":
assert (
self.basis[i] == "Array"
), """Polynomial error: if the direction is 'Array', the basis must be
'Array' too."""
def _checkEndpoints(self, endpoints: tuple[bool,...]) -> None:
assert isinstance(
endpoints, tuple
), "Polynomial error: endpoints must be a tuple or a bool."
assert (
len(endpoints) == self.rank
), 'Polynomial error: endpoints must be a tuple of length "rank".'
for x in endpoints: # pylint: disable=invalid-name
assert isinstance(
x, bool
), "Polynomial error: endpoints can only contain bool."
def _checkCoefficients(self, coefficients: np.ndarray) -> None:
for i, size in enumerate(coefficients.shape):
if self.basis[i] != "Array":
if self.direction[i] == "z":
assert (
size + 2 * (1 - self.endpoints[i]) == self.grid.M + 1
), f"""Polynomial error: coefficients with invalid size in
dimension {i}."""
elif self.direction[i] == "pz":
assert (
size + 2 * (1 - self.endpoints[i]) == self.grid.N + 1
), f"""Polynomial error: coefficients with invalid size in
dimension {i}."""
else:
assert (
size + (1 - self.endpoints[i]) == self.grid.N
), f"""Polynomial error: coefficients with invalid size in
dimension {i}."""
def _checkAxis(self, axis: tuple[int,...]) -> None:
assert isinstance(
axis, tuple
), "Polynomial error: axis must be a tuple or a int."
for x in axis: # pylint: disable=invalid-name
assert isinstance(x, int), "Polynomial error: axis must be a tuple of int."
assert 0 <= x < self.rank, "Polynomial error: axis out of range."
def _isBroadcastable(self, array1: np.ndarray, array2: np.ndarray) -> bool:
"""
Verifies that array1 and array2 are broadcastable, which mean that they
can be multiplied together.
Parameters
----------
array1 : array_like
First array.
array2 : array_like
Second array.
Returns
-------
bool
True if the two arrays are broadcastable, otherwise False.
"""
for a, b in zip( # pylint: disable=invalid-name
np.asanyarray(array1).shape[::-1], np.asanyarray(array2).shape[::-1]
):
if a == 1 or b == 1 or a == b:
pass
else:
return False
return True
[docs]
class SpectralConvergenceInfo:
"""
Carries information about the convergence of a polynomial expansion.
Assumes input is a 1d array of coefficients in the Chebyshev basis. Fits a slope to the logarithm of the absolute value of these coefficients. Also, finds the average value of the index, treating the coefficients as a
histogram.
"""
coefficients: np.ndarray
"""Coefficients of expansion in the Chebyshev basis, must be 1d."""
weightPower: int = 0
r"""Additional powers of :math:`n` to weight by in assessing convergence.
Default is zero."""
offset: int = 0
r"""Offest in :math:`n`. Default is zero."""
apparentConvergence: bool = False
"""True if spectral expansion appears to be converging, False otherwise."""
spectralPeak: int = 0
"""Positions of the peak of the spectral expansion."""
spectralExponent: float = 0.0
r"""Exponent :math:`\sigma` of :math:`A e^{\sigma n}` fit to spectral expansion."""
[docs]
def __init__(
self, coefficients: np.ndarray, weightPower: int=0, offset: int=0
) -> None:
"""Initialise given an array."""
assert len(coefficients.shape) == 1, "SpectralConvergenceInfo requires a 1d array"
self.coefficients = coefficients[:]
self.weightPower = weightPower
self.offset = offset
self._checkSpectralConvergence()
def _checkSpectralConvergence(self) -> None:
"""
Check for spectral convergence, performing fits and looking at the
position of the maximum.
"""
nCoeff = len(self.coefficients)
weight = (1 + self.offset + np.arange(nCoeff)) ** self.weightPower
absCoefficients = abs(self.coefficients)
if nCoeff < 2:
logging.warning("Spectral convergence tests not valid for n < 2.")
return
if nCoeff < 3:
strict = False
else:
strict = True
# fit slope to the log of the coefficients
if strict:
# enough points to fit slopes with errors
fit = np.polyfit(
np.arange(nCoeff) + self.offset,
np.log(absCoefficients),
# np.log(absCoefficients * weight),
1,
cov=True,
)
self.spectralExponent = fit[0][0]
# self.apparentConvergence = fit[0][0] < - np.sqrt(fit[1][0, 0])
self.apparentConvergence = fit[0][0] < 0
else:
# not enough points to get sensible errors
fit = np.polyfit(
np.arange(nCoeff) + self.offset,
np.log(absCoefficients),
# np.log(absCoefficients * weight),
1,
cov=False,
)
self.spectralExponent = fit[0]
self.apparentConvergence = fit[0] < 0
# Index weighted by absolute value in array
self.spectralPeak = int(
np.sum((np.arange(nCoeff) + self.offset) * weight * absCoefficients) / np.sum(absCoefficients * weight)
)
# Alternative convergence condition
# self.apparentConvergence = self.spectralPeak < self.offset + nCoeff // 2