Use fast convolution algorithm in cupy.poly1d.__pow__
@@ -147,7 +147,7 @@ def polymul(a1, a2):
return cupy.convolve(a1, a2)
-def _polypow(x, n):
+def _polypow_direct(x, n):
if n == 0:
return 1
if n == 1:
@@ -157,6 +157,32 @@ def _polypow(x, n):
return cupy.convolve(x, _polypow(cupy.convolve(x, x), (n - 1) // 2))
+def _polypow(x, n):
+ if n == 0:
+ return 1
+ if n == 1:
+ return x
+
+ method = cupy._math.misc._choose_conv_method(x, x, 'full')
+
+ if method == 'direct':
+ return _polypow_direct(x, n)
+ elif method == 'fft':
+ if x.dtype.kind == 'c':
+ fft, ifft = cupy.fft.fft, cupy.fft.ifft
+ else:
+ fft, ifft = cupy.fft.rfft, cupy.fft.irfft
+ out_size = (x.size - 1) * n + 1
+ import cupyx.scipy.fft
+ size = cupyx.scipy.fft.next_fast_len(out_size)
+ fx = fft(x, size)
+ fy = cupy.power(fx, n, fx)
+ y = ifft(fy, size)
+ return y[:out_size]
+ else:
+ assert False
+
+
def _polyfit_typecast(x):
if x.dtype.kind == 'c':
return x.astype(numpy.complex128, copy=False)
