def _findCarrier(self, band0, band1, mask):
band = band0 * band1
ixf = abs(fft.fftshift(fft.fft2(fft.fftshift(band))))
if self.debug:
plt.figure()
plt.title('Find carrier')
plt.imshow((ixf - gaussian_filter(ixf, 20)) * mask)
pyc0, pxc0 = self._findPeak((ixf - gaussian_filter(ixf, 20)) * mask)
ixfz, Kx, Ky = self._zoomf(band, self.N, self._kx[pyc0, pxc0],
self._ky[pyc0, pxc0], 50, self._dk * self.N)
pyc, pxc = self._findPeak(abs(ixfz))
if self.debug:
plt.figure()
plt.title('Zoon Find carrier')
plt.imshow(abs(ixfz))
kx = Kx[pxc]
ky = Ky[pyc]
otf_exclude_min_radius = 0.5
otf_exclude_max_radius = 1.5
kr = sqrt(self._kx**2 + self._ky**2)
m = (kr < 2)
otf = fft.fftshift(self._tfm(kr, m) + (1 - m))
otf_mask = (kr > otf_exclude_min_radius) & (kr <
otf_exclude_max_radius)
otf_mask_for_band_common_freq = fft.fftshift(
otf_mask & scipy.ndimage.shift(otf_mask, (pyc0 -
(self.N // 2 + 1), pxc0 -
(self.N // 2 + 1)),
order=0))
band0_common = fft.ifft2(
fft.fft2(band0) / otf * otf_mask_for_band_common_freq)
xx = np.arange(-self.N / 2 * self._dx,
self.N / 2 * self._dx,
self._dx,
dtype=np.single)
phase_shift_to_xpeak = exp(-1j * kx * xx * 2 * pi * self.NA /
self.wavelength)
phase_shift_to_ypeak = exp(-1j * ky * xx * 2 * pi * self.NA /
self.wavelength)
band1_common = fft.ifft2(
fft.fft2(band1) / otf * otf_mask_for_band_common_freq) * np.outer(
phase_shift_to_ypeak, phase_shift_to_xpeak)
scaling = 1 / np.sum(band0_common * np.conjugate(band0_common))
cross_corr_result = np.sum(band0_common * band1_common) * scaling
ampl = np.abs(cross_corr_result) * 2
phase = np.angle(cross_corr_result)
return kx, ky, phase, ampl
def fast_cfft2(x, axes=(-1, -2)):
if len(x.shape) == 2:
return fftshift(np.transpose(fft2(np.transpose(ifftshift(x)))))
elif len(x.shape) == 3:
y = ifftshift(x, axes=axes)
y = fft2(y, axes=axes)
y = fftshift(y, axes=axes)
return y
else:
raise ValueError("x must be 2D or 3D")
def _refineCarrier(self, band0, band1, kx_in, ky_in):
pxc0 = np.int(np.round(kx_in / self._dk) + self.N / 2)
pyc0 = np.int(np.round(ky_in / self._dk) + self.N / 2)
otf_exclude_min_radius = self.eta / 2
otf_exclude_max_radius = 1.5
m = (self._kr < 2)
otf = fft.fftshift(self._tfm(self._kr, m) + (1 - m) * 0.0001)
otf_mask = (self._kr > otf_exclude_min_radius) & (
self._kr < otf_exclude_max_radius)
otf_mask_for_band_common_freq = fft.fftshift(
otf_mask & scipy.ndimage.shift(otf_mask, (pyc0 -
(self.N // 2), pxc0 -
(self.N // 2)),
order=0))
band0_common = fft.ifft2(
fft.fft2(band0) / otf * otf_mask_for_band_common_freq)
band1_common = fft.ifft2(
fft.fft2(band1) / otf * otf_mask_for_band_common_freq)
band = band0_common * band1_common
mag = 25 * self.N / 256
ixfz, Kx, Ky = self._zoomf(band, self.N, np.single(self._k[pxc0]),
np.single(self._k[pyc0]), mag,
self._dk * self.N)
pyc, pxc = self._findPeak(abs(ixfz))
if self.debug:
plt.figure()
plt.title('Zoom Find carrier')
plt.imshow(abs(ixfz))
kx = Kx[pxc]
ky = Ky[pyc]
xx = np.arange(-self.N / 2 * self._dx,
self.N / 2 * self._dx,
self._dx,
dtype=np.double)
phase_shift_to_xpeak = exp(-1j * kx * xx * 2 * pi * self.NA /
self.wavelength)
phase_shift_to_ypeak = exp(-1j * ky * xx * 2 * pi * self.NA /
self.wavelength)
scaling = 1 / np.sum(band0_common * np.conjugate(band0_common))
cross_corr_result = np.sum(band0_common * band1_common * np.outer(
phase_shift_to_ypeak, phase_shift_to_xpeak)) * scaling
ampl = np.abs(cross_corr_result) * 2
phase = np.angle(cross_corr_result)
return kx, ky, phase, ampl
def reconstruct_fftw(self, img):
imf = fft.fft2(img) * self._prefilter
self._carray[:, 0:self.N // 2, 0:self.N // 2] = imf[:, 0:self.N // 2, 0:self.N // 2]
self._carray[:, 0:self.N // 2, 3 * self.N // 2:2 * self.N] = imf[:, 0:self.N // 2, self.N // 2:self.N]
self._carray[:, 3 * self.N // 2:2 * self.N, 0:self.N // 2] = imf[:, self.N // 2:self.N, 0:self.N // 2]
self._carray[:, 3 * self.N // 2:2 * self.N, 3 * self.N // 2:2 * self.N] = imf[:, self.N // 2:self.N,
self.N // 2:self.N]
img2 = np.sum(np.real(fft.ifft2(self._carray)).real * self._reconfactor, 0)
self._imgstore = img
self._bigimgstore = fft.ifft2(fft.fft2(img2) * self._postfilter).real
return self._bigimgstore
def solve_spectral(Winit, expU, expT):
B = fft2(Winit, axes=(0,1)) # (x,y,px,py) -> (λx,λy,px,py)
B *= expT
B = ifft2(B, axes=(0,1)) # (λx,λy,px,py) -> (x,y,px,py)
B = fft2(B, axes=(2,3)) # (x,y,px,py) -> (x,y,θx,θy)
B *= expU
B = ifft2(B, axes=(2,3)) # (x,y,θx,θy) -> (x,y,px,py)
B = fft2(B, axes=(0,1)) # (x,y,px,py) -> (λx,λy,px,py)
B *= expT
B = ifft2(B, axes=(0,1)) # (λx,λy,px,py) -> (x,y,px,py)
return real(B) # to avoid python warning
def translation(im0, im1):
"""Return translation vector to register images."""
shape = im0.shape
f0 = fft2(im0)
f1 = fft2(im1)
ir = abs(ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1))))
t0, t1 = np.unravel_index(np.argmax(ir), shape)
if t0 > shape[0] // 2:
t0 -= shape[0]
if t1 > shape[1] // 2:
t1 -= shape[1]
return [t0, t1]
def fast_cfft2(x):
if len(x.shape) == 2:
return fftshift(np.transpose(fft2(np.transpose(ifftshift(x)))))
elif len(x.shape) == 3:
y = ifftshift(x, (1, 2))
y = np.transpose(y, (0, 2, 1))
y = fft2(y)
y = np.transpose(y, (0, 2, 1))
y = fftshift(y, (1, 2))
return y
else:
raise ValueError("x must be 2D or 3D")
def fourier_ring_correlation(obj,
ref,
step_size=1,
save_path='frc',
save_mask=False):
if not os.path.exists(save_path):
os.makedirs(save_path)
radius_max = int(min(obj.shape) / 2)
f_obj = np_fftshift(fft2(obj))
f_ref = np_fftshift(fft2(ref))
f_prod = f_obj * np.conjugate(f_ref)
f_obj_2 = np.real(f_obj * np.conjugate(f_obj))
f_ref_2 = np.real(f_ref * np.conjugate(f_ref))
radius_ls = np.arange(1, radius_max, step_size)
fsc_ls = []
np.save(os.path.join(save_path, 'radii.npy'), radius_ls)
for rad in radius_ls:
print(rad)
if os.path.exists(
os.path.join(save_path,
'mask_rad_{:04d}.tiff'.format(int(rad)))):
mask = dxchange.read_tiff(
os.path.join(save_path,
'mask_rad_{:04d}.tiff'.format(int(rad))))
else:
mask = generate_ring(obj.shape, rad, anti_aliasing=2)
if save_mask:
dxchange.write_tiff(
mask,
os.path.join(save_path,
'mask_rad_{:04d}.tiff'.format(int(rad))),
dtype='float32',
overwrite=True)
fsc = abs(np.sum(f_prod * mask))
fsc /= np.sqrt(np.sum(f_obj_2 * mask) * np.sum(f_ref_2 * mask))
fsc_ls.append(fsc)
np.save(os.path.join(save_path, 'fsc.npy'), fsc_ls)
matplotlib.rcParams['pdf.fonttype'] = 'truetype'
fontProperties = {
'family': 'serif',
'serif': ['Times New Roman'],
'weight': 'normal',
'size': 12
}
plt.rc('font', **fontProperties)
plt.plot(radius_ls.astype(float) / radius_ls[-1], fsc_ls)
plt.xlabel('Spatial frequency (1 / Nyquist)')
plt.ylabel('FRC')
plt.savefig(os.path.join(save_path, 'frc.pdf'), format='pdf')
def translation(im0, im1):
"""Return translation vector to register images."""
shape = im0.shape
f0 = fft2(im0)
f1 = fft2(im1)
ir = abs(ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1))))
t0, t1 = np.unravel_index(np.argmax(ir), shape)
if t0 > shape[0] // 2:
t0 -= shape[0]
if t1 > shape[1] // 2:
t1 -= shape[1]
return [t0, t1]
def reconstructframe_fftw(self, img, i):
diff = img - self._imgstore[i, :, :]
imf = fft.fft2(diff) * self._prefilter
self._carray[0, 0:self.N // 2, 0:self.N // 2] = imf[0:self.N // 2, 0:self.N // 2]
self._carray[0, 0:self.N // 2, 3 * self.N // 2:2 * self.N] = imf[0:self.N // 2, self.N // 2:self.N]
self._carray[0, 3 * self.N // 2:2 * self.N, 0:self.N // 2] = imf[self.N // 2:self.N, 0:self.N // 2]
self._carray[0, 3 * self.N // 2:2 * self.N, 3 * self.N // 2:2 * self.N] = imf[self.N // 2:self.N,
self.N // 2:self.N]
img2 = fft.ifft2(self._carray[0, :, :]).real * self._reconfactor[i, :, :]
self._imgstore[i, :, :] = img
self._bigimgstore = self._bigimgstore + fft.ifft2(fft.fft2(img2) * self._postfilter).real
return self._bigimgstore
def optimal_binary_image_subtraction(R,N,Pr,Pn,sr,sn):
R_hat = fft.fft2(R)
N_hat = fft.fft2(N)
Pn_hat = fft.fft2(Pn)
Pr_hat = fft.fft2(Pr)
G_hat = (Pr_hat*N_hat - Pn_hat*R_hat) / np.sqrt((sr**2*abs(Pn_hat**2) + sn**2*abs(Pr_hat**2)))
P_G_hat = (Pr_hat*Pn_hat) / np.sqrt((sr**2*abs(Pn_hat**2) + sn**2*abs(Pr_hat**2)))
S_hat = G_hat*conj(P_G_hat)
#S_hat = (conj(Pn_hat)*np.abs(Pr_hat)**2*N_hat - conj(Pr_hat)*np.abs(Pn_hat)**2*R_hat) / (sr**2*abs(Pn_hat**2) + sn**2*abs(Pr_hat**2))
S = fft.ifft2(S_hat)
G = fft.ifft2(G_hat)
P_G = real(fft.ifft2(P_G_hat))
return S/std(S[15::30,15::30]), G/std(G[15::30,15::30]), P_G / sum(P_G)
def multislice_propagate_batch_numpy(grid_delta_batch, grid_beta_batch, probe_real, probe_imag, energy_ev, psize_cm, free_prop_cm=None, obj_batch_shape=None):
minibatch_size = obj_batch_shape[0]
grid_shape = obj_batch_shape[1:]
voxel_nm = np.array([psize_cm] * 3) * 1.e7
wavefront = np.zeros([minibatch_size, obj_batch_shape[1], obj_batch_shape[2]], dtype='complex64')
wavefront += (probe_real + 1j * probe_imag)
lmbda_nm = 1240. / energy_ev
mean_voxel_nm = np.prod(voxel_nm) ** (1. / 3)
size_nm = np.array(grid_shape) * voxel_nm
n_slice = obj_batch_shape[-1]
delta_nm = voxel_nm[-1]
# h = get_kernel_ir(delta_nm, lmbda_nm, voxel_nm, grid_shape)
h = get_kernel(delta_nm, lmbda_nm, voxel_nm, grid_shape)
k = 2. * PI * delta_nm / lmbda_nm
probe_array = []
for i in range(n_slice):
delta_slice = grid_delta_batch[:, :, :, i]
beta_slice = grid_beta_batch[:, :, :, i]
c = np.exp(1j * k * delta_slice) * np.exp(-k * beta_slice)
wavefront = wavefront * c
if i < n_slice - 1:
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2]))
probe_array.append(wavefront)
if free_prop_cm is not None:
#1dxchange.write_tiff(abs(wavefront), '2d_1024/monitor_output/wv', dtype='float32', overwrite=True)
if free_prop_cm == 'inf':
wavefront = np_fftshift(fft2(wavefront), axes=[1, 2])
else:
dist_nm = free_prop_cm * 1e7
l = np.prod(size_nm)**(1. / 3)
crit_samp = lmbda_nm * dist_nm / l
algorithm = 'TF' if mean_voxel_nm > crit_samp else 'IR'
# print(algorithm)
algorithm = 'TF'
if algorithm == 'TF':
h = get_kernel(dist_nm, lmbda_nm, voxel_nm, grid_shape)
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2]))
else:
h = get_kernel_ir(dist_nm, lmbda_nm, voxel_nm, grid_shape)
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2]))
# dxchange.write_tiff(abs(wavefront), '2d_512/monitor_output/wv', dtype='float32', overwrite=True)
# dxchange.write_tiff(np.angle(h), '2d_512/monitor_output/h', dtype='float32', overwrite=True)
return wavefront, np.array(probe_array)
def get_kernel_ir(dist_nm, lmbda_nm, voxel_nm, grid_shape):
"""
Get Fresnel propagation kernel for IR algorithm.
Parameters:
-----------
simulator : :class:`acquisition.Simulator`
The Simulator object.
dist : float
Propagation distance in cm.
"""
size_nm = np.array(voxel_nm) * np.array(grid_shape)
k = 2 * PI / lmbda_nm
ymin, xmin = np.array(size_nm)[:2] / -2.
dy, dx = voxel_nm[0:2]
x = np.arange(xmin, xmin + size_nm[1], dx)
y = np.arange(ymin, ymin + size_nm[0], dy)
x, y = np.meshgrid(x, y)
try:
h = np.exp(1j * k * dist_nm) / (1j * lmbda_nm * dist_nm) * np.exp(
1j * k / (2 * dist_nm) * (x**2 + y**2))
H = np_fftshift(fft2(h)) * voxel_nm[0] * voxel_nm[1]
dxchange.write_tiff(x,
'2d_512/monitor_output/x',
dtype='float32',
overwrite=True)
except:
h = tf.exp(1j * k * dist_nm) / (1j * lmbda_nm * dist_nm) * tf.exp(
1j * k / (2 * dist_nm) * (x**2 + y**2))
# h = tf.convert_to_tensor(h, dtype='complex64')
H = fftshift(tf.fft2d(h)) * voxel_nm[0] * voxel_nm[1]
return H
def _rapsd(X):
"""Compute radially averaged PSD of input field X.
"""
if X.shape[0] != X.shape[1]:
raise ValueError("a square array expected, but the shape of X is (%d,%d)" % \
(X.shape[0], X.shape[1]))
L = X.shape[0]
if L % 2 == 1:
XC,YC = np.ogrid[-int(L/2):int(L/2)+1, -int(L/2):int(L/2)+1]
else:
XC,YC = np.ogrid[-int(L/2):int(L/2), -int(L/2):int(L/2)]
R = np.sqrt(XC*XC + YC*YC).astype(int)
F = fft.fftshift(fft.fft2(X, **fft_kwargs))
F = abs(F)**2
if L % 2 == 0:
r_range = np.arange(0, int(L/2)+1)
else:
r_range = np.arange(0, int(L/2))
result = []
for r in r_range:
MASK = R == r
F_vals = F[MASK]
result.append(np.mean(F_vals))
return np.array(result)
def downsample(stack, n, mask=None, stack_in_fourier=False):
""" Use Fourier methods to change the sample interval and/or aspect ratio
of any dimensions of the input image 'img'. If the optional argument
stack is set to True, then the *first* dimension of 'img' is interpreted as the index of
each image in the stack. The size argument side is an integer, the size of the
output images. Let the size of a stack
of 2D images 'img' be n1 x n1 x k. The size of the output will be side x side x k.
If the optional mask argument is given, this is used as the
zero-centered Fourier mask for the re-sampling. The size of mask should
be the same as the output image size. For example for downsampling an
n0 x n0 image with a 0.9 x nyquist filter, do the following:
msk = fuzzymask(n,2,.45*n,.05*n)
out = downsample(img, n, 0, msk)
The size of the mask must be the size of output. The optional fx output
argument is the padded or cropped, masked, FT of in, with zero
frequency at the origin.
"""
size_in = np.square(stack.shape[1])
size_out = np.square(n)
mask = 1 if mask is None else mask
num_images = stack.shape[0]
output = np.zeros((num_images, n, n), dtype='float32')
images_batches = np.array_split(np.arange(num_images), 10)
for batch in images_batches:
curr_batch = np.array(stack[batch])
curr_batch = curr_batch if stack_in_fourier else fft2(curr_batch)
fx = common.crop(np.fft.fftshift(curr_batch, axes=(-2, -1)),
(-1, n, n)) * mask
output[batch] = ifft2(np.fft.ifftshift(
fx, axes=(-2, -1))) * (size_out / size_in)
print('finished {}/{}'.format(batch[-1] + 1, num_images))
return output
def decomposition_fft(X, filter, **kwargs):
"""Decompose a 2d input field into multiple spatial scales by using the Fast
Fourier Transform (FFT) and a bandpass filter.
Parameters
----------
X : array_like
Two-dimensional array containing the input field. All values are required
to be finite.
filter : dict
A filter returned by any method implemented in bandpass_filters.py.
Other Parameters
----------------
MASK : array_like
Optional mask to use for computing the statistics for the cascade levels.
Pixels with MASK==False are excluded from the computations.
Returns
-------
out : ndarray
A dictionary described in the module documentation. The parameter n is
determined from the filter (see bandpass_filters.py).
"""
MASK = kwargs.get("MASK", None)
if len(X.shape) != 2:
raise ValueError("the input is not two-dimensional array")
if MASK is not None and MASK.shape != X.shape:
raise ValueError("dimension mismatch between X and MASK: X.shape=%s, MASK.shape=%s" % \
(str(X.shape), str(MASK.shape)))
if X.shape != filter["weights_2d"].shape[1:3]:
raise ValueError(
"dimension mismatch between X and filter: X.shape=%s, filter['weights_2d'].shape[1:3]=%s"
% (str(X.shape), str(filter["weights_2d"].shape[1:3])))
if np.any(~np.isfinite(X)):
raise ValueError("X contains non-finite values")
result = {}
means = []
stds = []
F = fft.fftshift(fft.fft2(X, **fft_kwargs))
X_decomp = []
for k in range(len(filter["weights_1d"])):
W_k = filter["weights_2d"][k, :, :]
X_ = np.real(fft.ifft2(fft.ifftshift(F * W_k), **fft_kwargs))
X_decomp.append(X_)
if MASK is not None:
X_ = X_[MASK]
means.append(np.mean(X_))
stds.append(np.std(X_))
result["cascade_levels"] = np.stack(X_decomp)
result["means"] = means
result["stds"] = stds
return result
def _rapsd(X):
"""Compute radially averaged PSD of input field X.
"""
if len(X.shape) != 2:
raise ValueError("%i dimensions are found, but the number of dimensions should be 2" % \
len(X.shape))
M, N = X.shape
YC, XC = _compute_centred_coord_array(M, N)
R = np.sqrt(XC * XC + YC * YC).astype(int)
F = fft.fftshift(fft.fft2(X, **fft_kwargs))
F = abs(F)**2
L = max(X.shape[0], X.shape[1])
if L % 2 == 0:
r_range = np.arange(0, int(L / 2) + 1)
else:
r_range = np.arange(0, int(L / 2))
result = []
for r in r_range:
MASK = R == r
F_vals = F[MASK]
result.append(np.mean(F_vals))
return np.array(result)
def _get_ang_scale(ims, bgval, exponent='inf', constraints=None, reports=None):
"""
Given two images, return their scale and angle difference.
Args:
ims (2-tuple-like of 2D ndarrays): The images
bgval: We also pad here in the :func:`map_coordinates`
exponent (float or 'inf'): The exponent stuff, see :func:`similarity`
constraints (dict, optional)
reports (optional)
Returns:
tuple: Scale, angle. Describes the relationship of
the subject image to the first one.
"""
assert len(ims) == 2, \
"Only two images are supported as input"
shape = ims[0].shape
ims_apod = [utils._apodize(im) for im in ims]
dfts = [
fft.fftshift(
fft.fft2(im,
threads=4,
overwrite_input=True,
auto_align_input=True,
auto_contiguous=True,
planner_effort='FFTW_ESTIMATE')) for im in ims_apod
]
filt = _logpolar_filter(shape)
dfts = [dft * filt for dft in dfts]
# High-pass filtering used to be here, but we have moved it to a higher
# level interface
pcorr_shape = _get_pcorr_shape(shape)
log_base = _get_log_base(shape, pcorr_shape[1])
stuffs = [_logpolar(np.abs(dft), pcorr_shape, log_base) for dft in dfts]
(arg_ang, arg_rad), success = _phase_correlation(stuffs[0], stuffs[1],
utils.argmax_angscale,
log_base, exponent,
constraints, reports)
angle = -np.pi * arg_ang / float(pcorr_shape[0])
angle = np.rad2deg(angle)
angle = utils.wrap_angle(angle, 360)
scale = log_base**arg_rad
angle = -angle
scale = 1.0 / scale
if not 0.5 < scale < 2:
raise ValueError(
"Images are not compatible. Scale change %g too big to be true." %
scale)
return scale, angle
def phrt(im, plan, method=4, nr_threads=2):
"""Process a tomographic projection image with the selected phase retrieval algorithm.
Parameters
----------
im : array_like
Flat corrected image data as numpy array.
plan : structure
Structure with pre-computed data (see prepare_plan function)
method : int
Phase retrieval filter {1 = TIE (default), 2 = CTF, 3 = CTF first-half sine, 4 = Quasiparticle,
5 = Quasiparticle first half sine}.
nr_threads : int
Number of threads to be used in the computation of FFT by PyFFTW
Credits
-------
Julian Moosmann, KIT (Germany) is acknowledged for this code
"""
# Extract plan values:
dim0_o = plan['dim0']
dim1_o = plan['dim1']
n_pad0 = plan['npad0']
n_pad1 = plan['npad1']
marg0 = (n_pad0 - dim0_o) / 2
marg1 = (n_pad1 - dim1_o) / 2
filter = plan['filter']
# Pad image (if required):
im = padImage(im, n_pad0, n_pad1)
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = fft2(im - 1, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = fft2(im - 1)
# Apply phase retrieval filter:
im = filter * im
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = real(ifft2(im, threads=nr_threads))
# (Un)comment the following line to use NumPy:
#im = real(ifft2(im))
# Return the negative:
im = -im
# Return cropped output:
return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def phrt(im, plan, method=4, nr_threads=2):
"""Process a tomographic projection image with the selected phase retrieval algorithm.
Parameters
----------
im : array_like
Flat corrected image data as numpy array.
plan : structure
Structure with pre-computed data (see prepare_plan function)
method : int
Phase retrieval filter {1 = TIE (default), 2 = CTF, 3 = CTF first-half sine, 4 = Quasiparticle,
5 = Quasiparticle first half sine}.
nr_threads : int
Number of threads to be used in the computation of FFT by PyFFTW
Credits
-------
Julian Moosmann, KIT (Germany) is acknowledged for this code
"""
# Extract plan values:
dim0_o = plan['dim0']
dim1_o = plan['dim1']
n_pad0 = plan['npad0']
n_pad1 = plan['npad1']
marg0 = (n_pad0 - dim0_o) / 2
marg1 = (n_pad1 - dim1_o) / 2
filter = plan['filter']
# Pad image (if required):
im = padImage(im, n_pad0, n_pad1)
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = fft2(im - 1, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = fft2(im - 1)
# Apply phase retrieval filter:
im = filter * im
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = real(ifft2(im, threads=nr_threads))
# (Un)comment the following line to use NumPy:
#im = real(ifft2(im))
# Return the negative:
im = - im
# Return cropped output:
return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def coad_func(j,base_psfcat, base_sexcat):
# j is the jth coadded image
psf_dat, psf_hed, sexcat2, psfcat2 = get_psf(j)
x_fratio, y_fratio, fratio, dx, dy = get_fratio(base_psfcat, psfcat2, base_sexcat, sexcat2)
R_j = fits.getdata(j)
R_j_hat = fft.fft2(R_j)
_, PSF = chop_kern(R_j, psf_dat, psf_hed,1,1)
P_j = PSF[0]
#print(P_j.shape)
P_j_hat = fft.fft2(P_j)
P_j_hat_bar = np.conj(P_j_hat)
F_j = 1.0/np.median(fratio) #Zeropoint flux
sig_j_2 = np.std(R_j)**2
Nomin = (F_j/sig_j_2) * P_j_hat_bar * R_j_hat
Denom = (F_j**2/sig_j_2) * np.abs(P_j_hat**2)
subprocess.call(['rm', sexcat2, psfcat2, sexcat2.replace('.fits', '.psf')])
return(Nomin, Denom)
def generate_signals(
shape: tuple,
specs: list[CompanionSpec],
template: np.ndarray,
derotation_angles: np.ndarray = None,
template_scale_factors: Union[np.ndarray, float, None] = None,
):
'''Inject signals for companions specified using
optional derotation angles to *counter*rotate the coordinates
before injection, such that rotation CCW (e.g. by `derotate_cube`)
aligns 0 deg PA with +Y
Parameters
----------
shape : tuple[int,int,int]
specs : list[CompanionSpec]
template : np.ndarray
derotation_angles : Optional[np.ndarray]
template_scale_factors : Union[np.ndarray,float,None]
Scale factor relative to 1.0 being the average brightness
of the primary over the observation, used to scale the
template image to reflect particularly sharp or poor
AO correction
Returns
-------
outcube : np.ndarray
'''
outcube = np.zeros(shape, dtype=template.dtype)
n_obs = shape[0]
template = improc.shift2(template, 0, 0, output_shape=shape[1:])
ft_template = fft.fft2(template)
xfreqs = fft.fftfreq(shape[2])
yfreqs = fft.fftfreq(shape[1])
if template_scale_factors is None:
template_scale_factors = np.ones(n_obs)
if np.isscalar(template_scale_factors):
template_scale_factors = np.repeat(np.array([template_scale_factors]),
n_obs)
if derotation_angles is None:
derotation_angles = np.zeros(n_obs)
for spec in specs:
theta = np.deg2rad(90 + spec.pa_deg - derotation_angles)
for i in range(n_obs):
dx = spec.r_px * np.cos(theta[i])
dy = spec.r_px * np.sin(theta[i])
shifter = np.exp(2j * np.pi * ((-dx * xfreqs[np.newaxis, :]) +
(-dy * yfreqs[:, np.newaxis])))
cube_contribution = fft.ifft2(ft_template * shifter).real
cube_contribution *= template_scale_factors[i] * spec.scale
outcube[i] += cube_contribution
return outcube
def _coarseFindCarrier(self, band0, band1, mask):
otf_exclude_min_radius = self.eta / 2 # Min Radius of the circular region around DC that is to be excluded from the cross-correlation calculation
maskhpf = fft.fftshift(self._kr > otf_exclude_min_radius)
band0_common = fft.ifft2(fft.fft2(band0) * maskhpf)
# band1_common = fft.ifft2(fft.fft2(band1)*maskhpf)
ix = band0_common * band1
ixf = np.abs(fft.fftshift(fft.fft2(fft.fftshift(ix))))
if self.debug:
plt.figure()
plt.title('Find carrier')
plt.imshow(ixf, cmap=plt.get_cmap('gray'))
# pyc0, pxc0 = self._findPeak((ixf - gaussian_filter(ixf, 20)) * mask)
pyc0, pxc0 = self._findPeak(ixf * mask)
kx = self._dk * (pxc0 - self.N / 2)
ky = self._dk * (pyc0 - self.N / 2)
return kx, ky
def fresnel_propagate_numpy(wavefront, energy_ev, psize_cm, dist_cm):
lmbda_nm = 1240. / energy_ev
lmbda_cm = 0.000124 / energy_ev
psize_nm = psize_cm * 1e7
dist_nm = dist_cm * 1e7
if dist_cm == 'inf':
wavefront = np_fftshift(fft2(wavefront))
else:
n = np.mean(wavefront.shape)
z_crit_cm = (psize_cm * n)**2 / (lmbda_cm * n)
algorithm = 'TF' if dist_cm < z_crit_cm else 'IR'
if algorithm == 'TF':
h = get_kernel(dist_nm, lmbda_nm, [psize_nm, psize_nm],
wavefront.shape)
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront)) * h))
else:
h = get_kernel_ir(dist_nm, lmbda_nm, [psize_nm, psize_nm],
wavefront.shape)
wavefront = np_ifftshift(ifft2(np_fftshift(fft2(wavefront)) * h))
return wavefront
def im_cov(im): # image isotrope covariance function based on fft
cov = np.real(fft.fftshift(fft.ifft2(np.absolute(fft.fft2(im))**
2))) / im.size
print(cov.shape)
cov1 = cov[int(np.shape(cov)[0] / 2 + 0.5), int(np.shape(cov)[1] / 2):]
cov2 = np.flip(cov[int(np.shape(cov)[0] / 2 +
0.5), :int(np.shape(cov)[1] / 2 + 0.5)])
cov3 = cov[int(np.shape(cov)[1] / 2):, int(np.shape(cov)[0] / 2 + 0.5)]
cov4 = np.flip(cov[:int(np.shape(cov)[1] / 2 + 0.5),
int(np.shape(cov)[0] / 2 + 0.5)])
cov = (cov1 + cov2 + cov3 + cov4) / 4
lags = np.arange(np.shape(cov)[0]) + 0.5
return lags, cov # x and y of the cov function
def lens(wf, focal=660, lens_method='proper', offset_before=0, offset_after=0, **conf):
# propagation before lens
proper.prop_propagate(wf, focal + offset_before)
# Fourier transform of an image using a lens
if lens_method == 'proper':
proper.prop_lens(wf, focal)
elif lens_method == 'numpy':
wf._wfarr = fft.fft2(wf._wfarr)/wf._ngrid
elif lens_method == 'pyfftw':
wf._wfarr = fftw.fft2(wf._wfarr)/wf._ngrid
# propagation after lens
proper.prop_propagate(wf, focal + offset_after)
def usecorrelation( im1, im2 ):
"""Assess the offset (to be used for e.g. the assessment of the center of rotation or the
ovarlap) by computation the peak of the correlation between the two input images.
Parameters
----------
im1 : array_like
Image data as numpy array.
im2 : array_like
Image data as numpy array.
Return value
----------
An integer value of the location of the maximum peak correlation.
"""
# Fourier transform both images:
f_im1 = fft2( im1.astype(float32), threads=2 );
f_im2 = fft2( im2.astype(float32), threads=2 );
# Perform phase correlation (amplitude is normalized):
fc = f_im1 * ( f_im2.conjugate() );
fcn = fc / abs(fc);
# Inverse fourier of peak correlation matrix and max location:
peak_correlation_matrix = real( ifft2( fcn, threads=2 ));
# Calculate actual translation:
max_ix = argmax(peak_correlation_matrix.flatten())
(row, col) = unravel_index(max_ix, peak_correlation_matrix.shape)
if ( col < (peak_correlation_matrix.shape[1]/2) ):
col = - (col - 1);
else:
col = peak_correlation_matrix.shape[1] - (col - 1);
return col / 2;
def worker():
# generate Gaussian white noise field, multiply it with the standard
# deviation of the observed field and apply the precipitation mask
N = randstates[k].randn(R.shape[0], R.shape[1])
N = np.real(fft.ifft2(fft.fft2(N) * R_fft))
N = N / np.std(N) * sigma + mu
N[~MASK] = R_thr_2
# subtract the mean and decompose the masked noise field into a
# cascade
N -= mu
decomp_N = decomp_method(N, F, MASK=MASK_)
return decomp_N["stds"]
def _phase_correlation(im0, im1, callback=None, *args):
"""
Computes phase correlation between im0 and im1
Args:
im0
im1
callback (function): Process the cross-power spectrum (i.e. choose
coordinates of the best element, usually of the highest one).
Defaults to :func:`imreg_dft.utils.argmax2D`
Returns:
tuple: The translation vector (Y, X). Translation vector of (0, 0)
means that the two images match.
"""
if callback is None:
callback = utils._argmax2D
# TODO: Implement some form of high-pass filtering of PHASE correlation
f0, f1 = [
fft.fft2(arr,
threads=4,
overwrite_input=True,
auto_align_input=True,
auto_contiguous=True,
planner_effort='FFTW_ESTIMATE') for arr in (im0, im1)
]
# spectrum can be filtered (already),
# so we have to take precaution against dividing by 0
eps = abs(f1).max() * 1e-15
# cps == cross-power spectrum of im0 and im1
cps = abs(
fft.ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1) + eps),
threads=4,
overwrite_input=True,
auto_align_input=True,
auto_contiguous=True,
planner_effort='FFTW_ESTIMATE'))
# scps = shifted cps
scps = fft.fftshift(cps)
(t0, t1), success = callback(scps, *args)
ret = np.array((t0, t1))
# _compensate_fftshift is not appropriate here, this is OK.
t0 -= f0.shape[0] // 2
t1 -= f0.shape[1] // 2
ret -= np.array(f0.shape, int) // 2
return ret, success
def exp_harm_period(img,
harmonicPeriod,
harmonic_ij='00',
searchRegion=10,
isFFT=False,
verbose=True):
"""
Function to obtain the position (in pixels) in the reciprocal space
of the first harmonic ().
"""
(nRows, nColumns) = img.shape
harV = int(harmonic_ij[0])
harH = int(harmonic_ij[1])
periodVert = harmonicPeriod[0]
periodHor = harmonicPeriod[1]
# adjusts for 1D grating
if periodVert <= 0 or periodVert is None:
periodVert = nRows
if verbose:
wpu.print_blue("MESSAGE: Assuming Horizontal 1D Grating")
if periodHor <= 0 or periodHor is None:
periodHor = nColumns
if verbose:
wpu.print_blue("MESSAGE: Assuming Vertical 1D Grating")
# _check_harmonic_inside_image(harV, harH, nRows, nColumns,
# periodVert, periodHor)
if isFFT:
imgFFT = img
else:
imgFFT = np.fft.fftshift(fft2(img, norm='ortho'))
del_i, del_j = _error_harmonic_peak(imgFFT, harV, harH, periodVert,
periodHor, searchRegion)
if verbose:
wpu.print_blue("MESSAGE: error experimental harmonics " +
"vertical: {:d}".format(del_i))
wpu.print_blue("MESSAGE: error experimental harmonics " +
"horizontal: {:d}".format(del_j))
return periodVert + del_i, periodHor + del_j
def initialize_nonparam_2d_fft_filter(X, **kwargs):
"""Takes a 2d input field and produces a fourier filter by using the Fast
Fourier Transform (FFT).
Parameters
----------
X : array-like
Two-dimensional array containing the input field. All values are required
to be finite.
Other Parameters
----------------
win_type : string
Optional tapering function to be applied to X.
Default : flat-hanning
donorm : bool
Option to normalize the real and imaginary parts.
Default : False
Returns
-------
F : array-like
A two-dimensional array containing the non-parametric filter.
It can be passed to generate_noise_2d_fft_filter().
"""
if len(X.shape) != 2:
raise ValueError("the input is not two-dimensional array")
if np.any(~np.isfinite(X)):
raise ValueError("X contains non-finite values")
# defaults
win_type = kwargs.get('win_type', 'flat-hanning')
donorm = kwargs.get('donorm', False)
X = X.copy()
if win_type is not None:
X -= X.min()
tapering = build_2D_tapering_function(X.shape, win_type)
else:
tapering = np.ones_like(X)
F = fft.fft2(X * tapering, **fft_kwargs)
# normalize the real and imaginary parts
if donorm:
F.imag = (F.imag - np.mean(F.imag)) / np.std(F.imag)
F.real = (F.real - np.mean(F.real)) / np.std(F.real)
return np.abs(F)
def ft_shift2(image: np.ndarray,
dy: float,
dx: float,
flux_tol: Union[None, float] = 1e-15,
output_shape=None):
"""
Fast Fourier subpixel shifting
Parameters
----------
dy : float
Translation in +Y direction (i.e. a feature at (x, y) moves to (x, y + dy))
dx : float
Translation in +X direction (i.e. a feature at (x, y) moves to (x + dx, y))
flux_tol : float
Fractional flux change permissible
``(sum(output) - sum(image)) / sum(image) < flux_tol``
(default: 1e-15)
output_shape : tuple
shape of output array (default: same as input)
"""
if output_shape is None:
output_shape = image.shape
xfreqs = fft.fftfreq(output_shape[1])
yfreqs = fft.fftfreq(output_shape[0])
xform = fft.fft2(image, s=output_shape)
if output_shape is not None:
# compute center-to-center displacement such that
# supplying dx == dy == 0.0 will be a no-op (aside
# from changing shape)
orig_ctr_x, orig_ctr_y = (image.shape[1] - 1) / 2, (image.shape[0] -
1) / 2
new_ctr_x, new_ctr_y = (output_shape[1] - 1) / 2, (output_shape[0] -
1) / 2
base_dx, base_dy = new_ctr_x - orig_ctr_x, new_ctr_y - orig_ctr_y
else:
base_dx = base_dy = 0
modified_xform = xform * np.exp(2j * np.pi * (
(-(dx + base_dx) * xfreqs)[np.newaxis, :] +
(-(dy + base_dy) * yfreqs)[:, np.newaxis]))
new_image = fft.ifft2(modified_xform).real
frac_diff_flux = (np.sum(image) - np.sum(new_image)) / np.sum(image)
if flux_tol is not None and frac_diff_flux > flux_tol:
raise RuntimeError(
f"Flux conservation violated by {frac_diff_flux} fractional difference (more than {flux_tol})"
)
return new_image
def exp_harm_period(img, harmonicPeriod,
harmonic_ij='00', searchRegion=10,
isFFT=False, verbose=True):
"""
Function to obtain the position (in pixels) in the reciprocal space
of the first harmonic ().
"""
(nRows, nColumns) = img.shape
harV = int(harmonic_ij[0])
harH = int(harmonic_ij[1])
periodVert = harmonicPeriod[0]
periodHor = harmonicPeriod[1]
# adjusts for 1D grating
if periodVert <= 0 or periodVert is None:
periodVert = nRows
if verbose:
wpu.print_blue("MESSAGE: Assuming Horizontal 1D Grating")
if periodHor <= 0 or periodHor is None:
periodHor = nColumns
if verbose:
wpu.print_blue("MESSAGE: Assuming Vertical 1D Grating")
# _check_harmonic_inside_image(harV, harH, nRows, nColumns,
# periodVert, periodHor)
if isFFT:
imgFFT = img
else:
imgFFT = np.fft.fftshift(fft2(img, norm='ortho'))
del_i, del_j = _error_harmonic_peak(imgFFT, harV, harH,
periodVert, periodHor,
searchRegion)
if verbose:
wpu.print_blue("MESSAGE: error experimental harmonics " +
"vertical: {:d}".format(del_i))
wpu.print_blue("MESSAGE: error experimental harmonics " +
"horizontal: {:d}".format(del_j))
return periodVert + del_i, periodHor + del_j
def velocity_abs_exp(alpha, timestep, spacing, wavenumber):
"""
Absorption coefficient exponent.
:param alpha: :math:`\\alpha_{\\xi}`
:param timestep: Timestep :math:`\\Delta t`
This value is calculated according to
.. math:: e^{-\\alpha_{\\xi} \\Delta t / 2}
However, since the velocity field is shifted by half a spacing, a correction needs to be applied.
.. math:: \\mathcal{F}^{-1} \\left[ e^{+j k_{\\xi} \\Delta \\xi / 2} \\mathcal{F} \\left( e^{-\\alpha_{\\xi} \Delta t / 2} \\right) \\right]
"""
j = 1j
return ifft2(ne.evaluate("exp(+j * wavenumber*spacing/2.0)") * fft2(ne.evaluate("exp(-alpha * timestep / 2.0)")))
def update(d):
"""
Calculation steps to be taken every step.
:param d: Dictionary containing simulation data.
.. note:: This method should only contain calculation steps.
"""
step = d['step']
# Calculate FFT of pressure
pressure_fft = fft2(d['field'][('pressure', None)]) # Apply atmospheric absorption here?
#out_x = d['executor'].submit(sync_steps, d['temp']['p_x'], d['field'][('velocity', 'x')], pressure_fft, d['k_x'],
#d['kappa'], d['spacing'], d['timestep'], d['density'], d['soundspeed'],
#d['abs_exp']['p']['x'], d['abs_exp']['v']['x'], d['source'][('pressure', None)], d['source'][('velocity', 'x')])
#out_y = d['executor'].submit(sync_steps, d['temp']['p_y'], d['field'][('velocity', 'y')], pressure_fft, d['k_y'],
#d['kappa'], d['spacing'], d['timestep'], d['density'], d['soundspeed'],
#d['abs_exp']['p']['y'], d['abs_exp']['v']['y'], d['source'][('pressure', None)], d['source'][('velocity', 'y')])
#out_x = out_x.result()
#out_y = out_y.result()
out_x = sync_steps(d['temp']['p_x'], d['field'][('velocity', 'x')], pressure_fft, d['k_x'],
d['kappa'], d['spacing'], d['timestep'], d['density'], d['soundspeed'],
d['abs_exp']['p']['x'], d['abs_exp']['v']['x'], d['source'][('pressure', None)], d['source'][('velocity', 'x')])
out_y = sync_steps(d['temp']['p_y'], d['field'][('velocity', 'y')], pressure_fft, d['k_y'],
d['kappa'], d['spacing'], d['timestep'], d['density'], d['soundspeed'],
d['abs_exp']['p']['y'], d['abs_exp']['v']['y'], d['source'][('pressure', None)], d['source'][('velocity', 'y')])
d['temp']['p_x'], d['field'][('velocity', 'x')] = out_x
d['temp']['p_y'], d['field'][('velocity', 'y')] = out_y
d['field'][('pressure', None)] = d['temp']['p_x'] + d['temp']['p_y']
return d
def _phase_correlation(im0, im1, callback=None, *args):
"""
Computes phase correlation between im0 and im1
Args:
im0
im1
callback (function): Process the cross-power spectrum (i.e. choose
coordinates of the best element, usually of the highest one).
Defaults to :func:`imreg_dft.utils.argmax2D`
Returns:
tuple: The translation vector (Y, X). Translation vector of (0, 0)
means that the two images match.
"""
if callback is None:
callback = utils._argmax2D
# TODO: Implement some form of high-pass filtering of PHASE correlation
f0, f1 = [fft.fft2(arr) for arr in (im0, im1)]
# spectrum can be filtered (already),
# so we have to take precaution against dividing by 0
eps = abs(f1).max() * 1e-15
# cps == cross-power spectrum of im0 and im1
cps = abs(fft.ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1) + eps)))
# scps = shifted cps
scps = fft.fftshift(cps)
(t0, t1), success = callback(scps, *args)
ret = np.array((t0, t1))
# _compensate_fftshift is not appropriate here, this is OK.
t0 -= f0.shape[0] // 2
t1 -= f0.shape[1] // 2
ret -= np.array(f0.shape, int) // 2
return ret, success
def plot_harmonic_grid(img, harmonicPeriod=None, isFFT=False):
"""
Takes the FFT of single 2D grating Talbot imaging and plot the grid from
where we extract the harmonic in a image of the
Parameters
----------
img : ndarray – Data (data_exchange format)
Experimental image, whith proper blank image, crop and rotation already
applied.
harmonicPeriod : integer or list of integers
If list, it must be in the format ``[periodVert, periodHor]``. If
integer, then [periodVert = periodHor``.
``periodVert`` and ``periodVert`` are the period of the harmonics in
the reciprocal space in pixels. For the checked board grating,
``periodVert = sqrt(2) * pixel Size / grating Period * number of
rows in the image``
isFFT : Boolean
Flag that tells if the input image ``img`` is in the reciprocal
(``isFFT=True``) or in the real space (``isFFT=False``)
"""
if not isFFT:
imgFFT = np.fft.fftshift(fft2(np.fft.fftshift(img),
norm='ortho'))
else:
imgFFT = img
(nRows, nColumns) = imgFFT.shape
periodVert = harmonicPeriod[0]
periodHor = harmonicPeriod[1]
# adjusts for 1D grating
if periodVert <= 0 or periodVert is None:
periodVert = nRows
if periodHor <= 0 or periodHor is None:
periodHor = nColumns
plt.figure(figsize=(8, 7))
plt.imshow(np.log10(np.abs(imgFFT)), cmap='inferno',
extent=wpu.extent_func(imgFFT))
plt.xlabel('Pixels')
plt.ylabel('Pixels')
harV_min = -(nRows + 1) // 2 // periodVert
harV_max = (nRows + 1) // 2 // periodVert
harH_min = -(nColumns + 1) // 2 // periodHor
harH_max = (nColumns + 1) // 2 // periodHor
for harV in range(harV_min + 1, harV_max + 2):
idxPeak_ij = _idxPeak_ij(harV, 0, nRows, nColumns,
periodVert, periodHor)
plt.axhline(idxPeak_ij[0] - periodVert//2 - nRows//2,
lw=2, color='r')
for harH in range(harH_min + 1, harH_max + 2):
idxPeak_ij = _idxPeak_ij(0, harH, nRows, nColumns,
periodVert, periodHor)
plt.axvline(idxPeak_ij[1] - periodHor // 2 - nColumns//2,
lw=2, color='r')
for harV in range(harV_min, harV_max + 1):
for harH in range(harH_min, harH_max + 1):
idxPeak_ij = _idxPeak_ij(harV, harH,
nRows, nColumns,
periodVert, periodHor)
plt.plot(idxPeak_ij[1] - nColumns//2,
idxPeak_ij[0] - nRows//2,
'ko', mew=2, mfc="None", ms=15)
plt.annotate('{:d}{:d}'.format(-harV, harH),
(idxPeak_ij[1] - nColumns//2,
idxPeak_ij[0] - nRows//2,),
color='red', fontsize=20)
plt.xlim(-nColumns//2, nColumns - nColumns//2)
plt.ylim(-nRows//2, nRows - nRows//2)
plt.title('log scale FFT magnitude, Hamonics Subsets and Indexes',
fontsize=16, weight='bold')
def plot_harmonic_peak(img, harmonicPeriod=None, isFFT=False, fname=None):
"""
Funtion to plot the profile of the harmonic peaks ``10`` and ``01``.
It is ploted 11 profiles of the 11 nearest vertical (horizontal)
lines to the peak ``01`` (``10``)
img : ndarray – Data (data_exchange format)
Experimental image, whith proper blank image, crop and rotation already
applied.
harmonicPeriod : integer or list of integers
If list, it must be in the format ``[periodVert, periodHor]``. If
integer, then [periodVert = periodHor``.
``periodVert`` and ``periodVert`` are the period of the harmonics in
the reciprocal space in pixels. For the checked board grating,
``periodVert = sqrt(2) * pixel Size / grating Period * number of
rows in the image``
isFFT : Boolean
Flag that tells if the input image ``img`` is in the reciprocal
(``isFFT=True``) or in the real space (``isFFT=False``)
"""
if not isFFT:
imgFFT = np.fft.fftshift(fft2(np.fft.fftshift(img),
norm='ortho'))
else:
imgFFT = img
(nRows, nColumns) = img.shape
periodVert = harmonicPeriod[0]
periodHor = harmonicPeriod[1]
# adjusts for 1D grating
if periodVert <= 0 or periodVert is None:
periodVert = nRows
if periodHor <= 0 or periodHor is None:
periodHor = nColumns
fig = plt.figure(figsize=(8, 7))
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)
idxPeak_ij = _idxPeak_ij(0, 1,
nRows, nColumns,
periodVert, periodHor)
for i in range(-5, 5):
ax1.plot(np.abs(imgFFT[idxPeak_ij[0] - 100:idxPeak_ij[0] + 100,
idxPeak_ij[1]-i]),
lw=2, label='01 Vert ' + str(i))
ax1.grid()
idxPeak_ij = _idxPeak_ij(1, 0,
nRows, nColumns,
periodVert, periodHor)
for i in range(-5, 5):
ax2.plot(np.abs(imgFFT[idxPeak_ij[0]-i,
idxPeak_ij[1] - 100:idxPeak_ij[1] + 100]),
lw=2, label='10 Horz ' + str(i))
ax2.grid()
ax1.set_xlabel('Pixels')
ax1.set_ylabel(r'$| FFT |$ ')
ax1.legend(loc=1, fontsize='xx-small')
ax1.title.set_text('Horz')
ax2.set_xlabel('Pixels')
ax2.set_ylabel(r'$| FFT |$ ')
ax2.legend(loc=1, fontsize='xx-small')
ax2.title.set_text('Vert')
plt.show(block=False)
if fname is not None:
plt.savefig(fname, transparent=True)
def single_grating_harmonic_images(img, harmonicPeriod,
searchRegion=10,
plotFlag=False, verbose=False):
"""
Auxiliary function to process the data of single 2D grating Talbot imaging.
It obtain the (real space) harmonic images 00, 01 and 10.
Parameters
----------
img : ndarray – Data (data_exchange format)
Experimental image, whith proper blank image, crop and rotation already
applied.
harmonicPeriod : list of integers in the format [periodVert, periodHor]
``periodVert`` and ``periodVert`` are the period of the harmonics in
the reciprocal space in pixels. For the checked board grating,
periodVert = sqrt(2) * pixel Size / grating Period * number of
rows in the image. For 1D grating, set one of the values to negative or
zero (it will set the period to number of rows or colunms).
searchRegion: int
search for the peak will be in a region of harmonicPeriod/searchRegion
around the theoretical peak position. See also
`:py:func:`wavepy.grating_interferometry.plot_harmonic_grid`
plotFlag: boolean
Flag to plot the image in the reciprocal space and to show the position
of the found peaked and the limits of the harmonic image
verbose: Boolean
verbose flag.
Returns
-------
three 2D ndarray data
Images obtained from the harmonics 00, 01 and 10.
"""
imgFFT = np.fft.fftshift(fft2(img, norm='ortho'))
if plotFlag:
plot_harmonic_grid(imgFFT, harmonicPeriod=harmonicPeriod, isFFT=True)
plt.show(block=False)
imgFFT00 = extract_harmonic(imgFFT,
harmonicPeriod=harmonicPeriod,
harmonic_ij='00',
searchRegion=searchRegion,
isFFT=True,
plotFlag=plotFlag,
verbose=verbose)
imgFFT01 = extract_harmonic(imgFFT,
harmonicPeriod=harmonicPeriod,
harmonic_ij=['0', '1'],
searchRegion=searchRegion,
isFFT=True,
plotFlag=plotFlag,
verbose=verbose)
imgFFT10 = extract_harmonic(imgFFT,
harmonicPeriod=harmonicPeriod,
harmonic_ij=['1', '0'],
searchRegion=searchRegion,
isFFT=True,
plotFlag=plotFlag,
verbose=verbose)
# Plot Fourier image (intensity)
if plotFlag:
# Intensity is Fourier Space
intFFT00 = np.log10(np.abs(imgFFT00))
intFFT01 = np.log10(np.abs(imgFFT01))
intFFT10 = np.log10(np.abs(imgFFT10))
fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(14, 5))
for dat, ax, textTitle in zip([intFFT00, intFFT01, intFFT10],
axes.flat,
['FFT 00', 'FFT 01', 'FFT 10']):
# The vmin and vmax arguments specify the color limits
im = ax.imshow(dat, cmap='inferno', vmin=np.min(intFFT00),
vmax=np.max(intFFT00),
extent=wpu.extent_func(dat))
ax.set_title(textTitle)
if textTitle == 'FFT 00':
ax.set_ylabel('Pixels')
ax.set_xlabel('Pixels')
# Make an axis for the colorbar on the right side
cax = fig.add_axes([0.92, 0.1, 0.03, 0.8])
fig.colorbar(im, cax=cax)
plt.suptitle('FFT subsets - Intensity', fontsize=18, weight='bold')
plt.show(block=False)
img00 = ifft2(np.fft.ifftshift(imgFFT00), norm='ortho')
# non existing harmonics will return NAN, so here we check NAN
if np.all(np.isfinite(imgFFT01)):
img01 = ifft2(np.fft.ifftshift(imgFFT01), norm='ortho')
else:
img01 = imgFFT01
if np.all(np.isfinite(imgFFT10)):
img10 = ifft2(np.fft.ifftshift(imgFFT10), norm='ortho')
else:
img10 = imgFFT10
return (img00, img01, img10)
def visib_1st_harmonics(img, harmonicPeriod, searchRegion=20,
unFilterSize=1, verbose=False):
"""
This function obtain the visibility in a grating imaging experiment by the
ratio of the amplitudes of the first and zero harmonics. See
https://doi.org/10.1364/OE.22.014041 .
Note
----
Note that the absolute visibility also depends on the higher harmonics, and
for a absolute value of visibility all of them must be considered.
Parameters
----------
img : ndarray – Data (data_exchange format)
Experimental image, whith proper blank image, crop and rotation already
applied.
harmonicPeriod : list of integers in the format [periodVert, periodHor]
``periodVert`` and ``periodVert`` are the period of the harmonics in
the reciprocal space in pixels. For the checked board grating,
periodVert = sqrt(2) * pixel Size / grating Period * number of
rows in the image. For 1D grating, set one of the values to negative or
zero (it will set the period to number of rows or colunms).
searchRegion: int
search for the peak will be in a region of harmonicPeriod/searchRegion
around the theoretical peak position. See also
`:py:func:`wavepy.grating_interferometry.plot_harmonic_grid`
verbose: Boolean
verbose flag.
Returns
-------
(float, float)
horizontal and vertical visibilities respectivelly from
harmonics 01 and 10
"""
imgFFT = np.fft.fftshift(fft2(img, norm='ortho'))
_idxPeak_ij_exp00 = _idxPeak_ij_exp(imgFFT, 0, 0,
harmonicPeriod[0], harmonicPeriod[1],
searchRegion)
_idxPeak_ij_exp10 = _idxPeak_ij_exp(imgFFT, 1, 0,
harmonicPeriod[0], harmonicPeriod[1],
searchRegion)
_idxPeak_ij_exp01 = _idxPeak_ij_exp(imgFFT, 0, 1,
harmonicPeriod[0], harmonicPeriod[1],
searchRegion)
from scipy.ndimage.filters import uniform_filter
arg_imgFFT = np.abs(imgFFT)
if unFilterSize > 1:
arg_imgFFT = uniform_filter(arg_imgFFT, unFilterSize)
peak00 = arg_imgFFT[_idxPeak_ij_exp00[0], _idxPeak_ij_exp00[1]]
peak10 = arg_imgFFT[_idxPeak_ij_exp10[0], _idxPeak_ij_exp10[1]]
peak01 = arg_imgFFT[_idxPeak_ij_exp01[0], _idxPeak_ij_exp01[1]]
return (2*peak10/peak00, 2*peak01/peak00, _idxPeak_ij_exp00, _idxPeak_ij_exp10, _idxPeak_ij_exp01)
def compute_bispectrum(self, show_progress=True, use_pyfftw=False,
threads=1, nsamples=100, seed=1000,
mean_subtract=False, **pyfftw_kwargs):
'''
Do the computation.
Parameters
----------
show_progress : optional, bool
Show progress bar while sampling the bispectrum.
use_pyfftw : bool, optional
Enable to use pyfftw, if it is installed.
threads : int, optional
Number of threads to use in FFT when using pyfftw.
nsamples : int, optional
Sets the number of samples to take at each vector
magnitude.
seed : int, optional
Sets the seed for the distribution draws.
mean_subtract : bool, optional
Subtract the mean from the data before computing. This removes the
"zero frequency" (i.e., constant) portion of the power, resulting
in a loss of phase coherence along the k_1=k_2 line.
pyfft_kwargs : Passed to
`~turbustat.statistics.rfft_to_fft.rfft_to_fft`. See
`here <https://hgomersall.github.io/pyFFTW/pyfftw/interfaces/interfaces.html#interfaces-additional-args>`_
for a list of accepted kwargs.
'''
if mean_subtract:
norm_data = self.data - self.data.mean()
else:
norm_data = self.data
if use_pyfftw:
if PYFFTW_FLAG:
if pyfftw_kwargs.get('threads') is not None:
pyfftw_kwargs.pop('threads')
fftarr = fft2(norm_data,
threads=threads,
**pyfftw_kwargs)
else:
warn("pyfftw not installed. Reverting to using numpy.")
use_pyfftw = False
if not use_pyfftw:
fftarr = np.fft.fft2(norm_data)
conjfft = np.conj(fftarr)
bispec_shape = (int(self.shape[0] / 2.), int(self.shape[1] / 2.))
self._bispectrum = np.zeros(bispec_shape, dtype=np.complex)
self._bicoherence = np.zeros(bispec_shape, dtype=np.float)
self._tracker = np.zeros(self.shape, dtype=np.int16)
biconorm = np.ones_like(self.bispectrum, dtype=float)
if show_progress:
bar = ProgressBar(np.prod(fftarr.shape) / 4.)
prod = product(range(int(fftarr.shape[0] / 2.)),
range(int(fftarr.shape[1] / 2.)))
with NumpyRNGContext(seed):
for n, (k1mag, k2mag) in enumerate(prod):
phi1 = ra.uniform(0, 2 * np.pi, nsamples)
phi2 = ra.uniform(0, 2 * np.pi, nsamples)
k1x = np.asarray([int(k1mag * np.cos(angle))
for angle in phi1])
k2x = np.asarray([int(k2mag * np.cos(angle))
for angle in phi2])
k1y = np.asarray([int(k1mag * np.sin(angle))
for angle in phi1])
k2y = np.asarray([int(k2mag * np.sin(angle))
for angle in phi2])
k3x = np.asarray([int(k1mag * np.cos(ang1) +
k2mag * np.cos(ang2))
for ang1, ang2 in zip(phi1, phi2)])
k3y = np.asarray([int(k1mag * np.sin(ang1) +
k2mag * np.sin(ang2))
for ang1, ang2 in zip(phi1, phi2)])
samps = fftarr[k1x, k1y] * fftarr[k2x, k2y] * conjfft[k3x, k3y]
self._bispectrum[k1mag, k2mag] = np.sum(samps)
biconorm[k1mag, k2mag] = np.sum(np.abs(samps))
# Track where we're sampling from in fourier space
self._tracker[k1x, k1y] += 1
self._tracker[k2x, k2y] += 1
self._tracker[k3x, k3y] += 1
if show_progress:
bar.update(n + 1)
self._bicoherence = (np.abs(self.bispectrum) / biconorm)
self._bispectrum_amp = np.log10(np.abs(self.bispectrum))
def phase_retrieval(im, plan, method=1, nr_threads=2):
"""Process a tomographic projection image with the selected phase retrieval algorithm.
Parameters
----------
im : array_like
Flat corrected image data as numpy array.
plan : structure
Structure with pre-computed data (see prepare_plan function)
method : int
Phase retrieval algorithm {1 = TIE (default), 2 = Born, 3 = Rytov, 4 = Wu}
nr_threads : int
Number of threads to be used in the computation of FFT by PyFFTW
"""
# Extract plan values:
dim0_o = plan['dim0']
dim1_o = plan['dim1']
n_pad0 = plan['npad0']
n_pad1 = plan['npad1']
marg0 = (n_pad0 - dim0_o) / 2
marg1 = (n_pad1 - dim1_o) / 2
den = plan['den']
mu = plan['mu']
# Pad image (if required):
im = padImage(im, n_pad0, n_pad1)
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = fft2(im, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = fft2(im)
# Apply formula:
if method == 1:
im = im / den
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = ifft2(im, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = ifft2(im)
im = im.astype(complex64)
im = real(im)
im = im.astype(float32)
im = -1 / mu * nplog(im)
#
# WARNING: The following methods are not tested
#
elif method == 2:
im = real(ifft2((im - 1.0) / 2.0) / den)
elif method == 3:
im = real(ifft2(nplog(im) / 2.0) / den)
elif method == 4:
im = real(ifft2(im / den))
im = -1 / 2 * (delta / beta) * nplog(im)
# Return cropped output:
return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def similarity(im0, im1):
"""Return similarity transformed image im1 and transformation parameters.
Transformation parameters are: isotropic scale factor, rotation angle (in
degrees), and translation vector.
A similarity transformation is an affine transformation with isotropic
scale and without shear.
Limitations:
Image shapes must be equal and square.
All image areas must have same scale, rotation, and shift.
Scale change must be less than 1.8.
No subpixel precision.
"""
if im0.shape != im1.shape:
raise ValueError("Images must have same shapes.")
elif len(im0.shape) != 2:
raise ValueError("Images must be 2 dimensional.")
f0 = fftshift(abs(fft2(im0)))
f1 = fftshift(abs(fft2(im1)))
h = highpass(f0.shape)
f0 *= h
f1 *= h
del h
f0, log_base = logpolar(f0)
f1, log_base = logpolar(f1)
f0 = fft2(f0)
f1 = fft2(f1)
r0 = abs(f0) * abs(f1)
ir = abs(ifft2((f0 * f1.conjugate()) / r0))
i0, i1 = np.unravel_index(np.argmax(ir), ir.shape)
angle = 180.0 * i0 / ir.shape[0]
scale = log_base ** i1
if scale > 1.8:
ir = abs(ifft2((f1 * f0.conjugate()) / r0))
i0, i1 = np.unravel_index(np.argmax(ir), ir.shape)
angle = -180.0 * i0 / ir.shape[0]
scale = 1.0 / (log_base ** i1)
if scale > 1.8:
raise ValueError("Images are not compatible. Scale change > 1.8")
if angle < -90.0:
angle += 180.0
elif angle > 90.0:
angle -= 180.0
im2 = ndii.zoom(im1, 1.0 / scale)
im2 = ndii.rotate(im2, angle)
if im2.shape < im0.shape:
t = np.zeros_like(im0)
t[:im2.shape[0], :im2.shape[1]] = im2
im2 = t
elif im2.shape > im0.shape:
im2 = im2[:im0.shape[0], :im0.shape[1]]
f0 = fft2(im0)
f1 = fft2(im2)
ir = abs(ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1))))
t0, t1 = np.unravel_index(np.argmax(ir), ir.shape)
if t0 > f0.shape[0] // 2:
t0 -= f0.shape[0]
if t1 > f0.shape[1] // 2:
t1 -= f0.shape[1]
im2 = ndii.shift(im2, [t0, t1])
# correct parameters for ndimage's internal processing
if angle > 0.0:
d = int((int(im1.shape[1] / scale) * math.sin(math.radians(angle))))
t0, t1 = t1, d + t0
elif angle < 0.0:
d = int((int(im1.shape[0] / scale) * math.sin(math.radians(angle))))
t0, t1 = d + t1, d + t0
scale = (im1.shape[1] - 1) / (int(im1.shape[1] / scale) - 1)
return im2, scale, angle, [-t0, -t1]
from __future__ import division
import warnings
import numpy as np
from functools import wraps
from menpo.feature.base import rebuild_feature_image
try:
# try importing pyfftw
from pyfftw.interfaces.numpy_fft import fft2, ifft2, fftshift, ifftshift
try:
# try calling fft2 on a 4-dimensional array (this is known to have
# problem in some linux distributions)
fft2(np.zeros((1, 1, 1, 1)))
except RuntimeError:
warnings.warn("pyfftw is known to be buggy on your system, numpy.fft "
"will be used instead. Consequently, all algorithms "
"using ffts will be running at a slower speed.",
RuntimeWarning)
from numpy.fft import fft2, ifft2, fftshift, ifftshift
except ImportError:
warnings.warn("pyfftw is not installed on your system, numpy.fft will be "
"used instead. Consequently, all algorithms using ffts "
"will be running at a slower speed. Consider installing "
"pyfftw (pip install pyfftw) to speed up your ffts.",
ImportWarning)
from numpy.fft import fft2, ifft2, fftshift, ifftshift
# TODO: Document me!
def _get_ang_scale(ims, bgval, exponent='inf', constraints=None, reports=None):
"""
Given two images, return their scale and angle difference.
Args:
ims (2-tuple-like of 2D ndarrays): The images
bgval: We also pad here in the :func:`map_coordinates`
exponent (float or 'inf'): The exponent stuff, see :func:`similarity`
constraints (dict, optional)
reports (optional)
Returns:
tuple: Scale, angle. Describes the relationship of
the subject image to the first one.
"""
assert len(ims) == 2, \
"Only two images are supported as input"
shape = ims[0].shape
ims_apod = [utils._apodize(im) for im in ims]
dfts = [fft.fftshift(fft.fft2(im)) for im in ims_apod]
filt = _logpolar_filter(shape)
dfts = [dft * filt for dft in dfts]
# High-pass filtering used to be here, but we have moved it to a higher
# level interface
pcorr_shape = _get_pcorr_shape(shape)
log_base = _get_log_base(shape, pcorr_shape[1])
stuffs = [_logpolar(np.abs(dft), pcorr_shape, log_base)
for dft in dfts]
(arg_ang, arg_rad), success = _phase_correlation(
stuffs[0], stuffs[1],
utils.argmax_angscale, log_base, exponent, constraints, reports)
angle = -np.pi * arg_ang / float(pcorr_shape[0])
angle = np.rad2deg(angle)
angle = utils.wrap_angle(angle, 360)
scale = log_base ** arg_rad
angle = - angle
scale = 1.0 / scale
if reports is not None:
reports["shape"] = filt.shape
reports["base"] = log_base
if reports.show("spectra"):
reports["dfts_filt"] = dfts
if reports.show("inputs"):
reports["ims_filt"] = [fft.ifft2(np.fft.ifftshift(dft))
for dft in dfts]
if reports.show("logpolar"):
reports["logpolars"] = stuffs
if reports.show("scale_angle"):
reports["amas-result-raw"] = (arg_ang, arg_rad)
reports["amas-result"] = (scale, angle)
reports["amas-success"] = success
extent_el = pcorr_shape[1] / 2.0
reports["amas-extent"] = (
log_base ** (-extent_el), log_base ** extent_el,
-90, 90
)
if not 0.5 < scale < 2:
raise ValueError(
"Images are not compatible. Scale change %g too big to be true."
% scale)
return scale, angle
def fft_convolve2d(x, f, mode='same', boundary='constant', fft_filter=False):
r"""
Performs fast 2d convolution in the frequency domain convolving each image
channel with its corresponding filter channel.
Parameters
----------
x : ``(channels, height, width)`` `ndarray`
Image.
f : ``(channels, height, width)`` `ndarray`
Filter.
mode : str {`full`, `same`, `valid`}, optional
Determines the shape of the resulting convolution.
boundary: str {`constant`, `symmetric`}, optional
Determines how the image is padded.
fft_filter: `bool`, optional
If `True`, the filter is assumed to be defined on the frequency
domain. If `False` the filter is assumed to be defined on the
spatial domain.
Returns
-------
c: ``(channels, height, width)`` `ndarray`
Result of convolving each image channel with its corresponding
filter channel.
"""
if fft_filter:
# extended shape is filter shape
ext_shape = np.asarray(f.shape[-2:])
# extend image and filter
ext_x = pad(x, ext_shape, boundary=boundary)
# compute ffts of extended image
fft_ext_x = fft2(ext_x)
fft_ext_f = f
else:
# extended shape
x_shape = np.asarray(x.shape[-2:])
f_shape = np.asarray(f.shape[-2:])
f_half_shape = (f_shape / 2).astype(int)
ext_shape = x_shape + f_half_shape - 1
# extend image and filter
ext_x = pad(x, ext_shape, boundary=boundary)
ext_f = pad(f, ext_shape)
# compute ffts of extended image and extended filter
fft_ext_x = fft2(ext_x)
fft_ext_f = fft2(ext_f)
# compute extended convolution in Fourier domain
fft_ext_c = fft_ext_f * fft_ext_x
# compute ifft of extended convolution
ext_c = np.real(ifftshift(ifft2(fft_ext_c), axes=(-2, -1)))
if mode is 'full':
return ext_c
elif mode is 'same':
return crop(ext_c, x_shape)
elif mode is 'valid':
return crop(ext_c, x_shape - f_half_shape + 1)
else:
raise ValueError(
"mode={}, is not supported. The only supported "
"modes are: 'full', 'same' and 'valid'.".format(mode))
def extract_harmonic(img, harmonicPeriod,
harmonic_ij='00', searchRegion=10, isFFT=False,
plotFlag=False, verbose=True):
"""
Function to extract one harmonic image of the FFT of single grating
Talbot imaging.
The function use the provided value of period to search for the harmonics
peak. The search is done in a rectangle of size
``periodVert*periodHor/searchRegion**2``. The final result is a rectagle of
size ``periodVert x periodHor`` centered at
``(harmonic_Vertical*periodVert x harmonic_Horizontal*periodHor)``
Parameters
----------
img : ndarray – Data (data_exchange format)
Experimental image, whith proper blank image, crop and rotation already
applied.
harmonicPeriod : list of integers in the format [periodVert, periodHor]
``periodVert`` and ``periodVert`` are the period of the harmonics in
the reciprocal space in pixels. For the checked board grating,
periodVert = sqrt(2) * pixel Size / grating Period * number of
rows in the image. For 1D grating, set one of the values to negative or
zero (it will set the period to number of rows or colunms).
harmonic_ij : string or list of string
string with the harmonic to extract, for instance '00', '01', '10'
or '11'. In this notation negative harmonics are not allowed.
Alternativelly, it accepts a list of string
``harmonic_ij=[harmonic_Vertical, harmonic_Horizontal]``, for instance
``harmonic_ij=['0', '-1']``
Note that since the original image contain only real numbers (not
complex), then negative and positive harmonics are symetric
related to zero.
isFFT : Boolean
Flag that tells if the input image ``img`` is in the reciprocal
(``isFFT=True``) or in the real space (``isFFT=False``)
searchRegion: int
search for the peak will be in a region of harmonicPeriod/searchRegion
around the theoretical peak position
plotFlag: Boolean
Flag to plot the image in the reciprocal space and to show the position
of the found peaked and the limits of the harmonic image
verbose: Boolean
verbose flag.
Returns
-------
2D ndarray
Copped Images of the harmonics ij
This functions crops a rectagle of size ``periodVert x periodHor`` centered
at ``(harmonic_Vertical*periodVert x harmonic_Horizontal*periodHor)`` from
the provided FFT image.
Note
----
* Note that it is the FFT of the image that is required.
* The search for the peak is only used to print warning messages.
**Q: Why not the real image??**
**A:** Because FFT can be time consuming. If we use the real image, it will
be necessary to run FFT for each harmonic. It is encourage to wrap this
function within a function that do the FFT, extract the harmonics, and
return the real space harmonic image.
See Also
--------
:py:func:`wavepy.grating_interferometry.plot_harmonic_grid`
"""
(nRows, nColumns) = img.shape
harV = int(harmonic_ij[0])
harH = int(harmonic_ij[1])
periodVert = harmonicPeriod[0]
periodHor = harmonicPeriod[1]
if verbose:
wpu.print_blue("MESSAGE: Extracting harmonic " +
harmonic_ij[0] + harmonic_ij[1])
wpu.print_blue("MESSAGE: Harmonic period " +
"Horizontal: {:d} pixels".format(periodHor))
wpu.print_blue("MESSAGE: Harmonic period " +
"Vertical: {:d} pixels".format(periodVert))
# adjusts for 1D grating
if periodVert <= 0 or periodVert is None:
periodVert = nRows
if verbose:
wpu.print_blue("MESSAGE: Assuming Horizontal 1D Grating")
if periodHor <= 0 or periodHor is None:
periodHor = nColumns
if verbose:
wpu.print_blue("MESSAGE: Assuming Vertical 1D Grating")
try:
_check_harmonic_inside_image(harV, harH, nRows, nColumns,
periodVert, periodHor)
except ValueError:
raise SystemExit
if isFFT:
imgFFT = img
else:
imgFFT = np.fft.fftshift(fft2(img, norm='ortho'))
intensity = (np.abs(imgFFT))
# Estimate harmonic positions
idxPeak_ij = _idxPeak_ij(harV, harH, nRows, nColumns,
periodVert, periodHor)
del_i, del_j = _error_harmonic_peak(imgFFT, harV, harH,
periodVert, periodHor,
searchRegion)
if verbose:
print("MESSAGE: extract_harmonic:" +
" harmonic peak " + harmonic_ij[0] + harmonic_ij[1] +
" is misplaced by:")
print("MESSAGE: {:d} pixels in vertical, {:d} pixels in hor".format(
del_i, del_j))
print("MESSAGE: Theoretical peak index: {:d},{:d} [VxH]".format(
idxPeak_ij[0], idxPeak_ij[1]))
if ((np.abs(del_i) > searchRegion // 2) or
(np.abs(del_j) > searchRegion // 2)):
wpu.print_red("ATTENTION: Harmonic Peak " + harmonic_ij[0] +
harmonic_ij[1] + " is too far from theoretical value.")
wpu.print_red("ATTENTION: {:d} pixels in vertical,".format(del_i) +
"{:d} pixels in hor".format(del_j))
if plotFlag:
from matplotlib.patches import Rectangle
plt.figure(figsize=(8, 7))
plt.imshow(np.log10(intensity), cmap='inferno', extent=wpu.extent_func(intensity))
plt.xlabel('Pixels')
plt.ylabel('Pixels')
xo = idxPeak_ij[1] - nColumns//2 - periodHor//2
yo = nRows//2 - idxPeak_ij[0] - periodVert//2
# xo yo are the lower left position of the reangle
plt.gca().add_patch(Rectangle((xo, yo),
periodHor, periodVert,
lw=2, ls='--', color='red',
fill=None, alpha=1))
plt.title('Selected Region ' + harmonic_ij[0] + harmonic_ij[1],
fontsize=18, weight='bold')
plt.show(block=False)
return imgFFT[idxPeak_ij[0] - periodVert//2:
idxPeak_ij[0] + periodVert//2,
idxPeak_ij[1] - periodHor//2:
idxPeak_ij[1] + periodHor//2]
def sync_steps(p, v, p_fft, k, kappa, spacing, timestep, density, soundspeed, abs_exp_p, abs_exp_v, source_p, source_v):
v = velocity_with_pml(v, ifft2(to_pressure_gradient(p_fft, k, kappa, spacing)), timestep, density, abs_exp_v, source_v)
p = pressure_with_pml(p, ifft2(to_velocity_gradient(fft2(v), k, kappa, spacing)), timestep, density, soundspeed, abs_exp_p, source_p)
return p, v
def fftFromLiteMap(liteMap, applySlepianTaper=False, nresForSlepian=3.0, threads=1):
"""
Create an fft2D object from a liteMap.
Parameters
----------
liteMap : liteMap.liteMap
The map object whose fft is being taken.
applySlepianTaper : bool, optional
If ``True``, apply the lowest order taper (to minimize edge-leakage).
Default is ``False``.
nresForSlepian : float, optional
If ``applySlepianTaper`` = ``True``, this specifies the resolution of
the taper to use. Default is 3.0.
threads : int, optional
Number of threads to use in pyFFTW calculations. Default is 1.
Returns
-------
ft : fftTools.fft2D
The fft2D object corresponding the input liteMap.
"""
ft = fft2D()
ft.Nx = liteMap.Nx
ft.Ny = liteMap.Ny
trace.issue("flipper.fftTools", 1, "Taking FFT of map with (Ny, Nx)= (%f, %f)" %(ft.Ny,ft.Nx))
ft.pixScaleX = liteMap.pixScaleX
ft.pixScaleY = liteMap.pixScaleY
lx = 2*numpy.pi*fftfreq(ft.Nx, d=ft.pixScaleX)
ly = 2*numpy.pi*fftfreq(ft.Ny, d=ft.pixScaleY)
ix = numpy.mod(numpy.arange(ft.Nx*ft.Ny), ft.Nx)
iy = numpy.arange(ft.Nx*ft.Ny)/ft.Nx
modLMap = numpy.zeros([ft.Ny, ft.Nx])
modLMap[iy,ix] = numpy.sqrt(lx[ix]**2 + ly[iy]**2)
ft.modLMap = modLMap
ft.lx = lx
ft.ly = ly
ft.ix = ix
ft.iy = iy
ft.thetaMap = numpy.zeros([ft.Ny, ft.Nx])
ft.thetaMap[iy[:], ix[:]] = numpy.arctan2(ly[iy[:]], lx[ix[:]])
ft.thetaMap *= 180./numpy.pi
mp = liteMap.data.copy()
taper = mp.copy()*0. + 1.0
if (applySlepianTaper):
try:
f = open(taperDir + os.path.sep + 'taper_Ny%d_Nx%d_Nres%3.1f' %(ft.Ny, ft.Nx, nresForSlepian))
taper = pickle.load(f)
f.close()
except:
taper = slepianTaper00(ft.Nx, ft.Ny,nresForSlepian)
f = open(taperDir + os.path.sep + 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny, ft.Nx, nresForSlepian), mode="w")
pickle.dump(taper,f)
f.close()
if have_pyFFTW:
ft.kMap = fft2(mp*taper, threads=threads)
else:
ft.kMap = fft2(mp*taper)
del mp, modLMap, lx, ly
return ft
def upgradePixelPitch(m, N=1, threads=1):
"""
Go to finer pixels with fourier interpolation.
Parameters
----------
m : liteMap
The liteMap object holding the data to upgrade the pixel size of.
N : int, optional
Go to 2^N times smaller pixels. Default is 1.
threads : int, optional
Number of threads to use in pyFFTW calculations. Default is 1.
Returns
-------
mNew : liteMap
The map with smaller pixels.
"""
if N < 1:
return m.copy()
Ny = m.Ny * 2 ** N
Nx = m.Nx * 2 ** N
npix = Ny * Nx
if have_pyFFTW:
ft = fft2(m.data, threads=threads)
else:
ft = fft2(m.data)
ftShifted = fftshift(ft)
newFtShifted = numpy.zeros((Ny, Nx), dtype=numpy.complex128)
# From the numpy.fft.fftshift help:
# """
# Shift zero-frequency component to center of spectrum.
#
# This function swaps half-spaces for all axes listed (defaults to all).
# If len(x) is even then the Nyquist component is y[0].
# """
#
# So in the case that we have an odd dimension in our map, we want to put
# the extra zero at the beginning
if m.Nx % 2 != 0:
offsetX = (Nx - m.Nx) / 2 + 1
else:
offsetX = (Nx - m.Nx) / 2
if m.Ny % 2 != 0:
offsetY = (Ny - m.Ny) / 2 + 1
else:
offsetY = (Ny - m.Ny) / 2
newFtShifted[offsetY : offsetY + m.Ny, offsetX : offsetX + m.Nx] = ftShifted
del ftShifted
ftNew = ifftshift(newFtShifted)
del newFtShifted
# Finally, deconvolve by the pixel window
mPix = numpy.copy(numpy.real(ftNew))
mPix[:] = 0.0
mPix[
mPix.shape[0] / 2 - (2 ** (N - 1)) : mPix.shape[0] / 2 + (2 ** (N - 1)),
mPix.shape[1] / 2 - (2 ** (N - 1)) : mPix.shape[1] / 2 + (2 ** (N - 1)),
] = (1.0 / (2.0 ** N) ** 2)
if have_pyFFTW:
ftPix = fft2(mPix, threads=threads)
else:
ftPix = fft2(mPix)
del mPix
inds = numpy.where(ftNew != 0)
ftNew[inds] /= numpy.abs(ftPix[inds])
if have_pyFFTW:
newData = ifft2(ftNew, threads=threads) * (2 ** N) ** 2
else:
newData = ifft2(ftNew) * (2 ** N) ** 2
del ftNew
del ftPix
x0_new, y0_new = m.pixToSky(0, 0)
m = m.copy() # don't overwrite original
m.wcs.header.update("NAXIS1", 2 ** N * m.wcs.header["NAXIS1"])
m.wcs.header.update("NAXIS2", 2 ** N * m.wcs.header["NAXIS2"])
m.wcs.header.update("CDELT1", m.wcs.header["CDELT1"] / 2.0 ** N)
m.wcs.header.update("CDELT2", m.wcs.header["CDELT2"] / 2.0 ** N)
m.wcs.updateFromHeader()
p_x, p_y = m.skyToPix(x0_new, y0_new)
m.wcs.header.update("CRPIX1", m.wcs.header["CRPIX1"] - p_x)
m.wcs.header.update("CRPIX2", m.wcs.header["CRPIX2"] - p_y)
m.wcs.updateFromHeader()
mNew = liteMapFromDataAndWCS(numpy.real(newData), m.wcs)
mNew.data[:] = numpy.real(newData[:])
return mNew
