add reikna-backed fftn, ifftn, and fftshift

[tlambert03]

adds fftn, ifftn, fftshit, and fftconvolve using reikna (added as a dependency)

tested to match the output of scipy.fftpack

[codecov-commenter]

Codecov Report

Merging #113 (5d0b560) into master (5579f99) will increase coverage by 0.08%.
The diff coverage is 94.11%.

@@            Coverage Diff             @@
##           master     #113      +/-   ##
==========================================
+ Coverage   91.40%   91.49%   +0.08%     
==========================================
  Files         487      492       +5     
  Lines        8719     8936     +217     
==========================================
+ Hits         7970     8176     +206     
- Misses        749      760      +11     

Impacted Files	Coverage Δ
setup.py	0.00% <ø> (ø)
pyclesperanto_prototype/_fft/_fftconvolve.py	87.71% <87.71%> (ø)
pyclesperanto_prototype/_fft/_fft.py	92.85% <92.85%> (ø)
pyclesperanto_prototype/_fft/_fftshift.py	92.85% <92.85%> (ø)
tests/test_fft.py	98.76% <98.76%> (ø)
pyclesperanto_prototype/__init__.py	100.00% <100.00%> (ø)
pyclesperanto_prototype/_fft/__init__.py	100.00% <100.00%> (ø)
pyclesperanto_prototype/_tier0/_push.py	100.00% <100.00%> (ø)
... and 5 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 5579f99...5d0b560. Read the comment docs.

[haesleinhuepf]

Hey Talley @tlambert03 ,

that's an awesome PR! Just a minor concern as we are introducing new functions to the API: What do you think about the names of the functions. In clij2-fft there are similar functions with similar names. I would be happy to have things as similar as possible. For example on Java side we have convolveFFT and inverseFFT and here it's fftconvolve and ifftn. Would it be possible to unify those, e.g. by introducing some aliases, e.g. convolve_fft and inverse_fft?

Cheers,
Robert

[tlambert03]

We can name them whatever you want. I was just going with the python conventions... but just tell me what you want them to be

[tlambert03]

hey @haesleinhuepf, what's the canonical way to pad an array with zeros in cle? the equivalent of:

new_array = np.zeros(new_shape)
new_array[: old_array.shape[0], : old_array.shape[1]] = old_array

[tlambert03]

added aliases

# clij2 aliases
convolve_fft = fftconvolve
inverse = ifftn
forward_fft = fftn
fft = fftn

[tlambert03]

this is all set for review. the one thing that isn't supported yet is convolve_fft from an existing OCLArray, (waiting to hear how to pad with zeros)

[haesleinhuepf]

what's the canonical way to pad an array with zeros in cle? the equivalent of:

cle.set(new_array, 0)
cle.paste(old_array, new_array)

If you want to paste the old_array in the center of new_array, you should provide destination_x, destination_y and destination_z when calling paste.

[tlambert03]

cle.set(new_array, 0)
cle.paste(old_array, new_array)

Thanks @haesleinhuepf! am I correct in assuming that this won't currently work with complex dtypes? I'm running into some limitations here in terms of accepting existing GPU buffers... even if they come in as complex type already, will it be possible to resize them before #115?

I think maybe we can just leave the fftconvolve at only accepting numpy arrays for now... that ok?

[haesleinhuepf]

Thanks @haesleinhuepf! am I correct in assuming that this won't currently work with complex dtypes?

Yes, I think so. But haven't tested. Also on clij-side support for complex types is poor.

I think maybe we can just leave the fftconvolve at only accepting numpy arrays for now... that ok?

fftconvolve is from reikna, and I think we shouldn't modify with its API. But having cle.fftconvolve that doesn't work with cle-images/buffers sounds like introducing inconsistencies and in a later version we would need to break the backwards compatibility when we fix this. Maybe we take an additonal moment and build a wrapper/independence layer that allows us later modify the internal API without breaking it to the outside world. E.g. together with my suggestion above, we could do something like the following:

@plugin_function # receive numpy arrays, convert to OCLArrays
def convolve_fft(image : Image, kernel : Image , output : Image = None):
    # more conversion if necessary ...
    result = fftconvolve(...) 
    # conversion...
    return result

I also could imagine that the napari-clesperanto-assistant would break if it receives fftconvolve to build a user interface (haven't tested), because it expects functions that are plugin_function annotated...

Let me know what you think!

[tlambert03]

fftconvolve is from reikna, and I think we shouldn't modify with its API...Maybe we take an additonal moment and build a wrapper/independence layer that allows us later modify the internal API without breaking it to the outside world

not sure I'm following what you mean here. fftconvolve is not offered by reikna, it's a function that I wrote in this PR using reikna FFTs in the background... It's right here. I could do the whole thing directly in pyopencl, it's cle's lack of complex support that is limiting me, and there, only because I can't (easily) expand the shape of an already provided OCLArray that is complex.

[tlambert03]

I also could imagine that the napari-clesperanto-assistant would break if it receives fftconvolve to build a user interface (haven't tested), because it expects functions that are plugin_function annotated...

not following here either... what part of them need to be decorated by plugin_function?

[haesleinhuepf]

Ok, I obviously need more time to look into this in detail... :-)

[haesleinhuepf]

Hey @tlambert03 ,

I just dived into it a bit. I was thinking of making these changes, but the tests don't run and I must admit I don't know enough about pyopencl. Obviously, #115 needs to be fixed first before we can go ahead here.

Things that are kind of neccessary before merging:

• convolve_fft(in, psf, out) should accept numpy arrays and/or ocl-arrays as any other function in cle. Therefore, we need to annotate it with @plugin_function as suggested in my branch linked above.

Things I'm not sure about yet:

• All cle functions with tier>0 are @plugin_function annotated. If we now introduce new functions that are not, we should be sure that this is the way to go.
• All cle functions have parameters like this: [input-images, output-images, other parameters], which makes it easy to generate a magicgui user interface from it. We used a similar patter in clij for similar reasons. fftconvolve violates that pattern.
• No cle function accepts string parameters yet. If we introduce a function with string parameters now, I'm afraid we need to fix issues downstream, e.g. in the napari-pyclesperanto-assistant. I'm also not sure if a string is a good way for passing something that should rather be an enum. We also don't have enums yet in cle btw.
• The parameter axes might cause issues related to axis auto-transposition between python and opencl #49. Is it necessary?

In very general, I think the cle API should be more flexible and more pythonic. Thus, I appreciate your PR. It pushes us in this direction. But I also would like to keep things simple and straight. I'm a bit confused by all the parameters that fftconvolve has.

Not sure how to go ahead. Do you have bigger plans for this PR? Building deconvolution on top?

Best,
Robert

[tlambert03]

yeah, I suspected this might be an issue. I'm not really sure how to proceed either.

my personal motivation is to have a python FFT interface that mimics the standard python ecosystem (i.e. numpy and scipy.fftpack). I've had great success easily swapping out CPU code for cupy code, and would like to be able to do similar stuff for cl. I had been starting to make such an interface for myself, but then figured I'd contribute it to cle instead.

In general, I'd love to see cle be that library for python cle stuff... but things like #49, adherence to function signature that don't match the rest of the ecosystem, and the general need to mimic the java api definitely make that much more challenging.

I definitely understand your hesitation, given the divergence of these functions from the rest of cle... so I don't blame you if you'd rather not merge, or wait to consider it more. that said, these are the signatures that I think most python users will want (i.e. it's a drop in replacement for scipy.fftpack, as shown by the tests. Maybe lets just punt on this decision for now. I'll put these in another repo, which can maybe be used eventually by cle if you want to put a more cle-specific signature wrapper on it?

[haesleinhuepf]

I've had great success easily swapping out CPU code for cupy code

I also started doing this recently more often. I actually became a big fan of cucim. And I'm also thinking about how to make cle usable like scipy and cupy. So far I was thinking of wrapping cl_something_py around cle.

Btw. the fact that cupy (and cucim) don't take numpy-arrays as parameters is enoying. I'm happy that cle does.

I'll put these in another repo, which can maybe be used eventually by cle if you want to put a more cle-specific signature wrapper on it?

That sounds like a good idea! There is one more option we might at least think about: We start a new repository which has pyopencl, reikna and pyclesperanto as dependency. Maybe also other cl-stuff. There, we collect our_cl_something_py.fftpack.fftconvolve and our_cl_something_py.ndimage.rotate (see also). The only constraints: The API should fit to scipy and consume either numpy arrays or ocl-arrays as images. No pressure, we just put the functions in we need for our projects, or maybe also snippets of stuff we find on the way.
What do you think?

[jni]

Hi both! Some thoughts on this discussion:

• This issue is super relevant. Not clear guidance there yet but something to keep an eye on. Leveraging NEP 18 for image processing and signal processing scipy/scipy#10204
• I agree with @tlambert03 that we really really need a Pythonic/SciPythonic API more than we need a not-quite-java-not-quite-python API.
• For that, axes= is indeed essential (though it's unclear to me whether that should be axes= or axis=).
• It's also unclear to me that the above should live in a separate package. vs a sub package here, but 🤷. Either way, can it go inside the clesperanto org? ;)
• re accepting "either NumPy or OpenCL buffers", I suggest instead "either NumPy-coercible or OpenCL buffers". It would be pretty awesome to just be able to chain the output of a PyTorch inference together with a CuPy array of seeds into the Clesperanto watershed implementation!

[haesleinhuepf]

Either way, can it go inside the clesperanto org? ;)

Absolutely yes! I would love to build it myself. However, as I'm not using scipy.ndimage on a daily basis (not even weekly), I might not be the right lead for that project. I don't know the API well.

re accepting "either NumPy or OpenCL buffers", I suggest instead "either NumPy-coercible or OpenCL buffers". It would be pretty awesome to just be able to chain the output of a PyTorch inference together with a CuPy array of seeds into the Clesperanto watershed implementation!

Yes, and I worked hard to make those things better. I also would love to multiply cle and cupy images using the * operator (short-term by converting, long term by wrapping). That's something that could live in clEsperanto itself. I created a separate issue. Also a related issue here: If you know how to fix that on clesperanto-side, that would be very helpful! Or, if we fix it on napari side, also cupy would profit...

[tlambert03]

Btw, I already made a new package here: https://github.com/tlambert03/anyfft

It's not published yet. But it lets you easily change between FFT backends (using the best method available for your platform) while using the scipy api.

The flexibility of not being tied to the clij api there was nice, so if anything, I'd say clesperanto can depend on it if desired... but it now feels out of scope here

[tlambert03]

?

[haesleinhuepf]

?

After partial fixing #115 I wanted to give this another shot (Just had some days of meetings and paperwork, and needed something useful to code ;-) ), I just merged it (into a new branch), fixed some minor things and cut down the API to a minimum. Cool: There is some intermediate speedup compared to convolution in real space when processing large images and kernels. See:
https://github.com/clEsperanto/pyclesperanto_prototype/blob/fft/demo/filter/convolution.ipynb

It feels like there is a bottleneck I'm not fully aware of... But I'm still digging in it...

[beniroquai]

Hey @tlambert03 @haesleinhuepf, I think this is the thread I was looking for. FFT in clesperanto. Wuhu! Was there any progress being made? Is it "officially" being supported now? Trying out the notebook now.

[haesleinhuepf]

Hi @beniroquai ,
great to see you here! This wasn't merged into master yet, because I think there's a performance bottleneck somewhere. I presume data travels between CPU and GPU memory, but I didn't dive deeper. It has pretty low priority on my side. But if you need it and could spend some time on it, I'm happy to support!

[tlambert03]

you can also use try https://github.com/tlambert03/anyfft, if you'd like, which is where I put this pr

[psobolewskiPhD]

Not sure what you need, but if you want easy, performant OpenCL FFT, try pyvkfft (don't mind the name, it's not just vulkan).
https://github.com/vincefn/pyvkfft
It has a pyopencl backend (as well as CUDA) and is simple to use:
Just
from pyvkfft.fft import fftn, ifftn, rfftn, irfftn
Make sure you have a CL context and queue and push our data to the GPU before running the FFT. See:
https://github.com/vincefn/pyvkfft/blob/master/examples/pyvkfft-fft.ipynb

I think there should be no issues with using this along with pyclesperanto?
Though I'm not sure how exactly the queue part works out...

[beniroquai]

Hey, thanks @psobolewskiPhD for your reply. I may have seen that in the forum.imagesc.com ;-)
Do you have any code snippet that may connect the clesperanto with the pyvkfft? Not sure how to reuse GPU allocated memory between different python packages.

Thanks @tlambert03, I'll also investigate your repo! :)

[psobolewskiPhD]

I've not actually used both together 🤣
It's on the list, but not gotten around to it yet.
Sorry!
I think it should work though. pyclesperanto cle.select_device() will generate the context and queue:
device = cle.select_device()
which you can access

In [4]: device.context
Out[4]: <pyopencl.Context at 0x600002a214a0 on <pyopencl.Device 'Apple M1' on 'Apple' at 0x1027f00>>

In [5]: device.queue
Out[5]: <pyopencl._cl.CommandQueue at 0x1254f16d0>

So this seems to work for me:

img_GPU = cle.push(img)
THB = cle.top_hat_box(img_GPU, radius_x = 10, radius_y = 10, radius_z = 10)
img_fft = rfftn(THB) # note this is R2C, so complex output, which is not supported by clesperanto

This works, it looks like pyvkFFT can consume the arrays manipulated by pyclesperanto (pyclesperanto_prototype._tier0._pycl.OCLArray)
You can verify that img_fft = rfftn(img_GPU) works and img_fft = rfftn(img) does not, so I assume that no extra mem transfers are occurring.

But the other direction doesn't work for me. If we reverse transform back (C2R):

out = irfftn(img_fft)
gaus = cle.gaussian_blur(out, sigma_x= 10, sigma_y= 10, sigma_z= 10)

The cle command returns a 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)
Untitled-3 in <module>
----> <a href='untitled:Untitled-3?line=45'>46</a> gaus = cle.gaussian_blur(out, sigma_x= 10, sigma_y= 10, sigma_z= 10)

