Fourier-domain BANE

[AlecThomson]

Hi @PaulHancock,

@tjgalvin and I have been experimenting with using a Fourier approach for BANE. The motivation for this was to boost the speed, especially when running on large cubes (such as those produced for polarisation products). We were also inspired by the Fourier approach used by WSClean for background/noise estimation.

I've put together a performant script that uses Numba for some very nice speed ups. This approach can blast through large cubes in short order. I've also added some options for handling 3D and 4D cubes which might be useful for regular BANE too.

The last outstanding issue is that is seems that output images are offset up and right by the size of the convolution kernel (a.k.a box size). I'll need to do some deeper digging to figure out how to fix that. In the meantime if you have any thoughts on that, please let me know!

Anyway, I'd be really keen to get your thoughts on this kind of approach :)

[PaulHancock]

Need to include numba and radio_beam as dependencies.

[PaulHancock]

I'm getting an error when running that i'm not sure how to fix:

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'rfft2' of type Module(<module 'numpy.fft' from '/data/alpha/hancock/AlecThomson/Aegean/env/lib/python3.8/site-packages/numpy/fft/__init__.py'>)

File "AegeanTools/BANE_fft.py", line 69:
def fft_average(image: np.ndarray, kernel: np.ndarray, kern_sum: float) -> np.ndarray:
    <source elided>
    """
    image_fft = fft.rfft2(image)
    ^

During: typing of get attribute at /data/alpha/hancock/AlecThomson/Aegean/AegeanTools/BANE_fft.py (69)

File "AegeanTools/BANE_fft.py", line 69:
def fft_average(image: np.ndarray, kernel: np.ndarray, kern_sum: float) -> np.ndarray:
    <source elided>
    """
    image_fft = fft.rfft2(image)
    ^

I'm not savvy enough with numba or python Typing to know how to address this.

[AlecThomson]

@PaulHancock what versions of Numpy and Numba are you rocking?

[AlecThomson]

Ah sorry @PaulHancock, I remembered some extra steps to get this to work. Using Numba and Numpy-FFT also requires rocket-fft.

It's a little sneaky since no import is required!

So the fix is a pip install rocket-fft - I'll add this to the deps as well

[PaulHancock]

Thanks @AlecThomson it's working now.