PyLops
Nasty bug in FFT2D when working with cupy arrays
Problem
The pylops.signalprocessing.FFT2D operator crashes when trying to deal with cupy arrays. This seems due to the fact that nfft is stored as tuple - note that this was not the case before the major refractoring of the FFT routines, reason why things used to work (too bad we can't have CI for cupy operators :( )
Basic example:
dt, dx = 0.005, 5
nt,nx = 100, 201
t = cp.arange(nt)*dt
x = cp.arange(nx)*dx
f0 = 10
nfft = 2**10
d = cp.outer(cp.sin(2*np.pi*f0*t), cp.arange(nx)+1)
FFTop = FFT2D(dims=[nt,nx], nffts=[nfft,nfft], sampling=[dt,dx])
dottest(FFTop, nfft*nfft, nt*nx, complexflag=2, verb=True, backend='cupy') # crashes
FFTop.nffts = list(FFTop.nffts)
dottest(FFTop, nfft*nfft, nt*nx, complexflag=2, verb=True, backend='cupy') # works
Error:
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-42-d1129b05654c> in <module>
8
9 FFTop = FFT2D(dims=[nt,nx], nffts=[nfft,nfft], sampling=[dt,dx])
---> 10 dottest(FFTop, nfft*nfft, nt*nx, complexflag=2, verb=True, backend='cupy')
11 FFTop.nffts = list(FFTop.nffts)
12
~/Documents/OpenSource/pylops/pylops/utils/dottest.py in dottest(Op, nr, nc, rtol, complexflag, raiseerror, verb, backend, atol, tol)
108 v = v + 1j * ncp.random.randn(nr).astype(rdtype)
109
--> 110 y = Op.matvec(u) # Op * u
111 x = Op.rmatvec(v) # Op'* v
112
~/Documents/OpenSource/pylops/pylops/LinearOperator.py in matvec(self, x)
141 raise ValueError("dimension mismatch")
142
--> 143 y = self._matvec(x)
144
145 if x.ndim == 1:
~/Documents/OpenSource/pylops/pylops/signalprocessing/FFT2D.py in _matvec(self, x)
72 else:
73 print(type(x))
---> 74 y = np.fft.fft2(x, s=self.nffts, axes=self.dirs, **self._norm_kwargs)
75 if self.norm is _FFTNorms.ONE_OVER_N:
76 y *= self._scale
<__array_function__ internals> in fft2(*args, **kwargs)
cupy/core/core.pyx in cupy.core.core.ndarray.__array_function__()
~/miniconda3/envs/pylops_cupy_cusignal/lib/python3.9/site-packages/cupy/fft/_fft.py in fft2(a, s, axes, norm)
640 """
641 func = _default_fft_func(a, s, axes)
--> 642 return func(a, s, axes, norm, cufft.CUFFT_FORWARD)
643
644
~/miniconda3/envs/pylops_cupy_cusignal/lib/python3.9/site-packages/cupy/fft/_fft.py in _fftn(a, s, axes, norm, direction, value_type, order, plan, overwrite_x, out)
532
533 # Note: need to call _cook_shape prior to sorting the axes
--> 534 a = _cook_shape(a, s, axes, value_type, order=order)
535
536 for n in a.shape:
~/miniconda3/envs/pylops_cupy_cusignal/lib/python3.9/site-packages/cupy/fft/_fft.py in _cook_shape(a, s, axes, value_type, order)
40
41 def _cook_shape(a, s, axes, value_type, order='C'):
---> 42 if s is None or s == a.shape:
43 return a
44 if (value_type == 'C2R') and (s[-1] is not None):
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-42-d1129b05654c> in <module>
8
9 FFTop = FFT2D(dims=[nt,nx], nffts=[nfft,nfft], sampling=[dt,dx])
---> 10 dottest(FFTop, nfft*nfft, nt*nx, complexflag=2, verb=True, backend='cupy')
11 FFTop.nffts = list(FFTop.nffts)
12
~/Documents/OpenSource/pylops/pylops/utils/dottest.py in dottest(Op, nr, nc, rtol, complexflag, raiseerror, verb, backend, atol, tol)
108 v = v + 1j * ncp.random.randn(nr).astype(rdtype)
109
--> 110 y = Op.matvec(u) # Op * u
111 x = Op.rmatvec(v) # Op'* v
112
~/Documents/OpenSource/pylops/pylops/LinearOperator.py in matvec(self, x)
141 raise ValueError("dimension mismatch")
142
--> 143 y = self._matvec(x)
144
145 if x.ndim == 1:
~/Documents/OpenSource/pylops/pylops/signalprocessing/FFT2D.py in _matvec(self, x)
72 else:
73 print(type(x))
---> 74 y = np.fft.fft2(x, s=self.nffts, axes=self.dirs, **self._norm_kwargs)
75 if self.norm is _FFTNorms.ONE_OVER_N:
76 y *= self._scale
<__array_function__ internals> in fft2(*args, **kwargs)
cupy/core/core.pyx in cupy.core.core.ndarray.__array_function__()
~/miniconda3/envs/pylops_cupy_cusignal/lib/python3.9/site-packages/cupy/fft/_fft.py in fft2(a, s, axes, norm)
640 """
641 func = _default_fft_func(a, s, axes)
--> 642 return func(a, s, axes, norm, cufft.CUFFT_FORWARD)
643
644
~/miniconda3/envs/pylops_cupy_cusignal/lib/python3.9/site-packages/cupy/fft/_fft.py in _fftn(a, s, axes, norm, direction, value_type, order, plan, overwrite_x, out)
532
533 # Note: need to call _cook_shape prior to sorting the axes
--> 534 a = _cook_shape(a, s, axes, value_type, order=order)
535
536 for n in a.shape:
~/miniconda3/envs/pylops_cupy_cusignal/lib/python3.9/site-packages/cupy/fft/_fft.py in _cook_shape(a, s, axes, value_type, order)
40
41 def _cook_shape(a, s, axes, value_type, order='C'):
---> 42 if s is None or s == a.shape:
43 return a
44 if (value_type == 'C2R') and (s[-1] is not None):
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
Bugfix
I suggest to reconsider the choice of using a tuple for the nfft parameter all together. Or, if this has some special reason (@cako) we may need to introduce the usual logic to check if x is a cupy array and convert it into a list
This is very similar to the bug in Intel MKL that I reported and submitted a PR for, but remains unfixed (https://github.com/IntelPython/mkl_fft/pull/69). IMO this is an upstream bug we should probably report, but for now I agree we can change nfft to list exceptionally for this operator.
The reason why I chose nfft as tuple and not list is because I wanted to keep it compatible with dimsd and dims which are tuples. It also made sense from a design point of view because once the operator is set up, nfft should not change, and therefore it should not be mutable in principle.
Makes sense. I guess here the situation is reversed as cupy doesn’t seem to like tuples.
I definitely see the point of using tuples for this parameter as it is not going to change after initialization. Let me like a GitHub issue.
However, given that it will take some time for this to change (provided the cupy folks allow for such a change), let’s go ahead and fix it on your side.
@cako do you want to do it, or shall I?
Wait, I got this wrong. Actually if nnfts was stored as tuple everything would be fine, the issue is that it is stored as np.array.
a = (100, 100)
F = pylops.signalprocessing.FFT2D(a)
print(type(F.nffts))
gives <class 'numpy.ndarray'>.
So yes you are right, this is the same issue and IntelPython/MKL but now I am more tempted to say this is a problem on our hand, not on theirs. Numpy documentation says sequence of ints which is a bit ambiguous, what is a sequence. That could be a list, tuple, a np.array. Cupy documentation is actually more precise and says tuple of ints. So I dont think we will get them to change to allow types that are less natural.
Let's just make sure our nffts and axes are always TUPLE. What do you think?
I fully agree. We should just fix it on our end, and making it a tuple should do the trick. I don't have a lot of time on my hands this week, so if you have some spare time go ahead!
I don't have a GPU readily available at the moment. Could you have a look at it?
Sure, let me do it.
I think I am going to make sure dims is always a tuple, then I’m sure it will work as this is what I do outside of the operator as workaround now ;)