File ~/Dev/miniforge3/envs/apple-TF28/lib/python3.9/site-packages/pyclesperanto_prototype/_tier0/_plugin_function.py:68, in plugin_function.<locals>.worker_function(*args, **kwargs)
     66     if key in sig.parameters and sig.parameters[key].annotation is Image and value is None:
     67         sig2 = inspect.signature(output_creator)
---> 68         bound.arguments[key] = output_creator(*bound.args[:len(sig2.parameters)])
     70 # call the decorated function
     71 return function(*bound.args, **bound.kwargs)

File ~/Dev/miniforge3/envs/apple-TF28/lib/python3.9/site-packages/pyclesperanto_prototype/_tier0/_create.py:32, in create_like(*args)
     30 elif isinstance(dimensions, np.ndarray):
     31     dimensions = dimensions.shape[::-1]
---> 32 return create(dimensions)

File ~/Dev/miniforge3/envs/apple-TF28/lib/python3.9/site-packages/pyclesperanto_prototype/_tier0/_create.py:21, in create(dimensions, dtype)
      7 """
      8 Convenience method for creating images on the GPU. This method basically does the same as in CLIJ:
      9 
   (...)
     13 :return: OCLArray, potentially with random values
     14 """
     16 dimensions = (
...
   1553 else:
-> 1554     raise ValueError("The truth value of an array with "
   1555             "more than one element is ambiguous. Use a.any() or a.all()")

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Can test this more directly with:

import pyopencl.array as cla
pushed = cla.to_device(device.queue, img)
gaus = cle.gaussian_blur(pushed, sigma_x= 10, sigma_y= 10, sigma_z= 10)

That yields an error.
@haesleinhuepf any ideas? does pyclesperanto not support pyopencl.array.Array? is there a way to coerce it without extra roundtrip (get & re-push).

[haesleinhuepf]

Hey @psobolewskiPhD ,

import pyopencl.array as cla
pushed = cla.to_device(device.queue, img)
gaus = cle.gaussian_blur(pushed, sigma_x= 10, sigma_y= 10, sigma_z= 10)

That yields an error.
@haesleinhuepf any ideas? does pyclesperanto not support pyopencl.array.Array? is there a way to coerce it without extra roundtrip (get & re-push).

Can you please try again with pyclesperanto-prototype 0.17.2 ? I just added support for that array type. If an error remains, please open an issue.

@beniroquai : In it feels uncomfortable to discuss your use-case on a closed pull request, feel free to open an issue and explain what you're trying to achieve. There, we can concentrate on discussing your project.

Thanks to you two!

Best,
Robert

[psobolewskiPhD]

Everything works perfectly now, based on the few tests I did above.
I won't have time for more in-depth testing, but it's working!
🔥🎉

[beniroquai]

Thanks @psobolewskiPhD and @haesleinhuepf for your code snippets. I think that's already enough to continue working on the prototype (Holographic Reconstruction - the Ho in HoLiSheet 😉). Once I face the next issue, I may open a new issue. I actually didn't expect it to be that straightforward to reuse memory across python libraries. Cool!! :)