orthax
orthax is a Python package for working with orthogonal (and other) polynomials in JAX.
It largely seeks to replicate the functionality of the numpy.polynomial package,
through there are some API differences due to limitations of JAX, primarily that
trailing zeros are not automatically trimmed from series, so you should do that
manually if it becomes a concern.
For full details of various options see the Documentation
Installation
orthax is installable with pip:
pip install orthax
API Documentation
Within the documentation for this package, a “finite power series,”
i.e., a polynomial (also referred to simply as a “series”) is represented
by a 1-D JAX array of the polynomial’s coefficients, ordered from lowest
order term to highest. For example, array([1,2,3]) represents
P_0 + 2*P_1 + 3*P_2, where P_n is the n-th order basis polynomial
applicable to the specific module in question, e.g., polynomial (which
“wraps” the “standard” basis) or chebyshev. For optimal performance,
all operations on polynomials, including evaluation at an argument, are
implemented as operations on the coefficients. Additional (module-specific)
information can be found in the docstring for the module of interest.
This package provides functions for operations on each of six different kinds of polynomials: