image-match/python/py/Lib/site-packages/scipy/signal/_polyutils.py

"""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 warnings
import scipy._lib.array_api_extra as xpx
from scipy._lib._array_api import (
    xp_promote, xp_default_dtype, xp_size, xp_device, is_numpy
)

try:
    from numpy.exceptions import RankWarning
except ImportError:
    # numpy 1.x
    from numpy import RankWarning


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]


# For numpy arrays, use scipy.linalg.lstsq;
# For other backends,
#   - use xp.linalg.lstsq, if available (cupy, torch, jax.numpy);
#   - otherwise manually compute pseudoinverse via SVD factorization
def _lstsq(a, b, xp=None, rcond=None):
    a, b = xp_promote(a, b, force_floating=True, xp=xp)

    if rcond is None:
        rcond = xp.finfo(a.dtype).eps * max(a.shape[-1], a.shape[-2])

    if is_numpy(xp):
        from scipy.linalg import lstsq as s_lstsq
        return s_lstsq(a, b, cond=rcond)
    elif lstsq_func := getattr(xp.linalg, "lstsq", None):
        # cupy, torch, jax.numpy all have xp.linalg.lstsq
        return lstsq_func(a, b, rcond=rcond)
    else:
        # unknown array library: LSQ solve via pseudoinverse
        u, s, vt = xp.linalg.svd(a, full_matrices=False)

        sing_val_mask = s > rcond
        s = xpx.apply_where(sing_val_mask, (s,), lambda x: 1. / x, fill_value=0.)

        sigma = xp.eye(s.shape[0]) * s    # == np.diag(s)
        x = vt.T @ sigma @ u.T @ b

        rank = xp.count_nonzero(sing_val_mask)

        # XXX actually compute residuals, when there's a use case
        residuals = xp.asarray([])
        return x, residuals, rank, s


# ### 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`
    """
    p = xp.asarray(p)
    x = xp.asarray(x)
    y = xp.zeros_like(x)

    # NB: cannot do `for pv in p` since array API iteration
    # is only defined for 1D arrays.
    for j in range(p.shape[0]):
        y = y * x + p[j, ...]
    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 xp.asarray(1.0, dtype=xp.real(seq_of_zeros).dtype)

    # 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


# https://github.com/numpy/numpy/blob/v2.3.3/numpy/lib/_polynomial_impl.py#L459
def polyfit(x, y, deg, *, xp, rcond=None):
    # only reproduce the variant with full=False, w=None, cov=False
    order = int(deg) + 1
    x = xp.asarray(x)
    y = xp.asarray(y)
    x, y = xp_promote(x, y, force_floating=True, xp=xp)

    # check arguments.
    if deg < 0:
        raise ValueError("expected deg >= 0")
    if x.ndim != 1:
        raise TypeError("expected 1D vector for x")
    if xp_size(x) == 0:
        raise TypeError("expected non-empty vector for x")
    if y.ndim < 1 or y.ndim > 2:
        raise TypeError("expected 1D or 2D array for y")
    if x.shape[0] != y.shape[0]:
        raise TypeError("expected x and y to have same length")

    # set rcond
    if rcond is None:
        rcond = x.shape[0] * xp.finfo(x.dtype).eps

    # set up least squares equation for powers of x: lhs = vander(x, order)
    powers = xp.flip(xp.arange(order, dtype=x.dtype, device=xp_device(x)))
    lhs = x[:, None] ** powers[None, :]

    # scale lhs to improve condition number and solve
    scale = xp.sqrt(xp.sum(lhs * lhs, axis=0))
    lhs /= scale

    c, _, rank, _ = _lstsq(lhs, y, rcond=rcond, xp=xp)
    c = (c.T / scale).T  # broadcast scale coefficients

    # warn on rank reduction, which indicates an ill conditioned matrix
    if rank != order:
        msg = "Polyfit may be poorly conditioned"
        warnings.warn(msg, RankWarning, stacklevel=2)

    return c


# ### 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)