def inverse_psf_rfft(psf, shape=None, l=20, mode='laplacian'):
"""
Computes the real FFT of a regularized inversed 2D PSF (or projected 3D)
This follows the convention of fft.rfft: only half the spectrum is computed.
Parameters
----------
psf : array [ZXY] or [XY]
The 2D PSF (if 3D, will be projected on Z axis).
shape : tuple (int, int), optional
Shape of the full-sized desired PSF
(if None, will be the same as the PSF), by default None.
l : int, optional
Regularization lambda, by default 20
mode : str, optional
The regularizer used, by default laplacian.
One of: ['laplacian', 'constant']
Returns
-------
array [XY]
The real FFT of the inverse PSF.
Raises
------
ValueError
If the PSF has incorrect number of dimensions.
If the regularizer is unknown.
"""
if psf.ndim == 3:
psf = psf.sum(0)
elif psf.ndim != 2:
raise ValueError("Invalid dimensions for PSF: {}".format(psf.ndim))
if shape is None:
shape = psf.shape
psf_fft = fft.rfft2(psf, s=shape)
# We need to shift the PSF so that the center is located at the (0, 0) pixel
# otherwise deconvolving will shift every pixel
freq = fft.rfftfreq(shape[1])
phase_shift = freq * 2 * np.pi * ((psf.shape[1] - 1) // 2)
psf_fft *= np.exp(1j * phase_shift[None, :])
freq = fft.fftfreq(shape[0])
phase_shift = freq * 2 * np.pi * ((psf.shape[0] - 1) // 2)
psf_fft *= np.exp(1j * phase_shift[:, None])
if mode == 'laplacian':
# Laplacian regularization filter to avoid NaN
filt = [[0, -0.25, 0], [-0.25, 1, -0.25], [0, -0.25, 0]]
reg = np.abs(fft.rfft2(filt, s=shape))**2
elif mode == 'constant':
reg = 1
else:
raise ValueError('Unknown regularizer: {}'.format(mode))
return psf_fft.conjugate() / (np.abs(psf_fft)**2 + l * reg)
def correlation_func(cor_win_1,
cor_win_2,
window_size,
correlation_method='circular'):
'''This function is doing the cross-correlation. Right now circular cross-correlation
That means no zero-padding is done
the .real is to cut off possible imaginary parts that remains due to finite numerical accuarcy
'''
if correlation_method == 'linear':
cor_win_1 = cor_win_1 - cor_win_1.mean(axis=(1, 2)).reshape(
cor_win_1.shape[0], 1, 1)
cor_win_2 = cor_win_2 - cor_win_2.mean(axis=(1, 2)).reshape(
cor_win_1.shape[0], 1, 1)
cor_win_1[cor_win_1 < 0] = 0
cor_win_2[cor_win_2 < 0] = 0
corr = fftshift(irfft2(
np.conj(rfft2(cor_win_1, s=(2 * window_size, 2 * window_size))) *
rfft2(cor_win_2, s=(2 * window_size, 2 * window_size))).real,
axes=(1, 2))
corr = corr[:, window_size // 2:3 * window_size // 2,
window_size // 2:3 * window_size // 2]
else:
corr = fftshift(irfft2(np.conj(rfft2(cor_win_1)) *
rfft2(cor_win_2)).real,
axes=(1, 2))
return corr
def fft_correlate_strided_images(image_a, image_b):
"""FFT based cross correlation
of two images with multiple views of np.stride_tricks()
The 2D FFT should be applied to the last two axes (-2,-1) and the
zero axis is the number of the interrogation window
This should also work out of the box for rectangular windows.
Parameters
----------
image_a : 3d np.ndarray, first dimension is the number of windows,
and two last dimensions are interrogation windows of the first image
image_b : similar
"""
s1 = np.array(image_a.shape[-2:])
s2 = np.array(image_b.shape[-2:])
size = s1 + s2 - 1
fsize = 2**np.ceil(np.log2(size)).astype(int)
fslice = tuple([slice(0, image_a.shape[0])] +
[slice(0, int(sz)) for sz in size])
f2a = rfft2(image_a, fsize, axes=(-2, -1))
f2b = rfft2(image_b[:, ::-1, ::-1], fsize, axes=(-2, -1))
corr = irfft2(f2a * f2b, axes=(-2, -1)).real[fslice]
return corr
def test_rfft2(self):
x = random((30, 20))
expect = fft.fft2(x)[:, :11]
assert_array_almost_equal(expect, fft.rfft2(x))
assert_array_almost_equal(expect, fft.rfft2(x, norm="backward"))
assert_array_almost_equal(expect / np.sqrt(30 * 20),
fft.rfft2(x, norm="ortho"))
assert_array_almost_equal(expect / (30 * 20),
fft.rfft2(x, norm="forward"))
def fft_correlate_images(image_a,
image_b,
correlation_method="circular",
normalized_correlation=True):
""" FFT based cross correlation
of two images with multiple views of np.stride_tricks()
The 2D FFT should be applied to the last two axes (-2,-1) and the
zero axis is the number of the interrogation window
This should also work out of the box for rectangular windows.
Parameters
----------
image_a : 3d np.ndarray, first dimension is the number of windows,
and two last dimensions are interrogation windows of the first image
image_b : similar
correlation_method : string
one of the three methods implemented: 'circular' or 'linear'
[default: 'circular].
normalized_correlation : string
decides wetehr normalized correlation is done or not: True or False
[default: True].
"""
if normalized_correlation:
# remove the effect of stronger laser or
# longer exposure for frame B
# image_a = match_histograms(image_a, image_b)
# remove mean background, normalize to 0..1 range
image_a = normalize_intensity(image_a)
image_b = normalize_intensity(image_b)
s1 = np.array(image_a.shape[-2:])
s2 = np.array(image_b.shape[-2:])
if correlation_method == "linear":
# have to be normalized, mainly because of zero padding
size = s1 + s2 - 1
fsize = 2**np.ceil(np.log2(size)).astype(int)
fslice = (slice(0, image_a.shape[0]),
slice((fsize[0] - s1[0]) // 2, (fsize[0] + s1[0]) // 2),
slice((fsize[1] - s1[1]) // 2, (fsize[1] + s1[1]) // 2))
f2a = rfft2(image_a, fsize, axes=(-2, -1)).conj()
f2b = rfft2(image_b, fsize, axes=(-2, -1))
corr = fftshift(irfft2(f2a * f2b).real, axes=(-2, -1))[fslice]
elif correlation_method == "circular":
corr = fftshift(irfft2(rfft2(image_a).conj() * rfft2(image_b)).real,
axes=(-2, -1))
else:
print("method is not implemented!")
if normalized_correlation:
corr = corr / (s2[0] * s2[1]) # for extended search area
corr = np.clip(corr, 0, 1)
return corr
def deconvolve_sinogram(sinogram, psf, l=20, mode='laplacian', clip=True):
"""
Deconvolve a sinogram with given PSF.
Parameters
----------
sinogram : numpy.ndarray [TPY]
The blurred sinogram.
psf : numpy.ndarray [ZXY] or [XY]
The PSF used to deconvolve.
l : float, optional
Strength of the regularization.
clip : bool, optional
Clip negative values to 0, default is True.
Returns
-------
numpy.ndarray [TPY]
The deconvolved sinogram.
"""
fft_shape = [fft_size(s) for s in sinogram.shape[1:]]
inverse = inverse_psf_rfft(psf, shape=fft_shape, l=l, mode=mode)
s_fft = fft.rfft2(sinogram, s=fft_shape)
i_fft = fft.irfft2(s_fft * inverse, s=fft_shape, overwrite_x=True)
i_fft = i_fft[:, :sinogram.shape[1], :sinogram.shape[2]]
if clip:
np.clip(i_fft, 0, None, out=i_fft)
return i_fft
def __init__(
self,
mrc,
pixelsize,
voltage,
spherical_abberation,
amplitude_contrast,
low_cutoff_res=30,
high_cutoff_res=5,
):
self.pixelsize = pixelsize
self.voltage = voltage
self.spherical_abberation = spherical_abberation
self.amplitude_contrast = amplitude_contrast
self.electron_wavelength = wavelength_from_voltage(voltage)
if type(mrc) == str:
with mrcfile.open(mrc, permissive=True) as f:
img = f.data
elif type(mrc) == np.ndarray:
img = mrc
self.spectrum = np.log(
np.abs(fft.rfft2(img, s=(max(img.shape), max(img.shape)))))
# keep a copy of the full spectrum
self.full_spectrum = self.spectrum
# reduce the spectrum to a desired range and subtract background
self.spectrum = self._preprocess_spectrum(self.spectrum,
low_cutoff_res,
high_cutoff_res)
def fourier(B, dx, dt):
from scipy.fft import rfft2, rfftfreq
Bwk = rfft2(B)
w = rfftfreq(B.shape[1], d=dt)
k = rfftfreq(B.shape[0], d=dx)
k = k * 2 * np.pi
w = w * 2 * np.pi
return Bwk, w, k
def test_dirac(self):
for psf in self.psfs:
i_psf = fpsopt.inverse_psf_rfft(psf, l=0)
psfft = fft.rfft2(psf.sum(0))
dirac = fft.irfft2(psfft * i_psf, s=psf.shape[1:])
ref = np.zeros_like(dirac)
ref[(psf.shape[0] - 1) // 2, (psf.shape[0] - 1) // 2] = 1
np.testing.assert_allclose(dirac, ref, atol=1e-12)
def fft_correlate_windows(window_a, window_b):
""" FFT based cross correlation
it is a so-called linear convolution based,
since we increase the size of the FFT to
reduce the edge effects.
This should also work out of the box for rectangular windows.
Parameters
----------
window_a : 2d np.ndarray
a two dimensions array for the first interrogation window,
window_b : 2d np.ndarray
a two dimensions array for the second interrogation window.
# from Stackoverflow:
from scipy import linalg
import numpy as np
# works for rectangular windows as well
x = [[1 , 0 , 0 , 0] , [0 , -1 , 0 , 0] , [0 , 0 , 3 , 0] ,
[0 , 0 , 0 , 1], [0 , 0 , 0 , 1]]
x = np.array(x,dtype=np.float)
y = [[4 , 5] , [3 , 4]]
y = np.array(y)
print ("conv:" , signal.convolve2d(x , y , 'full'))
s1 = np.array(x.shape)
s2 = np.array(y.shape)
size = s1 + s2 - 1
fsize = 2 ** np.ceil(np.log2(size)).astype(int)
fslice = tuple([slice(0, int(sz)) for sz in size])
new_x = np.fft.fft2(x , fsize)
new_y = np.fft.fft2(y , fsize)
result = np.fft.ifft2(new_x*new_y)[fslice].copy()
print("fft for my method:" , np.array(result.real, np.int32))
"""
s1 = np.array(window_a.shape)
s2 = np.array(window_b.shape)
size = s1 + s2 - 1
fsize = 2 ** np.ceil(np.log2(size)).astype(int)
fslice = tuple([slice(0, int(sz)) for sz in size])
f2a = rfft2(window_a, fsize)
f2b = rfft2(window_b[::-1, ::-1], fsize)
corr = irfft2(f2a * f2b).real[fslice]
return corr
def roi_iter(self, data, flip=False):
s = self.roi.size
for j, i in self.roi.positions:
r = data[i - s:i + s, j - s:j + s]
r = r - r.mean()
r *= self.window[:, None]
r *= self.window[None, :]
n = np.sqrt((r * r).sum())
if n > 0:
r /= n # normalize
if flip:
r = r[::-1, ::-1]
yield fft.rfft2(r, (4 * s, 4 * s))
def texshadeFFT(x: np.ndarray, alpha: float) -> np.ndarray:
"""FFT-based texture shading elevation
Given an array `x` of elevation data and an `alpha` > 0, apply the
texture-shading algorithm using the full (real-only) FFT: the entire `x` array
will be FFT'd.
`alpha` is the shading detail factor, i.e., the power of the
fractional-Laplacian operator. `alpha=0` means no detail (output is the
input). `alpha=2.0` is the full (non-fractional) Laplacian operator and is
probably too high. `alpha <= 1.0` seem aesthetically pleasing.
Returns an array the same dimensions as `x` that contains the texture-shaded
version of the input array.
If `x` is memory-mapped and/or your system doesn't have 5x `x`'s memory
available, consider using `texshade.texshadeSpatial`, which implements a
low-memory version of the algorithm by approximating the frequency response of
the fractional-Laplacian filter with a finite impulse response filter applied
in the spatial-domain.
Implementation note: this function uses Scipy's FFTPACK routines (in
`scipy.fftpack`) instead of Numpy's FFT (`numpy.fft`) because the former can
return single-precision float32. In newer versions of Numpy/Scipy, this
advantage may have evaporated [1], [2].
[1] https://github.com/numpy/numpy/issues/6012
[2] https://github.com/scipy/scipy/issues/2487
"""
Nyx = [nextprod([2, 3, 5, 7], x) for x in x.shape]
# Generate filter in the frequency domain
fy = sf.fftfreq(Nyx[0])[:, np.newaxis].astype(x.dtype)
fx = sf.rfftfreq(Nyx[1])[np.newaxis, :].astype(x.dtype)
H2 = (fx**2 + fy**2)**(alpha / 2.0)
# Compute the FFT of the input and apply the filter
xr = sf.rfft2(x, s=Nyx) * H2
H2 = None # potentially trigger GC here to reclaim H2's memory
xr = sf.irfft2(xr)
# Return the same size as input
return xr[:x.shape[0], :x.shape[1]]
return i_fft
if __name__ == '__main__':
from cbi_toolbox.simu import optics, primitives, imaging
import cbi_toolbox.splineradon as spl
import napari
TEST_SIZE = 64
s_psf = optics.gaussian_psf(numerical_aperture=0.3,
npix_axial=TEST_SIZE + 1,
npix_lateral=TEST_SIZE + 1)
i_psf = inverse_psf_rfft(s_psf, l=1e-15, mode='constant')
psfft = fft.rfft2(s_psf.sum(0))
dirac = fft.irfft2(psfft * i_psf, s=s_psf.shape[1:])
sample = primitives.boccia(TEST_SIZE,
radius=(0.8 * TEST_SIZE) // 2,
n_stripes=4)
s_theta = np.arange(90)
s_radon = spl.radon(sample, theta=s_theta, circle=True)
s_fpsopt = imaging.fps_opt(sample, s_psf, theta=s_theta)
s_deconv = deconvolve_sinogram(s_fpsopt, s_psf, l=0)
viewer = napari.view_image(s_radon)
viewer.add_image(s_fpsopt)
viewer.add_image(s_deconv)
def test_rfft2(self):
x = random((30, 20))
assert_array_almost_equal(fft.fft2(x)[:, :11], fft.rfft2(x))
assert_array_almost_equal(fft.rfft2(x) / np.sqrt(30 * 20),
fft.rfft2(x, norm="ortho"))
def test_irfft2(self):
x = random((30, 20))
assert_array_almost_equal(x, fft.irfft2(fft.rfft2(x)))
assert_array_almost_equal(
x, fft.irfft2(fft.rfft2(x, norm="ortho"), norm="ortho"))
def test_irfft2(self):
x = random((30, 20))
assert_array_almost_equal(x, fft.irfft2(fft.rfft2(x)))
for norm in ["backward", "ortho", "forward"]:
assert_array_almost_equal(
x, fft.irfft2(fft.rfft2(x, norm=norm), norm=norm))
