Repository: cupy/cupy

Issue:
Add `cupy.poly`
blocked by #3828


Discussion:

Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -538,6 +538,7 @@ def isscalar(element):
 # Polynomial functions
 # ------------------------------------------------------------------------------
 from cupy.lib._polynomial import poly1d  # NOQA
+from cupy.lib._routines_poly import poly  # NOQA
 from cupy.lib._routines_poly import polyadd  # NOQA
 from cupy.lib._routines_poly import polysub  # NOQA
 from cupy.lib._routines_poly import polymul  # NOQA
diff --git a/cupy/_math/misc.py b/cupy/_math/misc.py
--- a/cupy/_math/misc.py
+++ b/cupy/_math/misc.py
@@ -80,9 +80,9 @@ def convolve(a, v, mode='full'):
 def _fft_convolve(a1, a2, mode):
 
     offset = 0
-    if a1.size < a2.size:
+    if a1.shape[-1] < a2.shape[-1]:
         a1, a2 = a2, a1
-        offset = 1 - a2.size % 2
+        offset = 1 - a2.shape[-1] % 2
 
     # if either of them is complex, the dtype after multiplication will also be
     if a1.dtype.kind == 'c' or a2.dtype.kind == 'c':
@@ -91,7 +91,7 @@ def _fft_convolve(a1, a2, mode):
         fft, ifft = cupy.fft.rfft, cupy.fft.irfft
 
     dtype = cupy.result_type(a1, a2)
-    n1, n2 = a1.size, a2.size
+    n1, n2 = a1.shape[-1], a2.shape[-1]
     out_size = cupyx.scipy.fft.next_fast_len(n1 + n2 - 1)
     fa1 = fft(a1, out_size)
     fa2 = fft(a2, out_size)
@@ -108,7 +108,7 @@ def _fft_convolve(a1, a2, mode):
         raise ValueError(
             'acceptable mode flags are `valid`, `same`, or `full`.')
 
-    out = out[start:end]
+    out = out[..., start:end]
 
     if dtype.kind in 'iu':
         out = cupy.around(out)
diff --git a/cupy/lib/_routines_poly.py b/cupy/lib/_routines_poly.py
--- a/cupy/lib/_routines_poly.py
+++ b/cupy/lib/_routines_poly.py
@@ -34,6 +34,49 @@ def wrapper(*args):
     return functools.update_wrapper(wrapper, func)
 
 
+def poly(seq_of_zeros):
+    """Computes the coefficients of a polynomial with the given roots sequence.
+
+    Args:
+        seq_of_zeros (cupy.ndarray): a sequence of polynomial roots.
+
+    Returns:
+        cupy.ndarray: polynomial coefficients from highest to lowest degree.
+
+    .. warning::
+
+        This function doesn't support general 2d square arrays currently.
+        Only complex Hermitian and real symmetric 2d arrays are allowed.
+
+    .. seealso:: :func:`numpy.poly`
+
+    """
+    x = seq_of_zeros
+    if x.ndim == 2 and x.shape[0] == x.shape[1] and x.shape[0] != 0:
+        if cupy.array_equal(x, x.conj().T):
+            x = cupy.linalg.eigvalsh(x)
+        else:
+            raise NotImplementedError('Only complex Hermitian and real '
+                                      'symmetric 2d arrays are supported '
+                                      'currently')
+    elif x.ndim == 1:
+        x = x.astype(cupy.mintypecode(x.dtype.char), copy=False)
+    else:
+        raise ValueError('Input must be 1d or non-empty square 2d array.')
+
+    if x.size == 0:
+        return 1.0
+
+    size = 2 ** (x.size - 1).bit_length()
+    a = cupy.zeros((size, 2), x.dtype)
+    a[:, 0].fill(1)
+    cupy.negative(x, out=a[:x.size, 1])
+    while size > 1:
+        size = size // 2
+        a = cupy._math.misc._fft_convolve(a[:size], a[size:], 'full')
+    return a[0, :x.size + 1]
+
+
 @_wraps_polyroutine
 def polyadd(a1, a2):
     """Computes the sum of two polynomials.


Test patch:
diff --git a/tests/cupy_tests/lib_tests/test_polynomial.py b/tests/cupy_tests/lib_tests/test_polynomial.py
--- a/tests/cupy_tests/lib_tests/test_polynomial.py
+++ b/tests/cupy_tests/lib_tests/test_polynomial.py
@@ -98,6 +98,14 @@ def test_poly1d_zero_size(self, xp, dtype):
         assert out.variable == (self.variable or 'x')
         return out
 
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-5, atol=1e-5)
+    def test_poly1d_roots(self, xp, dtype):
+        a = testing.shaped_arange((4,), xp, dtype)
+        out = xp.poly1d(a, True, variable=self.variable)
+        assert out.variable == (self.variable or 'x')
+        return out.coeffs
+
 
 @testing.gpu
 class TestPoly1d:
@@ -236,6 +244,63 @@ def test_poly1d_call(self, xp, dtype):
         return b(a)
 
 
+@testing.gpu
+class TestPoly:
+
+    @testing.for_all_dtypes(no_bool=True)
+    @testing.numpy_cupy_allclose(rtol=1e-4)
+    def test_poly_1d(self, xp, dtype):
+        a = testing.shaped_arange((5,), xp, dtype)
+        return xp.poly(a)
+
+    @testing.for_all_dtypes(no_bool=True, no_float16=True, no_complex=True)
+    @testing.numpy_cupy_allclose(rtol=1e-4)
+    def test_poly_2d_symmetric_real(self, xp, dtype):
+        a = xp.array([[6, 3, 1],
+                      [3, 0, 5],
+                      [1, 5, 6]], dtype)
+        return xp.poly(a)
+
+    @testing.for_complex_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_poly_2d_hermitian_complex(self, xp, dtype):
+        a = xp.array([[2, -1j], [1j, 1]], dtype)
+        return xp.poly(a)
+
+    @testing.for_all_dtypes(no_bool=True)
+    def test_poly_2d_square(self, dtype):
+        a = testing.shaped_arange((3, 3), cupy, dtype)
+        with pytest.raises(NotImplementedError):
+            cupy.poly(a)
+
+    @testing.for_all_dtypes()
+    def test_poly_2d_general(self, dtype):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((2, 4), xp, dtype)
+            with pytest.raises(ValueError):
+                xp.poly(a)
+
+    @testing.for_all_dtypes()
+    def test_poly_ndim(self, dtype):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((5, 4, 3), xp, dtype)
+            with pytest.raises(ValueError):
+                xp.poly(a)
+
+    @testing.for_all_dtypes(no_bool=True)
+    def test_poly_zero_dim(self, dtype):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((), xp, dtype)
+            with pytest.raises(TypeError):
+                numpy.poly(a)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_equal()
+    def test_poly_empty(self, xp, dtype):
+        a = xp.zeros((0), dtype)
+        return xp.poly(a)
+
+
 @testing.gpu
 @testing.parameterize(*testing.product({
     'shape': [(), (0,), (5,)],