IDL code to compute a diffraction pattern (and a Python port)

aperture.pro

; Test routine to generate a circular pupil and compute the
; PSF and OTF associated with it.
pro aperture,nd,rado,radi,flag


;nd   - array size
;rado - outer radius (pixels)
;radi - inner radius (pixels)
;flag - 0 - circular aperture
;     - 1 - square aperture
;     - 2 - hexagonal aperture


; Set Parameters


nd2=nd/2.
nd4=nd/4.
;print,nd2+nd4/2.
pi=3.14159
ns=rado*2.


; Generate Pupil from input parameters


pupil=fltarr(nd,nd)
if(flag eq 0) then begin                 ;Circular Pupil
 pupil(where(radmap(pupil) le rado))=1.0
 if(radi eq 0.) then begin
  pupil=pupil
 endif else begin
  pupili=fltarr(nd,nd)
  pupili(where(radmap(pupili) le radi))=1.0
  pupil=pupil-pupili
 endelse
endif else begin                 ;Square Pupil
 if(flag eq 1) then begin
  pupil(nd2-rado:nd2-1+rado,nd2-rado:nd2-1+rado)=1
  if(radi eq 0) then begin
   pupil=pupil
  endif else begin
   pupili=fltarr(nd,nd)
   pupili(nd2-radi:nd2-1+radi,nd2-radi:nd2-1+radi)=1
   pupil=pupil-pupili
  endelse
 endif else begin
  pupil = hex(nd,rado)
  pupili = hex(nd,radi)
  pupil=pupil-pupili
 endelse
endelse


; Plot pupil


window,0,xsize=nd,ysize=nd,title="Pupil",xpos=0,ypos=0
tvscl,pupil
print,pupil


; Compute and display Pupil Power Spectrum (to produce PSF)


a=abs(fft(pupil,1))^2
psf=fftshift(a,nd)
window,1,xsize=nd,ysize=nd,title="PSF",xpos=0,ypos=nd+26
shade_surf,psf(nd4:nd-nd4,nd4:nd-nd4)
window,5,xsize=nd*1.5,ysize=nd,title="PSF",xpos=nd+8,ypos=0+(nd+26)
plot,psf(nd2:nd2+nd4/2.,nd2:nd2)
oplot,psf(nd2:nd2,nd2:nd2+nd4/2.),color=0,linestyle=5
window,7,xsize=nd,ysize=nd,title="PSF",xpos=nd+8,ypos=0
tvscl,sqrt(sqrt(psf))

; Compute and display MTF from PSF


mtf=abs(fft(a,/inverse))
mtfs=fftshift(mtf,nd)
window,2,xsize=nd,ysize=nd,title="MTF",xpos=0,ypos=2*(nd+26)
shade_surf,mtfs
window,6,xsize=nd*1.5,ysize=nd,title="MTF",xpos=nd+8,ypos=2*(nd+26)
if(ns lt nd2) then begin
 plot,mtfs(nd2:nd2+ns,nd2:nd2)
 oplot,mtfs(nd2:nd2,nd2:nd2+ns),color=0,linestyle=5
endif else begin
 plot,mtfs(nd2:nd-1,nd2:nd2)
 oplot,mtfs(nd2:nd2,nd2:nd-1),color=0,linestyle=5
endelse
window,8,xsize=nd,ysize=nd,title="MTF",xpos=2*(nd+8),ypos=0
tvscl,sqrt(sqrt(mtfs))


return
end


; This function rearranges the the FFT to put the DC point at the center of the array


function fftshift,data,nd
nd2=nd/2.
a=fltarr(nd,nd)
a(0:nd2-1,0:nd2-1)=data(nd2:nd-1,nd2:nd-1)
a(nd2:nd-1,0:nd2-1)=data(0:nd2-1,nd2:nd-1)
a(0:nd2-1,nd2:nd-1)=data(nd2:nd-1,0:nd2-1)
a(nd2:nd-1,nd2:nd-1)=data(0:nd2-1,0:nd2-1)
return,a
end


; This function will generate a hexagonal aperture


function hex,nd,rad


