Add cupy.poly1d

cupy/lib/polynomial.pyx

@@ -67,13 +67,17 @@ cdef class poly1d:
         return self.order
 
     def __init__(self, c_or_r, r=False, variable=None):
+        if isinstance(c_or_r, numpy.poly1d):
+           self._coeffs = cupy.empty(c_or_r.coeffs.shape,
+                                     c_or_r.coeffs.dtype)
+           self._coeffs.set(c_or_r.coeffs)
         if isinstance(c_or_r, poly1d):
-            self._variable = c_or_r._variable
             self._coeffs = c_or_r._coeffs
-
-            if variable is not None:
+        if isinstance(c_or_r, (numpy.poly1d, poly1d)):
+           self._variable = c_or_r._variable
+           if variable is not None:
                 self._variable = variable
-            return
+           return
         # TODO(Dahlia-Chehata): if r: c_or_r = poly(c_or_r)
         if c_or_r.ndim < 1:
             c_or_r = cupy.atleast_1d(c_or_r)
@@ -198,7 +202,7 @@ cdef class poly1d:
     # -------------------------------------------------------------------------
     # Cupy specific attributes and methods
     # -------------------------------------------------------------------------
-    def get(self, stream=None, out=None):
+    cpdef get(self, stream=None, out=None):
         """Returns a copy of poly1d object on host memory.
 
         Args:
@@ -230,7 +234,7 @@ cdef class poly1d:
         return numpy.poly1d(self.coeffs.get(stream=stream),
                             variable=self.variable)
 
-    def set(self, polyin, stream=None):
+    cpdef set(self, polyin, stream=None):
         """Copies a poly1d object on the host memory to :class:`cupy.poly1d`.
 
         Args:

tests/cupy_tests/lib_tests/test_polynomial.py

@@ -25,11 +25,9 @@ def test_poly1d_cp_poly1d(self, xp, dtype):
 
     @testing.for_all_dtypes(no_bool=True)
     def test_poly1d_np_poly1d(self, dtype):
-        arr1 = testing.shaped_arange((10,), cupy, dtype)
-        arr2 = numpy.ones(10, dtype=dtype)
-        a = cupy.poly1d(arr1)
-        b = numpy.poly1d(arr2)
-        a.set(b)
+        arr = testing.shaped_arange((10,), numpy, dtype)
+        a = numpy.poly1d(arr)
+        b = cupy.poly1d(a)
         testing.assert_array_equal(a.coeffs, b.coeffs)
 
     @testing.for_all_dtypes()
@@ -106,6 +104,15 @@ def test_poly1d_get2(self, dtype):
         b2 = numpy.poly1d(a2)
         assert b1 == b2
 
+    @testing.for_all_dtypes(no_bool=True)
+    def test_poly1d_set(self, dtype):
+        arr1 = testing.shaped_arange((10,), cupy, dtype)
+        arr2 = numpy.ones(10, dtype=dtype)
+        a = cupy.poly1d(arr1)
+        b = numpy.poly1d(arr2)
+        a.set(b)
+        testing.assert_array_equal(a.coeffs, b.coeffs)
+
     @testing.for_all_dtypes()
     @testing.numpy_cupy_equal()
     def test_poly1d_repr(self, xp, dtype):