pythonscipyconvolutionfftwpyfftw

Convolution product in pyFFTW different from scipy


I'm trying to implement a FFT convolution that mimics scipy.fftconvolve using pyfftw for performance and pictures as input :

import numpy as np
import pyfftw
a = np.ones((6000, 4000), dtype='float32')
b = np.kaiser(25, 8)
b = np.outer(b, b).astype('float32')

class fftconvolve:
    def __init__(self, A, B, domain, threads=8):
        MK =  B.shape[0]
        NK = B.shape[1]
        M = A.shape[0]
        N = A.shape[1]

        if domain =="same":
            Y = M
            X = N
        elif domain == "valid":
            Y = M - MK + 1
            X = N - NK + 1
        elif domain == "full":
            Y = M + MK - 1
            X = N + NK - 1

        self.fft_A_obj = pyfftw.builders.rfft2(A, s=(M + MK -1, N + NK -1), threads=threads)
        self.fft_B_obj = pyfftw.builders.rfft2(B, s=(M + MK -1, N + NK -1), threads=threads)
        self.ifft_obj = pyfftw.builders.irfft2(self.fft_A_obj.output_array, s=(Y, X), threads=threads)

    def __call__(self, A, B):
        return self.ifft_obj(np.fft.ifftshift(
            np.fft.fftshift(self.fft_A_obj(A)) * np.fft.fftshift(self.fft_B_obj(B))
        ))

Calling it :

plan = fftconvolve(a, b, "full", threads=8)
c_1 = plan(a, b)
c_1

Output :

array([[  3.89971137e-06,   3.51986018e-05,   1.24518745e-04, ...,
          1.25271297e-04,   3.56316777e-05,   4.04627326e-06],
       [  4.91737483e-05,   2.60021159e-04,   8.61040782e-04, ...,
          8.63055116e-04,   2.61142646e-04,   4.95371969e-05],
       [  1.26523402e-04,   8.49825097e-04,   2.90915114e-03, ...,
          2.90881563e-03,   8.49568460e-04,   1.26304061e-04],
       ..., 
       [  1.28503540e-04,   8.52331228e-04,   2.91197700e-03, ...,
          2.91016186e-03,   8.51134886e-04,   1.28111642e-04],
       [  2.14206957e-05,   2.32703838e-04,   8.34190170e-04, ...,
          8.34319100e-04,   2.32750244e-04,   2.14206957e-05],
       [ -8.42595455e-06,   2.29651105e-05,   1.12404508e-04, ...,
          1.12760317e-04,   2.31778213e-05,  -8.35505125e-06]], dtype=float32)

Calling the scipy equivalent :

c_2 = scipy.signal.fftconvolve(a, b, "full").astype(np.float32)
c_2

Output :

array([[  5.47012860e-06,   3.68362089e-05,   1.26135841e-04, ...,
          1.26135841e-04,   3.68362089e-05,   5.47012769e-06],
       [  3.68362089e-05,   2.48057506e-04,   8.49407224e-04, ...,
          8.49407224e-04,   2.48057506e-04,   3.68362089e-05],
       [  1.26135841e-04,   8.49407224e-04,   2.90856976e-03, ...,
          2.90856976e-03,   8.49407224e-04,   1.26135841e-04],
       ..., 
       [  1.26135841e-04,   8.49407224e-04,   2.90856976e-03, ...,
          2.90856976e-03,   8.49407224e-04,   1.26135841e-04],
       [  3.68362089e-05,   2.48057506e-04,   8.49407224e-04, ...,
          8.49407224e-04,   2.48057506e-04,   3.68362089e-05],
       [  5.47012814e-06,   3.68362089e-05,   1.26135841e-04, ...,
          1.26135841e-04,   3.68362089e-05,   5.47012814e-06]], dtype=float32)

Checking the outputs :

c_1 == c_2

Gives :

array([[False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       ..., 
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False],
       [False, False, False, ..., False, False, False]], dtype=bool)

And :

np.allclose(c_1, c_2)

gives :

False

So the output is not correct. Removing the fftshift does not change a thing.

In my project, the scipy version gives a correct image, the pyfftw with my implementation gives blurry output.

Edit

I tested also in double type (np.float64) and while the raw result of the convolution is close enough (actually, scipy does the convolution in double), the picture is still bad :

Deconvolved with scipy : enter image description here

Deconvolved with custom convolution here : it's not only blurry but you get fringes on the edges : enter image description here

a = np.ones((6000, 4000), dtype='float64')
b = np.kaiser(25, 8)
b = np.outer(b, b).astype('float64')

Now :

np.allclose(c_1, c_2)

Returns :

True

What could give this result ?


Solution

  • This works as expected :

    class fftconvolve:
        def __init__(self, A, B, domain, threads=8):
            MK =  B.shape[0]
            NK = B.shape[1]
            M = A.shape[0]
            N = A.shape[1]
    
            if domain =="same":
                self.Y = M
                self.X = N
            elif domain == "valid":
                self.Y = M - MK + 1
                self.X = N - NK + 1
            elif domain == "full":
                self.Y = M + MK - 1
                self.X = N + NK - 1
    
            self.M = M + MK - 1
            self.N = N + NK - 1
    
            a = np.pad(A, ((0, MK - 1), (0, NK - 1)), mode='constant')
            b = np.pad(B, ((0, M - 1), (0, N - 1)), mode='constant')
    
            self.fft_A_obj = pyfftw.builders.rfft2(a, s=(self.M, self.N), threads=threads)
            self.fft_B_obj = pyfftw.builders.rfft2(b, s=(self.M, self.N), threads=threads)
            self.ifft_obj = pyfftw.builders.irfft2(self.fft_A_obj.output_array, s=(self.M, self.N), threads=threads)
    
    
            self.offset_Y = int(np.floor((self.M - self.Y)/2))
            self.offset_X = int(np.floor((self.N - self.X)/2))
    
        def __call__(self, A, B):
            MK =  B.shape[0]
            NK = B.shape[1]
            M = A.shape[0]
            N = A.shape[1]
    
            a = np.pad(A, ((0, MK - 1), (0, NK - 1)), mode='constant')
            b = np.pad(B, ((0, M - 1), (0, N - 1)), mode='constant')
    
            return self.ifft_obj(self.fft_A_obj(a) * self.fft_B_obj(b))[self.offset_Y:self.offset_Y + self.Y, self.offset_X:self.offset_X + self.X]