;Test routine to generate an array with a binary hexagonal
;mask based on straight lines.


pi = 3.14159
dtr = pi / 180.
nd2 = nd / 2.
x1 = nd2 - rad
x2 = x1 + rad * cos(60. * dtr)
y1 = nd2 + rad * cos(30. * dtr)
m = (rad * cos(30. * dtr)) / (x2 - x1)
;print,rad, x1, y1, m
test = fltarr(nd,nd)


for i = nd2, y1 do begin
 for j = x1, x2 do begin
  y = m * (j - x1)
  if(i-nd2 lt y) then begin
   test(j,i) = 1.0
  endif else begin
   test(j,i) = 0.0
  endelse
 endfor
 for j= x2, nd2 do begin
  test(j,i) = 1.0
 endfor
endfor
for i = nd2, nd-1 do begin
 for j = 1, nd2-1 do begin
  jj = nd - j -1
  test(jj,i) = test(j,i)
 endfor
endfor
for i = nd2, nd-1 do begin
 ii = nd -i - 1
 for j = 1, nd-1 do begin
  test(j,ii) = test(j,i)
 endfor
endfor
;window,0,xsize=nd,ysize=nd,title="Pupil",xpos=0,ypos=750
;tvscl,test


return,test
end

function radmap,inarray,x0,y0
;function to generate a radius map to match array inarray from center x0,y0

sizein=size(inarray)
if n_elements(x0) eq 0 then x0=sizein(1)/2.0 - 0.5
if n_elements(y0) eq 0 then y0=sizein(2)/2.0 - 0.5
xs=sizein(1)
ys=sizein(2)

dists=fltarr(xs,ys)
x=fltarr(xs,ys)
y=fltarr(xs,ys)
for i=0,ys-1 do x(*,i)=findgen(xs)
for i=0,xs-1 do y(i,*)=findgen(ys)
dsq=(x-x0)^2+(y-y0)^2
dists=sqrt(dsq)

;for i=0,xs-1 do begin
;  for j=0,ys-1 do begin
;    dists(i,j)=norm([float(i)-x0,float(j)-y0])
;  endfor
;endfor

return,dists
end

aperture.py

import numpy as np
from matplotlib.pylab import *
pylab.rcParams['figure.figsize'] = (10.0, 8.0)
dim = 400
outer_radius = 100
inner_radius = 25 # unused
pupil = np.array(
    [[1 if (a-dim/2)**2 + (b-dim/2)**2 <= outer_radius**2 else 0 for b in range(0,dim)] for a in range(0,dim)], # ugly one liners...
    dtype="f"
)
plt.imshow(pupil)

def shift_half(pupilfft):
    dim = pupilfft.shape[0]
    assert dim % 2 == 0, "Need even-numbered pixel dimensions, got {0}".format(dim)
    
    shifted = np.zeros(pupilfft.shape)
    shifted[0:dim/2,0:dim/2] = pupilfft[dim/2:dim,dim/2:dim]
    shifted[dim/2:dim,0:dim/2] = pupilfft[0:dim/2,dim/2:dim]
    shifted[0:dim/2,dim/2:dim] = pupilfft[dim/2:dim,0:dim/2]
    shifted[dim/2:dim,dim/2:dim] = pupilfft[0:dim/2,0:dim/2]
    return shifted


pupilfft = np.fft.ifft2(pupil) # 2nd positional arg in idl > 0 means inverse
absfft = np.abs(pupilfft)**2
psf = shift_half(absfft)

plt.imshow(sqrt(sqrt(psf)), cmap=cm.Greys_r) # img^(1/4) scaling applied like in IDL

fftshift.pro

function fftshift,data,nd
nd2=nd/2.
a=fltarr(nd,nd)
a(0:nd2-1,0:nd2-1)=data(nd2:nd-1,nd2:nd-1)
a(nd2:nd-1,0:nd2-1)=data(0:nd2-1,nd2:nd-1)
a(0:nd2-1,nd2:nd-1)=data(nd2:nd-1,0:nd2-1)
a(nd2:nd-1,nd2:nd-1)=data(0:nd2-1,0:nd2-1)
return,a
end