def fftshift(x, *args, **kwargs):
out = _fftshift(x, *args, **kwargs)
if hasattr(x, "unit"):
return out * x.unit
else:
return out
def fftshift(x, *args, **kwargs):
"""
The same as numpy's fftshift, except that it preserves units (if Astropy quantities are used)
All extra arguments are passed directly to numpy's `fftshift`.
"""
out = _fftshift(x, *args, **kwargs)
if hasattr(x, "unit"):
return out * x.unit
else:
return out
def ASM_crop(Fin, z, factor=5):
wavelength = Fin.lam
# N = Fin.N*factor
size = Fin.siz
k = 2 * _np.pi / wavelength
a1 = Fin.field
A1 = _fftshift(_fft2(_fftshift(a1)))
r1 = _np.linspace(-a1.shape[0] / 2, a1.shape[1] / 2 - 1, a1.shape[0])
s1 = _np.linspace(-a1.shape[0] / 2, a1.shape[1] / 2 - 1, a1.shape[0])
deltaFX = 1 / size * r1
deltaFY = 1 / size * s1
meshgrid = _np.meshgrid(deltaFX, deltaFY)
H = _np.exp(1.0j * k * z *
_np.sqrt(1 - _np.power(wavelength * meshgrid[0], 2) -
_np.power(wavelength * meshgrid[1], 2)))
U = _np.multiply(A1, H)
u = _ifftshift(_ifft2(_ifftshift(U)))
Fout = Fin
Fout.field = crop_intensity(u, N)
return Fout
def ASM_Ext_fail(Fin, z, factor=5):
wavelength = Fin.lam
N = Fin.N
size = Fin.siz
deltax = size / N
deltay = size / N
F_ext = Interpol(Fin, size, N * factor)
k = 2 * _np.pi / wavelength
a1 = F_ext.field
A1 = _fftshift(_fft2(_fftshift(a1)))
r1 = _np.linspace(-a1.shape[0] / 2, a1.shape[1] / 2 - 1, a1.shape[0])
s1 = _np.linspace(-a1.shape[0] / 2, a1.shape[1] / 2 - 1, a1.shape[0])
deltaFX = 1 / (N * deltax) * r1
deltaFY = 1 / (N * deltay) * s1
meshgrid = _np.meshgrid(deltaFX, deltaFY)
H = _np.exp(1.0j * k * z *
_np.sqrt(1 - _np.power(wavelength * meshgrid[0], 2) -
_np.power(wavelength * meshgrid[1], 2)))
U = _np.multiply(A1, H)
u = _ifftshift(_ifft2(_ifftshift(U)))
Fout = Fin
Fout.field = crop_intensity(u, N)
return Fout
def ASM(Fin, z):
# Fout = Field.shallowcopy(Fin)
Fout = Fin
wavelength = Fout.lam
N = Fout.N
size = Fout.siz
# deltaX = size/N
# deltaY = size/N
k = 2 * _np.pi / wavelength
a1 = Fout.field
A1 = _fftshift(_fft2(_fftshift(a1)))
r1 = _np.linspace(-a1.shape[0] / 2, a1.shape[1] / 2 - 1, a1.shape[0])
s1 = _np.linspace(-a1.shape[0] / 2, a1.shape[1] / 2 - 1, a1.shape[0])
deltaFX = 1 / size * r1
deltaFY = 1 / size * s1
meshgrid = _np.meshgrid(deltaFX, deltaFY)
H = _np.exp(1.0j * k * z *
_np.sqrt(1 - _np.power(wavelength * meshgrid[0], 2) -
_np.power(wavelength * meshgrid[1], 2)))
U = _np.multiply(A1, H)
u = _ifftshift(_ifft2(_ifftshift(U)))
Fout.field = u
return Fout