172 lines
5.2 KiB
Python
172 lines
5.2 KiB
Python
"""Partial replacements for numpy polynomial routines, with Array API compatibility.
|
|
|
|
This module contains both "old-style", np.poly1d, routines from the main numpy
|
|
namespace, and "new-style", np.polynomial.polynomial, routines.
|
|
|
|
To distinguish the two sets, the "new-style" routine names start with `npp_`
|
|
"""
|
|
import scipy._lib.array_api_extra as xpx
|
|
from scipy._lib._array_api import xp_promote, xp_default_dtype
|
|
|
|
|
|
def _sort_cmplx(arr, xp):
|
|
# xp.sort is undefined for complex dtypes. Here we only need some
|
|
# consistent way to sort a complex array, including equal magnitude elements.
|
|
arr = xp.asarray(arr)
|
|
if xp.isdtype(arr.dtype, 'complex floating'):
|
|
sorter = abs(arr) + xp.real(arr) + xp.imag(arr)**3
|
|
else:
|
|
sorter = arr
|
|
|
|
idxs = xp.argsort(sorter)
|
|
return arr[idxs]
|
|
|
|
|
|
def polyroots(coef, *, xp):
|
|
"""numpy.roots, best-effor replacement
|
|
"""
|
|
if coef.shape[0] < 2:
|
|
return xp.asarray([], dtype=coef.dtype)
|
|
|
|
root_func = getattr(xp, 'roots', None)
|
|
if root_func:
|
|
# NB: cupy.roots is broken in CuPy 13.x, but CuPy is handled via delegation
|
|
# so we never hit this code path with xp being cupy
|
|
return root_func(coef)
|
|
|
|
# companion matrix
|
|
n = coef.shape[0]
|
|
a = xp.eye(n - 1, n - 1, k=-1, dtype=coef.dtype)
|
|
a[:, -1] = -xp.flip(coef[1:]) / coef[0]
|
|
|
|
# non-symmetric eigenvalue problem is not in the spec but is available on e.g. torch
|
|
if hasattr(xp.linalg, 'eigvals'):
|
|
return xp.linalg.eigvals(a)
|
|
else:
|
|
import numpy as np
|
|
return xp.asarray(np.linalg.eigvals(np.asarray(a)))
|
|
|
|
|
|
# https://github.com/numpy/numpy/blob/v2.1.0/numpy/lib/_function_base_impl.py#L1874-L1925
|
|
def _trim_zeros(filt, trim='fb'):
|
|
first = 0
|
|
trim = trim.upper()
|
|
if 'F' in trim:
|
|
for i in filt:
|
|
if i != 0.:
|
|
break
|
|
else:
|
|
first = first + 1
|
|
last = filt.shape[0]
|
|
if 'B' in trim:
|
|
for i in filt[::-1]:
|
|
if i != 0.:
|
|
break
|
|
else:
|
|
last = last - 1
|
|
return filt[first:last]
|
|
|
|
|
|
# ### Old-style routines ###
|
|
|
|
|
|
# https://github.com/numpy/numpy/blob/v2.2.0/numpy/lib/_polynomial_impl.py#L1232
|
|
def _poly1d(c_or_r, *, xp):
|
|
""" Constructor of np.poly1d object from an array of coefficients (r=False)
|
|
"""
|
|
c_or_r = xpx.atleast_nd(c_or_r, ndim=1, xp=xp)
|
|
if c_or_r.ndim > 1:
|
|
raise ValueError("Polynomial must be 1d only.")
|
|
c_or_r = _trim_zeros(c_or_r, trim='f')
|
|
if c_or_r.shape[0] == 0:
|
|
c_or_r = xp.asarray([0], dtype=c_or_r.dtype)
|
|
return c_or_r
|
|
|
|
|
|
# https://github.com/numpy/numpy/blob/v2.2.0/numpy/lib/_polynomial_impl.py#L702-L779
|
|
def polyval(p, x, *, xp):
|
|
""" Old-style polynomial, `np.polyval`
|
|
"""
|
|
y = xp.zeros_like(x)
|
|
|
|
for pv in p:
|
|
y = y * x + pv
|
|
return y
|
|
|
|
|
|
# https://github.com/numpy/numpy/blob/v2.2.0/numpy/lib/_polynomial_impl.py#L34-L157
|
|
def poly(seq_of_zeros, *, xp):
|
|
# Only reproduce the 1D variant of np.poly
|
|
seq_of_zeros = xp.asarray(seq_of_zeros)
|
|
seq_of_zeros = xpx.atleast_nd(seq_of_zeros, ndim=1, xp=xp)
|
|
|
|
if seq_of_zeros.shape[0] == 0:
|
|
return 1.0
|
|
|
|
# prefer np.convolve etc, if available
|
|
convolve_func = getattr(xp, 'convolve', None)
|
|
if convolve_func is None:
|
|
from scipy.signal import convolve as convolve_func
|
|
|
|
dt = seq_of_zeros.dtype
|
|
a = xp.ones((1,), dtype=dt)
|
|
one = xp.ones_like(seq_of_zeros[0])
|
|
for zero in seq_of_zeros:
|
|
a = convolve_func(a, xp.stack((one, -zero)), mode='full')
|
|
|
|
if xp.isdtype(a.dtype, 'complex floating'):
|
|
# if complex roots are all complex conjugates, the roots are real.
|
|
roots = xp.asarray(seq_of_zeros, dtype=xp.complex128)
|
|
if xp.all(xp.sort(xp.imag(roots)) == xp.sort(xp.imag(xp.conj(roots)))):
|
|
a = xp.asarray(xp.real(a), copy=True)
|
|
|
|
return a
|
|
|
|
|
|
# https://github.com/numpy/numpy/blob/v2.2.0/numpy/lib/_polynomial_impl.py#L912
|
|
def polymul(a1, a2, *, xp):
|
|
a1, a2 = _poly1d(a1, xp=xp), _poly1d(a2, xp=xp)
|
|
|
|
# prefer np.convolve etc, if available
|
|
convolve_func = getattr(xp, 'convolve', None)
|
|
if convolve_func is None:
|
|
from scipy.signal import convolve as convolve_func
|
|
|
|
val = convolve_func(a1, a2)
|
|
return val
|
|
|
|
|
|
# ### New-style routines ###
|
|
|
|
|
|
# https://github.com/numpy/numpy/blob/v2.2.0/numpy/polynomial/polynomial.py#L663
|
|
def npp_polyval(x, c, *, xp, tensor=True):
|
|
if xp.isdtype(c.dtype, 'integral'):
|
|
c = xp.astype(c, xp_default_dtype(xp))
|
|
|
|
c = xpx.atleast_nd(c, ndim=1, xp=xp)
|
|
if isinstance(x, tuple | list):
|
|
x = xp.asarray(x)
|
|
if tensor:
|
|
c = xp.reshape(c, (c.shape + (1,)*x.ndim))
|
|
|
|
c0, _ = xp_promote(c[-1, ...], x, broadcast=True, xp=xp)
|
|
for i in range(2, c.shape[0] + 1):
|
|
c0 = c[-i, ...] + c0*x
|
|
return c0
|
|
|
|
|
|
# https://github.com/numpy/numpy/blob/v2.2.0/numpy/polynomial/polynomial.py#L758-L842
|
|
def npp_polyvalfromroots(x, r, *, xp, tensor=True):
|
|
r = xpx.atleast_nd(r, ndim=1, xp=xp)
|
|
# if r.dtype.char in '?bBhHiIlLqQpP':
|
|
# r = r.astype(np.double)
|
|
|
|
if isinstance(x, tuple | list):
|
|
x = xp.asarray(x)
|
|
|
|
if tensor:
|
|
r = xp.reshape(r, r.shape + (1,) * x.ndim)
|
|
elif x.ndim >= r.ndim:
|
|
raise ValueError("x.ndim must be < r.ndim when tensor == False")
|
|
return xp.prod(x - r, axis=0)
|