NumPy FFT producing off centre output

user_16464660 0
Viewed 118

TL;DR: NumPy FFT creates non uniform output when output is wanted to be uniform. I want the output to be a uniform corona.

I am trying to eventually run a Gerchberg-Saxton phase retrieval algorithm. I have been trying to make sure that I understand how FFT works in NumPy. I have used fftshift to create the correct looking output but the image does not have uniform intensity afterwards.

My input image is a circle, output should be a coronagraph looking thing from the circle aperture. I am trying to reproduce the results detailed in https://www.osapublishing.org/optica/fulltext.cfm?uri=optica-2-2-147&id=311836#articleSupplMat.

My algorithm to produce the error:

  1. Initial image, f
  2. FT(f)
  3. x exp ( i phase_mask)
  4. IFT(FT(f)x exp( i phase_mask)

Output, np.abs(g)**2

Phase mask, 2xtheta measured from origin

Input image, circle

Happy to clear anything up.

import numpy as np
import matplotlib.pyplot as plt
#Create 'pixels' for circle
pixels = 400
edge = np.linspace(-10, 10, pixels)
xv, yv = np.meshgrid(edge, edge)

def circle(x, y, r):
    '''
    x, y : dimensions of grid to place circle on
    r : radius
    Function defines aperture
    '''
    x0 = 0
    y0 = 0
    return np.select([((x-x0)**2+(y-y0)**2)>=r**2,
                      ((x-x0)**2+(y-y0)**2)<r**2],
                     [0,
                      1.])

#Create input and output images
radius = 4
input_img = circle(xv, yv, radius)
constraint_img = xcircle(xv, yv, radius)

img = input_img
constraint = 1 - img
max_iter = 10
re,im = np.mgrid[-1:1:400j, -1:1:400j] #Creates grid of values, 400=pixels
mask = 2*np.angle(re + 1j*im) #Gets angle from centre of grid
mask_i = mask


#Initial focal plane field, F. Initial image f.
f = np.sqrt(img)
F = np.fft.fftshift(np.fft.fft2(f)) * np.exp(mask * 1j)  #Focal plane field
F_1 = F 
am_f = np.abs(F_1) #Initial amplitude

g = np.fft.ifft2(F)
mask = np.angle(F/(F_1+1e-18))         #Final phase mask

recovery = (np.fft.ifft2(F*np.exp(-1j * mask)))
im3 = plt.imshow(np.abs(g)**2, cmap='gray')
plt.title('Recovered image')
plt.tight_layout()
plt.show()
plt.imshow(mask_i)
plt.colorbar()
plt.show()
fft signal-processing numpy python
1 Answers

Your issue is in this bit of code:

pixels = 400
edge = np.linspace(-10, 10, pixels)

as well as this one:

re,im = np.mgrid[-1:1:400j, -1:1:400j]

Because you use fftshift*, you need the origin to be at pixels//2. However, you don't sample the origin at all, it is in between two samples.

* You should really be using ifftshift instead, which moves the origin from pixels//2 to 0. fftshift moves the origin from 0 to pixels//2. For an even number of samples, these two do the same thing though.

To properly sample the origin, create edge as follows:

edge = np.linspace(-10, 10, pixels, endpoint=False)

We now see that edge[pixels//2] is equal to 0.

For np.mgrid there's no equivalent option. You will have to do this manually by creating one more sample, then deleting the last sample:

re,im = np.mgrid[-1:1:401j, -1:1:401j] #Creates grid of values, 400=pixels
mask = 2*np.angle(re + 1j*im) #Gets angle from centre of grid
mask = mask[:-1, :-1]

With these two changes, you will see a symmetric output.

user_73287820
Related