def main(args):
log = logging.getLogger('root')
hdlr = logging.StreamHandler(sys.stdout)
log.addHandler(hdlr)
log.setLevel(logging.getLevelName(args.loglevel.upper()))
pyfftw.interfaces.cache.enable()
vol1 = mrc.read(args.volume1, inc_header=False, compat="relion")
vol2 = mrc.read(args.volume2, inc_header=False, compat="relion")
if args.mask is not None:
mask = mrc.read(args.mask, inc_header=False, compat="relion")
vol1 *= mask
vol2 *= mask
f3d1 = fft.rfftn(vol1, threads=args.threads)
f3d2 = fft.rfftn(vol2, threads=args.threads)
nside = 2**args.healpix_order
x, y, z = pix2vec(nside, np.arange(12 * nside ** 2))
xhalf = x >= 0
hp = np.column_stack([x[xhalf], y[xhalf], z[xhalf]])
t0 = time.time()
fcor = calc_dfsc(f3d1, f3d2, hp, np.deg2rad(args.arc))
log.info("Computed CFSC in %0.2f s" % (time.time() - t0))
fsc = calc_fsc(f3d1, f3d2)
t0 = time.time()
log.info("Computed GFSC in %0.2f s" % (time.time() - t0))
freqs = np.fft.rfftfreq(f3d1.shape[0])
np.save(args.output, np.row_stack([freqs, fsc, fcor]))
return 0
def fftconvolve(in1, in2, mode="same", threads=1):
"""Same as above but with pyfftw added in"""
in1 = np.asarray(in1)
in2 = np.asarray(in2)
if in1.ndim == in2.ndim == 0: # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("in1 and in2 should have the same dimensionality")
elif in1.size == 0 or in2.size == 0: # empty arrays
return np.array([])
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
complex_result = (np.issubdtype(in1.dtype, complex)
or np.issubdtype(in2.dtype, complex))
shape = s1 + s2 - 1
# Check that input sizes are compatible with 'valid' mode
if sig._inputs_swap_needed(mode, s1, s2):
# Convolution is commutative; order doesn't have any effect on output
in1, s1, in2, s2 = in2, s2, in1, s1
# Speed up FFT by padding to optimal size for FFTPACK
fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in shape]
fslice = tuple([slice(0, int(sz)) for sz in shape])
# Pre-1.9 NumPy FFT routines are not threadsafe. For older NumPys, make
# sure we only call rfftn/irfftn from one thread at a time.
if not complex_result and (sig._rfft_mt_safe
or sig._rfft_lock.acquire(False)):
try:
sp1 = rfftn(in1, fshape, threads=threads)
sp2 = rfftn(in2, fshape, threads=threads)
ret = (irfftn(sp1 * sp2, fshape,
threads=threads)[fslice].copy())
finally:
if not sig._rfft_mt_safe:
sig._rfft_lock.release()
else:
# If we're here, it's either because we need a complex result, or we
# failed to acquire _rfft_lock (meaning rfftn isn't threadsafe and
# is already in use by another thread). In either case, use the
# (threadsafe but slower) SciPy complex-FFT routines instead.
sp1 = fftn(in1, fshape, threads=threads)
sp2 = fftn(in2, fshape, threads=threads)
ret = ifftn(sp1 * sp2, threads=threads)[fslice].copy()
if not complex_result:
ret = ret.real
if mode == "full":
return ret
elif mode == "same":
return sig._centered(ret, s1)
elif mode == "valid":
return sig._centered(ret, s1 - s2 + 1)
else:
raise ValueError("Acceptable mode flags are 'valid',"
" 'same', or 'full'.")
def fftconvolve(in1, in2, mode="same", threads=1):
"""Same as above but with pyfftw added in"""
in1 = np.asarray(in1)
in2 = np.asarray(in2)
if in1.ndim == in2.ndim == 0: # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("in1 and in2 should have the same dimensionality")
elif in1.size == 0 or in2.size == 0: # empty arrays
return np.array([])
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
complex_result = (np.issubdtype(in1.dtype, complex) or
np.issubdtype(in2.dtype, complex))
shape = s1 + s2 - 1
# Check that input sizes are compatible with 'valid' mode
if sig._inputs_swap_needed(mode, s1, s2):
# Convolution is commutative; order doesn't have any effect on output
in1, s1, in2, s2 = in2, s2, in1, s1
# Speed up FFT by padding to optimal size for FFTPACK
fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in shape]
fslice = tuple([slice(0, int(sz)) for sz in shape])
# Pre-1.9 NumPy FFT routines are not threadsafe. For older NumPys, make
# sure we only call rfftn/irfftn from one thread at a time.
if not complex_result and (sig._rfft_mt_safe or sig._rfft_lock.acquire(False)):
try:
sp1 = rfftn(in1, fshape, threads=threads)
sp2 = rfftn(in2, fshape, threads=threads)
ret = (irfftn(sp1 * sp2, fshape, threads=threads)[fslice].copy())
finally:
if not sig._rfft_mt_safe:
sig._rfft_lock.release()
else:
# If we're here, it's either because we need a complex result, or we
# failed to acquire _rfft_lock (meaning rfftn isn't threadsafe and
# is already in use by another thread). In either case, use the
# (threadsafe but slower) SciPy complex-FFT routines instead.
sp1 = fftn(in1, fshape, threads=threads)
sp2 = fftn(in2, fshape, threads=threads)
ret = ifftn(sp1 * sp2, threads=threads)[fslice].copy()
if not complex_result:
ret = ret.real
if mode == "full":
return ret
elif mode == "same":
return sig._centered(ret, s1)
elif mode == "valid":
return sig._centered(ret, s1 - s2 + 1)
else:
raise ValueError("Acceptable mode flags are 'valid',"
" 'same', or 'full'.")
def fftconvolvefast(in1, in2):
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
shape = s1 + s2 - 1
fshape = [fftpack.helper.next_fast_len(int(d)) for d in shape]
fslice = tuple([slice(0, int(sz)) for sz in shape])
sp1 = rfftn(in1, fshape)
sp2 = rfftn(in2, fshape)
ret = (irfftn(sp1 * sp2, fshape)[fslice].copy())
return _centered(ret, s1 - s2 + 1)
def fft_gaussian_filter(img, sigma):
"""FFT gaussian convolution
Parameters
----------
img : ndarray
Image to convolve with a gaussian kernel
sigma : int or sequence
The sigma(s) of the gaussian kernel in _real space_
Returns
-------
filt_img : ndarray
The filtered image
"""
# This doesn't help agreement but it will make things faster
# pull the shape
s1 = np.array(img.shape)
# s2 = np.array([int(s * 4) for s in _normalize_sequence(sigma, img.ndim)])
shape = s1 # + s2 - 1
# calculate a nice shape
fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in shape]
# pad out with reflection
pad_img = fft_pad(img, fshape, "reflect")
# calculate the padding
padding = tuple(_calc_pad(o, n) for o, n in zip(img.shape, pad_img.shape))
# so that we can calculate the cropping, maybe this should be integrated
# into `fft_pad` ...
fslice = tuple(
slice(s, -e) if e != 0 else slice(s, None) for s, e in padding)
# fourier transfrom and apply the filter
kimg = rfftn(pad_img, fshape)
filt_kimg = fourier_gaussian(kimg, sigma, pad_img.shape[-1])
# inverse FFT and return.
return irfftn(filt_kimg, fshape)[fslice]
def rfftn(a, axes=None):
if _use_fftw:
from pyfftw.interfaces.numpy_fft import rfftn
_init_pyfftw()
return rfftn(a, axes=axes, **_fft_extra_args)
else:
return np.fft.rfftn(a, axes=axes)
def urdftn(inarray, ndim=None, *args, **kwargs):
"""N-dim real unitary discrete Fourier transform
This transform consider the Hermitian property of the transform on
real input
Parameters
----------
inarray : ndarray
The array to transform.
ndim : int, optional
The `ndim` last axis along wich to compute the transform. All
axes by default.
Returns
-------
outarray : array-like (the last ndim as N / 2 + 1 lenght)
"""
if not ndim:
ndim = inarray.ndim
return rfftn(inarray, axes=range(-ndim, 0), *args, **kwargs) / np.sqrt(
np.prod(inarray.shape[-ndim:]))
def fft_forward(self):
"""Compute the direct Fourier transform of the input data."""
logger.debug('Computing FFT of input with %s threads', self.nthreads)
data = self.data_raw
# Note: we don't use the defaults argument for the FFT as we
# don't want any planning effort going to this one
self.data_raw_f = fft.rfftn(data, threads=self.nthreads)
def vol_ft(vol, pfac=2, threads=1, normfft=1):
""" Returns a centered, Nyquist-limited, zero-padded, interpolation-ready 3D Fourier transform.
:param vol: Volume to be Fourier transformed.
:param pfac: Size factor for zero-padding.
:param threads: Number of threads for pyFFTW.
:param normfft: Normalization constant for Fourier transform.
"""
vol = grid_correct(vol.astype(np.float64), pfac=pfac, order=1)
padvol = np.pad(vol, (vol.shape[0] * pfac - vol.shape[0]) // 2, "constant")
ft = rfftn(np.fft.ifftshift(padvol), padvol.shape, threads=threads)
ftc = np.zeros((ft.shape[0] + 3, ft.shape[1] + 3, ft.shape[2]), dtype=np.complex128)
fill_ft(ft, ftc, vol.shape[0], normfft=normfft)
return ftc
def fftconvolve_fast(data, kernel, **kwargs):
"""A faster version of fft convolution
In this case the kernel ifftshifted before FFT but the data is not.
This can be done because the effect of fourier convolution is to
"wrap" around the data edges so whether we ifftshift before FFT
and then fftshift after it makes no difference so we can skip the
step entirely.
"""
# TODO: add error checking like in the above and add functionality
# for complex inputs. Also could add options for different types of
# padding.
dshape = np.array(data.shape)
kshape = np.array(kernel.shape)
# find maximum dimensions
maxshape = np.max((dshape, kshape), 0)
# calculate a nice shape
fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in maxshape]
# pad out with reflection
pad_data = fft_pad(data, fshape, "reflect")
# calculate padding
padding = tuple(
_calc_pad(o, n) for o, n in zip(data.shape, pad_data.shape))
# so that we can calculate the cropping, maybe this should be integrated
# into `fft_pad` ...
fslice = tuple(
slice(s, -e) if e != 0 else slice(s, None) for s, e in padding)
if kernel.shape != pad_data.shape:
# its been assumed that the background of the kernel has already been
# removed and that the kernel has already been centered
kernel = fft_pad(kernel, pad_data.shape, mode='constant')
k_kernel = rfftn(ifftshift(kernel), pad_data.shape, **kwargs)
k_data = rfftn(pad_data, pad_data.shape, **kwargs)
convolve_data = irfftn(k_kernel * k_data, pad_data.shape, **kwargs)
# return data with same shape as original data
return convolve_data[fslice]
def vol_ft(vol, pfac=2, threads=1, normfft=1):
""" Returns a centered, Nyquist-limited, zero-padded, interpolation-ready 3D Fourier transform.
:param vol: Volume to be Fourier transformed.
:param pfac: Size factor for zero-padding.
:param threads: Number of threads for pyFFTW.
:param normfft: Normalization constant for Fourier transform.
"""
vol = grid_correct(vol, pfac=pfac, order=1)
padvol = np.pad(vol, int((vol.shape[0] * pfac - vol.shape[0]) // 2), "constant")
ft = rfftn(np.fft.ifftshift(padvol), padvol.shape, threads=threads)
# ft = np.fft.fftn(np.fft.ifftshift(padvol), padvol.shape )
ftc = np.zeros((ft.shape[0] + 3, ft.shape[1] + 3, ft.shape[2]), dtype=ft.dtype)
vop.fill_ft(ft, ftc, vol.shape[0], normfft=normfft)
print( ftc.shape )
plot2dimage(np.log10(np.abs(ftc[:][:][230])).astype(np.float))
sys.exit()
return ftc
def from_array(cls, corr_array, is_cyclic=True):
"""Create an instance with the given correlations.
Parameters
----------
corr_array: array_like
The correlation of the first element of the domain with
each other element.
is_cyclic: bool
Returns
-------
HomogeneousIsotropicCorrelation
"""
corr_array = asarray(corr_array)
shape = corr_array.shape
ndims = corr_array.ndim
if is_cyclic:
computational_shape = shape
self = cls(shape, computational_shape)
corr_fourier = self._fft(corr_array)
else:
computational_shape = tuple(2 * (dim - 1) for dim in shape)
self = cls(shape, computational_shape)
for axis in reversed(range(ndims)):
corr_array = concatenate(
[corr_array, flip(corr_array[1:-1], axis)], axis=axis)
# Advantages over dctn: guaranteed same format and gets
# nice planning done for the later evaluations.
corr_fourier = rfftn(corr_array,
axes=arange(0, ndims, dtype=int),
threads=NUM_THREADS,
planner_effort=ADVANCE_PLANNER_EFFORT)
# The fft axes need to be a single chunk for the dask ffts
# It's in memory already anyway
# TODO: create a from_spectrum to delegate to
self._corr_fourier = (corr_fourier)
self._fourier_near_zero = (corr_fourier < FOURIER_NEAR_ZERO)
return self
def hartley(a, axes=None):
# Check if the axes provided are valid given the shape
if axes is not None and \
not all(axis < len(a.shape) for axis in axes):
raise ValueError("Provided axes do not match array shape")
if iscomplextype(a.dtype):
raise TypeError("Hartley transform requires real-valued arrays.")
tmp = rfftn(a, axes=axes)
def _fill_array(tmp, res, axes):
if axes is None:
axes = tuple(range(tmp.ndim))
lastaxis = axes[-1]
ntmplast = tmp.shape[lastaxis]
slice1 = (slice(None),)*lastaxis + (slice(0, ntmplast),)
np.add(tmp.real, tmp.imag, out=res[slice1])
def _fill_upper_half(tmp, res, axes):
lastaxis = axes[-1]
nlast = res.shape[lastaxis]
ntmplast = tmp.shape[lastaxis]
nrem = nlast - ntmplast
slice1 = [slice(None)]*lastaxis + [slice(ntmplast, None)]
slice2 = [slice(None)]*lastaxis + [slice(nrem, 0, -1)]
for i in axes[:-1]:
slice1[i] = slice(1, None)
slice2[i] = slice(None, 0, -1)
slice1 = tuple(slice1)
slice2 = tuple(slice2)
np.subtract(tmp[slice2].real, tmp[slice2].imag, out=res[slice1])
for i, ax in enumerate(axes[:-1]):
dim1 = (slice(None),)*ax + (slice(0, 1),)
axes2 = axes[:i] + axes[i+1:]
_fill_upper_half(tmp[dim1], res[dim1], axes2)
_fill_upper_half(tmp, res, axes)
return res
return _fill_array(tmp, np.empty_like(a), axes)
def fft_gaussian_filter(img, sigma):
"""FFT gaussian convolution
Parameters
----------
img : ndarray
Image to convolve with a gaussian kernel
sigma : int or sequence
The sigma(s) of the gaussian kernel in _real space_
Returns
-------
filt_img : ndarray
The filtered image
"""
# This doesn't help agreement but it will make things faster
# pull the shape
s1 = np.array(img.shape)
# if any of the sizes is 32 or smaller, revert to proper filter
if any(s1 < 33):
warnings.warn(("Input is small along a dimension,"
" will revert to `gaussian_filter`"))
return gaussian_filter(img, sigma)
# s2 = np.array([int(s * 4) for s in _normalize_sequence(sigma, img.ndim)])
shape = s1 # + s2 - 1
# calculate a nice shape
fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in shape]
# pad out with reflection
pad_img = fft_pad(img, fshape, "reflect")
# calculate the padding
padding = tuple(_calc_pad(o, n) for o, n in zip(img.shape, pad_img.shape))
# so that we can calculate the cropping, maybe this should be integrated
# into `fft_pad` ...
fslice = tuple(slice(s, -e) if e != 0 else slice(s, None)
for s, e in padding)
# fourier transfrom and apply the filter
kimg = rfftn(pad_img, fshape)
filt_kimg = fourier_gaussian(kimg, sigma, pad_img.shape[-1])
# inverse FFT and return.
return irfftn(filt_kimg, fshape)[fslice]
def my_fftn_r2c(a, axes=None):
# Check if the axes provided are valid given the shape
if axes is not None and \
not all(axis < len(a.shape) for axis in axes):
raise ValueError("Provided axes do not match array shape")
if iscomplextype(a.dtype):
raise TypeError("Transform requires real-valued input arrays.")
tmp = rfftn(a, axes=axes)
def _fill_complex_array(tmp, res, axes):
if axes is None:
axes = tuple(range(tmp.ndim))
lastaxis = axes[-1]
ntmplast = tmp.shape[lastaxis]
slice1 = [slice(None)]*lastaxis + [slice(0, ntmplast)]
res[tuple(slice1)] = tmp
def _fill_upper_half_complex(tmp, res, axes):
lastaxis = axes[-1]
nlast = res.shape[lastaxis]
ntmplast = tmp.shape[lastaxis]
nrem = nlast - ntmplast
slice1 = [slice(None)]*lastaxis + [slice(ntmplast, None)]
slice2 = [slice(None)]*lastaxis + [slice(nrem, 0, -1)]
for i in axes[:-1]:
slice1[i] = slice(1, None)
slice2[i] = slice(None, 0, -1)
# np.conjugate(tmp[slice2], out=res[slice1])
res[tuple(slice1)] = np.conjugate(tmp[tuple(slice2)])
for i, ax in enumerate(axes[:-1]):
dim1 = tuple([slice(None)]*ax + [slice(0, 1)])
axes2 = axes[:i] + axes[i+1:]
_fill_upper_half_complex(tmp[dim1], res[dim1], axes2)
_fill_upper_half_complex(tmp, res, axes)
return res
return _fill_complex_array(tmp, np.empty_like(a, dtype=tmp.dtype), axes)
def urfftn(inarray, dim=None):
"""N-dimensional real unitary Fourier transform.
This transform considers the Hermitian property of the transform on
real-valued input.
Parameters
----------
inarray : ndarray, shape (M, N, ..., P)
The array to transform.
dim : int, optional
The last axis along which to compute the transform. All
axes by default.
Returns
-------
outarray : ndarray, shape (M, N, ..., P / 2 + 1)
The unitary N-D real Fourier transform of ``inarray``.
Notes
-----
The ``urfft`` functions assume an input array of real
values. Consequently, the output has a Hermitian property and
redundant values are not computed or returned.
Examples
--------
>>> input = np.ones((5, 5, 5))
>>> output = urfftn(input)
>>> np.allclose(np.sum(input) / np.sqrt(input.size), output[0, 0, 0])
True
>>> output.shape
(5, 5, 3)
"""
if dim is None:
dim = inarray.ndim
outarray = rfftn(inarray, axes=range(-dim, 0))
return outarray / np.sqrt(np.prod(inarray.shape[-dim:]))
def generate_white_noise(self):
"""Generate a white noise with the relevant power spectrum.
Returns
-------
white_noise : np.ndarray
"""
# Compute the k grid
d = self.Lbox / self.dimensions / (2 * np.pi)
all_k = [fft.fftfreq(self.dimensions, d=d)] * (self.Ndim - 1) + [
fft.rfftfreq(self.dimensions, d=d)
]
self.kgrid = kgrid = np.array(np.meshgrid(*all_k, indexing="ij"))
self.knorm = knorm = np.sqrt(np.sum(kgrid**2, axis=0))
# Compute Pk
Pk = np.zeros_like(knorm)
mask = knorm > 0
Pk[mask] = self.Pk(knorm[mask] * 2)
# Compute white noise (in Fourier space)
mu = np.random.standard_normal([self.dimensions] * self.Ndim)
muk = fft.rfftn(mu)
deltak = muk * np.sqrt(Pk)
# Compute field in real space
white_noise = fft.irfftn(deltak)
# Normalize variance
deltak_smoothed = deltak * self.filter.W(knorm)
field = fft.irfftn(deltak_smoothed)
std = field.std()
self.white_noise_fft = deltak * self.sigma8 / std
self.white_noise = white_noise * self.sigma8 / std
return self.white_noise
def pattern_params(my_pat, size=2):
"""Find stuff"""
# REAL FFT!
# note the limited shifting, we don't want to shift the last axis
my_pat_fft = fftshift(rfftn(ifftshift(my_pat)),
axes=tuple(range(my_pat.ndim))[:-1])
my_abs_pat_fft = abs(my_pat_fft)
# find dc loc, center of FFT after shifting
sizeky, sizekx = my_abs_pat_fft.shape
# remember we didn't shift the last axis!
dc_loc = (sizeky // 2, 0)
# mask data and find next biggest peak
dc_power = my_abs_pat_fft[dc_loc]
my_abs_pat_fft[dc_loc] = 0
max_loc = np.unravel_index(my_abs_pat_fft.argmax(), my_abs_pat_fft.shape)
# pull the 3x3 region around the peak and fit
max_shift = localize_peak(my_abs_pat_fft[slice_maker(max_loc, 3)])
# calculate precise peak relative to dc
peak = np.array(max_loc) + np.array(max_shift) - np.array(dc_loc)
# correct location based on initial data shape
peak_corr = peak / np.array(my_pat.shape)
# calc angle
preciseangle = np.arctan2(*peak_corr)
# calc period
precise_period = 1 / norm(peak_corr)
# calc phase
phase = np.angle(my_pat_fft[max_loc[0], max_loc[1]])
# calc modulation depth
numerator = abs(my_pat_fft[slice_maker(max_loc, size)].sum())
mod = numerator / dc_power
return {"period": precise_period,
"angle": preciseangle,
"phase": phase,
"fft": my_pat_fft,
"mod": mod,
"max_loc": max_loc}
def rfft_to_fft(image,
keep_rfft=False,
use_pyfftw=False,
threads=1,
**pyfftw_kwargs):
'''
Perform a RFFT on the image (2 or 3D) and return the absolute value in
the same format as you would get with the fft (negative frequencies).
This avoids ever having to have the full complex cube in memory.
Inputs
------
image : numpy.ndarray
2 or 3D array.
keep_rfft : bool, optional
Return the rfft output instead of expanding to the
negative frequencies of the full FFT.
use_pyfftw : bool, optional
Try using pyfftw for the FFT.
threads : int, optional
Number of threads to use when using pyfftw. Default is 1.
pyfftw_kwargs : Passed to `~pyfftw.interfaces.numpy_fft.rfftn`.
Outputs
-------
fft_abs : absolute value of the fft.
'''
ndim = len(image.shape)
if ndim < 2 or ndim > 3:
raise TypeError("Dimension of image must be 2D or 3D.")
last_dim = image.shape[-1]
if use_pyfftw:
if PYFFTW_FLAG:
fft_abs = np.abs(rfftn(image, **pyfftw_kwargs))
else:
use_pyfftw = False
warn("pyfftw is not installed")
if not use_pyfftw:
fft_abs = np.abs(np.fft.rfftn(image))
if keep_rfft:
return fft_abs
if ndim == 2:
if last_dim % 2 == 0:
fftstar_abs = fft_abs.copy()[:, -2:0:-1]
else:
fftstar_abs = fft_abs.copy()[:, -1:0:-1]
fftstar_abs[1::, :] = fftstar_abs[:0:-1, :]
return np.concatenate((fft_abs, fftstar_abs), axis=1)
elif ndim == 3:
if last_dim % 2 == 0:
fftstar_abs = fft_abs.copy()[:, :, -2:0:-1]
else:
fftstar_abs = fft_abs.copy()[:, :, -1:0:-1]
fftstar_abs[1::, :, :] = fftstar_abs[:0:-1, :, :]
fftstar_abs[:, 1::, :] = fftstar_abs[:, :0:-1, :]
return np.concatenate((fft_abs, fftstar_abs), axis=2)
def ir2tf(imp_resp, shape, dim=None, is_real=True):
"""Compute the transfer function of an impulse response (IR).
This function makes the necessary correct zero-padding, zero
convention, correct fft2, etc... to compute the transfer function
of IR. To use with unitary Fourier transform for the signal (ufftn
or equivalent).
Parameters
----------
imp_resp : ndarray
The impulse responses.
shape : tuple of int
A tuple of integer corresponding to the target shape of the
transfer function.
dim : int, optional
The last axis along which to compute the transform. All
axes by default.
is_real : boolean (optional, default True)
If True, imp_resp is supposed real and the Hermitian property
is used with rfftn Fourier transform.
Returns
-------
y : complex ndarray
The transfer function of shape ``shape``.
See Also
--------
ufftn, uifftn, urfftn, uirfftn
Examples
--------
>>> np.all(np.array([[4, 0], [0, 0]]) == ir2tf(np.ones((2, 2)), (2, 2)))
True
>>> ir2tf(np.ones((2, 2)), (512, 512)).shape == (512, 257)
True
>>> ir2tf(np.ones((2, 2)), (512, 512), is_real=False).shape == (512, 512)
True
Notes
-----
The input array can be composed of multiple-dimensional IR with
an arbitrary number of IR. The individual IR must be accessed
through the first axes. The last ``dim`` axes contain the space
definition.
"""
if not dim:
dim = imp_resp.ndim
# Zero padding and fill
irpadded = np.zeros(shape)
irpadded[tuple([slice(0, s) for s in imp_resp.shape])] = imp_resp
# Roll for zero convention of the fft to avoid the phase
# problem. Work with odd and even size.
for axis, axis_size in enumerate(imp_resp.shape):
if axis >= imp_resp.ndim - dim:
irpadded = np.roll(irpadded,
shift=-int(np.floor(axis_size / 2)),
axis=axis)
if is_real:
return rfftn(irpadded, axes=range(-dim, 0))
else:
return fftn(irpadded, axes=range(-dim, 0))
def from_function(cls, corr_func, shape, is_cyclic=True):
"""Create an instance to apply the correlation function.
Parameters
----------
corr_func: callable(dist) -> float
The correlation of the first element of the domain with
each other element.
shape: tuple of int
The state is formally a vector, but the correlations are
assumed to depend on the layout in some other shape,
usually related to the physical layout. This is the other
shape.
is_cyclic: bool
Whether to assume the domain is periodic in all directions.
Returns
-------
HomogeneousIsotropicCorrelation
"""
shape = np.atleast_1d(shape)
if is_cyclic:
computational_shape = tuple(shape)
else:
computational_shape = tuple(
next_fast_len(2 * dim - 1) for dim in shape)
self = cls(tuple(shape), computational_shape)
shape = np.asarray(self._computational_shape)
ndims = len(shape)
broadcastable_shape = shape[:, newaxis]
while broadcastable_shape.ndim < ndims + 1:
broadcastable_shape = broadcastable_shape[..., newaxis]
def corr_from_index(*index):
"""Correlation of index with zero.
Turns a correlation function in terms of index distance
into one in terms of indices on a periodic domain.
Parameters
----------
index: tuple of int
Returns
-------
float[-1, 1]
The correlation of the given index with the origin.
See Also
--------
DistanceCorrelationFunction.correlation_from_index
"""
comp2_1 = square(index)
# Components of distance to shifted origin
comp2_2 = square(broadcastable_shape - index)
# use the smaller components to get the distance to the
# closest of the shifted origins
comp2 = fmin(comp2_1, comp2_2)
return corr_func(sqrt(array_sum(comp2, axis=0)))
corr_struct = fromfunction(corr_from_index,
shape=tuple(shape),
dtype=DTYPE)
# I should be able to generate this sequence with a type-I DCT
# For some odd reason complex/complex is faster than complex/real
# This also ensures the format here is the same as in _matmat
corr_fourier = rfftn(corr_struct,
axes=arange(ndims, dtype=int),
threads=NUM_THREADS,
planner_effort=ADVANCE_PLANNER_EFFORT)
self._corr_fourier = (corr_fourier)
# This is also affected by roundoff
abs_corr_fourier = abs(corr_fourier)
self._fourier_near_zero = (abs_corr_fourier <
FOURIER_NEAR_ZERO * abs_corr_fourier.max())
return self
def ir2fr(imp_resp, shape, center=None, real=True):
"""Return the frequency response from impulsionnal responses
This function make the necessary correct zero-padding, zero
convention, correct DFT etc. to compute the frequency response
from impulsionnal responses (IR).
The IR array is supposed to have the origin in the middle of the
array.
The Fourier transform is performed on the last `len(shape)`
dimensions.
Parameters
----------
imp_resp : ndarray
The impulsionnal responses.
shape : tuple of int
A tuple of integer corresponding to the target shape of the
frequency responses, without hermitian property.
center : tuple of int, optional
The origin index of the impulsionnal response. The middle by
default.
real : boolean (optionnal, default True)
If True, imp_resp is supposed real, the hermissian property is
used with rfftn DFT and the output has `shape[-1] / 2 + 1`
elements on the last axis.
Returns
-------
y : ndarray
The frequency responses of shape `shape` on the last
`len(shape)` dimensions.
Notes
-----
- For convolution, the result have to be used with unitary
discrete Fourier transform for the signal (udftn or equivalent).
- DFT are always peformed on last axis for efficiency.
- Results is always C-contiguous.
See Also
--------
udftn, uidftn, urdftn, uirdftn
"""
if len(shape) > imp_resp.ndim:
raise ValueError("length of shape must inferior to imp_resp.ndim")
if not center:
center = [int(np.floor(length / 2)) for length in imp_resp.shape]
if len(center) != len(shape):
raise ValueError("center and shape must have the same length")
# Place the provided IR at the beginning of the array
irpadded = np.zeros(shape)
irpadded[tuple([slice(0, s) for s in imp_resp.shape])] = imp_resp
# Roll, or circshift to place the origin at 0 index, the
# hypothesis of the DFT
for axe, shift in enumerate(center):
irpadded = np.roll(irpadded, -shift, imp_resp.ndim - len(shape) + axe)
# Perform the DFT on the last axes
if real:
return np.ascontiguousarray(
rfftn(irpadded,
axes=list(range(imp_resp.ndim - len(shape), imp_resp.ndim))))
else:
return np.ascontiguousarray(
fftn(irpadded,
axes=list(range(imp_resp.ndim - len(shape), imp_resp.ndim))))
def fft(*args, **kwargs):
return fftw.rfftn(*args, **kwargs, threads=2)
def drift_plot(fit, title=None, dt=0.1, dx=130, lf=-np.inf, hf=np.inf,
log=False, cmap='magma', xc='b', yc='r'):
'''
Plotting utility to show drift curves nicely
Parameters
----------
fit : pandas DataFrame
Assumes that it has attributes x0 and y0
title : str (optional)
Title of plot
dt : float (optional)
Sampling rate of data in seconds
dx : pixel size (optional)
Pixel size in nm
lf : float (optional)
Low frequency cutoff for fourier plot
hf : float (optional)
High frequency cutoff for fourier plot
log : bool (optional)
Take logarithm of FFT data before displaying
cmap : string or matplotlib.colors.cmap instance
Color map for scatter plot
xc : string or `color` instance
Color for x data
yc : string or `color` instance
Color for y data
Returns
-------
fig : figure object
The figure
axs : tuple of axes objects
In the following order, Real axis, FFT axis, Scatter axis
'''
# set up plot
fig = plt.figure()
fig.set_size_inches(8, 4)
axreal = plt.subplot(221)
axfft = plt.subplot(223)
axscatter = plt.subplot(122)
# label it
if title is not None:
fig.suptitle(title, y=1.02, fontweight='bold')
# detrend mean
ybar = fit.y0 - fit.y0.mean()
ybar *= dx
xbar = fit.x0 - fit.x0.mean()
xbar *= dx
# Plot Real space
t = np.arange(len(fit)) * dt
axreal.plot(t, xbar, xc, label=r"$x_0$")
axreal.plot(t, ybar, yc, label=r"$y_0$")
axreal.set_xlabel('Time (s)')
axreal.set_ylabel('Displacement (nm)')
# add legend to real axis
axreal.legend(loc='best')
# calc FFTs
Y = rfftn(ybar)
X = rfftn(xbar)
# calc FFT freq
k = rfftfreq(len(fit), dt)
# limit FFT display range
kg = np.logical_and(k > lf, k < hf)
# Plot FFT
if log:
axfft.semilogy(k[kg], abs(X[kg]), xc)
axfft.semilogy(k[kg], abs(Y[kg]), yc)
else:
axfft.plot(k[kg], abs(X[kg]), xc)
axfft.plot(k[kg], abs(Y[kg]), yc)
axfft.set_xlabel('Frequency (Hz)')
# Plot scatter
axscatter.scatter(xbar, ybar, c=t, cmap=cmap)
axscatter.set_xlabel('x')
axscatter.set_ylabel('y')
# make sure the scatter plot is square
lims = axreal.get_ylim()
axscatter.set_ylim(lims)
axscatter.set_xlim(lims)
# tight layout
axs = (axreal, axfft, axscatter)
fig.tight_layout()
# return fig, axs to user for further manipulation and/or saving if wanted
return fig, axs
def cross_correlation_3d(pixels1, pixels2):
'''Align the second image with the first using max cross-correlation
returns the z,y,x offsets to add to image1's indexes to align it with
image2
Many of the ideas here are based on the paper, "Fast Normalized
Cross-Correlation" by J.P. Lewis
(http://www.idiom.com/~zilla/Papers/nvisionInterface/nip.html)
which is frequently cited when addressing this problem.
'''
s = np.maximum(pixels1.shape, pixels2.shape)
fshape = s*2
#
# Calculate the # of pixels at a particular point
#
i,j,k = np.mgrid[-s[0]:s[0], -s[1]:s[1], -s[2]:s[2] ]
unit = np.abs(i*j*k).astype(float)
unit[unit<1]=1 # keeps from dividing by zero in some places
#
# Normalize the pixel values around zero which does not affect the
# correlation, keeps some of the sums of multiplications from
# losing precision and precomputes t(x-u,y-v) - t_mean
#
pixels1 = np.nan_to_num(pixels1-nanmean(pixels1))
pixels2 = np.nan_to_num(pixels2-nanmean(pixels2))
#
# Lewis uses an image, f and a template t. He derives a normalized
# cross correlation, ncc(u,v) =
# sum((f(x,y)-f_mean(u,v))*(t(x-u,y-v)-t_mean),x,y) /
# sqrt(sum((f(x,y)-f_mean(u,v))**2,x,y) * (sum((t(x-u,y-v)-t_mean)**2,x,y)
#
# From here, he finds that the numerator term, f_mean(u,v)*(t...) is zero
# leaving f(x,y)*(t(x-u,y-v)-t_mean) which is a convolution of f
# by t-t_mean.
#
fp1 = rfftn(pixels1.astype('float32'), fshape, axes=(0, 1, 2))
fp2 = rfftn(pixels2.astype('float32'), fshape, axes=(0, 1, 2))
corr12 = irfftn(fp1 * fp2.conj(), axes=(0, 1, 2)).real
#
# Use the trick of Lewis here - compute the cumulative sums
# in a fashion that accounts for the parts that are off the
# edge of the template.
#
# We do this in quadrants:
# q0 q1
# q2 q3
# For the first,
# q0 is the sum over pixels1[i:,j:] - sum i,j backwards
# q1 is the sum over pixels1[i:,:j] - sum i backwards, j forwards
# q2 is the sum over pixels1[:i,j:] - sum i forwards, j backwards
# q3 is the sum over pixels1[:i,:j] - sum i,j forwards
#
# The second is done as above but reflected lr and ud
#
def get_cumsums(im, fshape):
im_si = im.shape[0]
im_sj = im.shape[1]
im_sk = im.shape[2]
im_sum = np.zeros(fshape)
im_sum[:im_si,:im_sj,:im_sk] = cumsum_quadrant(im, False, False, False)
im_sum[:im_si,:im_sj,-im_sk:] = cumsum_quadrant(im, False, False, True)
im_sum[:im_si,-im_sj:,:im_sk] = cumsum_quadrant(im, False, True, True)
im_sum[:im_si,-im_sj:,-im_sk:] = cumsum_quadrant(im, False, True, False)
im_sum[-im_si:,:im_sj,:im_sk] = cumsum_quadrant(im, True, False, True)
im_sum[-im_si:,:im_sj,-im_sk:] = cumsum_quadrant(im, True, False, False)
im_sum[-im_si:,-im_sj:,:im_sk] = cumsum_quadrant(im, True, True, True)
im_sum[-im_si:,-im_sj:,-im_sk:] = cumsum_quadrant(im, True, True, False)
#
# Divide the sum over the # of elements summed-over
#
return im_sum / unit
p1_mean = get_cumsums(pixels1, fshape)
p2_mean = get_cumsums(pixels2, fshape)
#
# Once we have the means for u,v, we can caluclate the
# variance-like parts of the equation. We have to multiply
# the mean^2 by the # of elements being summed-over
# to account for the mean being summed that many times.
#
p1sd = np.sum(pixels1**2) - p1_mean**2 * np.product(s)
p2sd = np.sum(pixels2**2) - p2_mean**2 * np.product(s)
#
# There's always chance of roundoff error for a zero value
# resulting in a negative sd, so limit the sds here
#
sd = np.sqrt(np.maximum(p1sd * p2sd, 0))
with warnings.catch_warnings():
warnings.simplefilter("ignore")
corrnorm = corr12 / sd
#
# There's not much information for points where the standard
# deviation is less than 1/100 of the maximum. We exclude these
# from consideration.
#
corrnorm[(unit < np.product(s) / 2) &
(sd < np.mean(sd) / 100)] = 0
# Also exclude possibilites with few observed pixels.
corrnorm[unit < np.product(s) / 4] = 0
return corrnorm
def cross_correlation_3d(pixels1, pixels2):
'''Align the second image with the first using max cross-correlation
returns the z,y,x offsets to add to image1's indexes to align it with
image2
Many of the ideas here are based on the paper, "Fast Normalized
Cross-Correlation" by J.P. Lewis
(http://www.idiom.com/~zilla/Papers/nvisionInterface/nip.html)
which is frequently cited when addressing this problem.
'''
s = np.maximum(pixels1.shape, pixels2.shape)
fshape = s*2
#
# Calculate the # of pixels at a particular point
#
i, j, k = np.mgrid[-s[0]:s[0], -s[1]:s[1], -s[2]:s[2] ]
unit = np.abs(i*j*k).astype(float)
unit[unit < 1] = 1 # keeps from dividing by zero in some places
#
# Normalize the pixel values around zero which does not affect the
# correlation, keeps some of the sums of multiplications from
# losing precision and precomputes t(x-u,y-v) - t_mean
#
pixels1 = np.nan_to_num(pixels1-nanmean(pixels1))
pixels2 = np.nan_to_num(pixels2-nanmean(pixels2))
#
# Lewis uses an image, f and a template t. He derives a normalized
# cross correlation, ncc(u,v) =
# sum((f(x,y)-f_mean(u,v))*(t(x-u,y-v)-t_mean),x,y) /
# sqrt(sum((f(x,y)-f_mean(u,v))**2,x,y) * (sum((t(x-u,y-v)-t_mean)**2,x,y)
#
# From here, he finds that the numerator term, f_mean(u,v)*(t...) is zero
# leaving f(x,y)*(t(x-u,y-v)-t_mean) which is a convolution of f
# by t-t_mean.
#
fp1 = rfftn(pixels1.astype('float32'), fshape, axes=(0, 1, 2))
fp2 = rfftn(pixels2.astype('float32'), fshape, axes=(0, 1, 2))
corr12 = irfftn(fp1 * fp2.conj(), axes=(0, 1, 2)).real
#
# Use the trick of Lewis here - compute the cumulative sums
# in a fashion that accounts for the parts that are off the
# edge of the template.
#
# We do this in quadrants:
# q0 q1
# q2 q3
# For the first,
# q0 is the sum over pixels1[i:,j:] - sum i,j backwards
# q1 is the sum over pixels1[i:,:j] - sum i backwards, j forwards
# q2 is the sum over pixels1[:i,j:] - sum i forwards, j backwards
# q3 is the sum over pixels1[:i,:j] - sum i,j forwards
#
# The second is done as above but reflected lr and ud
#
def get_cumsums(im, fshape):
im_si = im.shape[0]
im_sj = im.shape[1]
im_sk = im.shape[2]
im_sum = np.zeros(fshape)
im_sum[:im_si, :im_sj, :im_sk] = cumsum_quadrant(im, False, False, False)
im_sum[:im_si, :im_sj, -im_sk:] = cumsum_quadrant(im, False, False, True)
im_sum[:im_si, -im_sj:, :im_sk] = cumsum_quadrant(im, False, True, True)
im_sum[:im_si, -im_sj:, -im_sk:] = cumsum_quadrant(im, False, True, False)
im_sum[-im_si:, :im_sj, :im_sk] = cumsum_quadrant(im, True, False, True)
im_sum[-im_si:, :im_sj, -im_sk:] = cumsum_quadrant(im, True, False, False)
im_sum[-im_si:, -im_sj:, :im_sk] = cumsum_quadrant(im, True, True, True)
im_sum[-im_si:, -im_sj:, -im_sk:] = cumsum_quadrant(im, True, True, False)
#
# Divide the sum over the # of elements summed-over
#
return old_div(im_sum, unit)
p1_mean = get_cumsums(pixels1, fshape)
p2_mean = get_cumsums(pixels2, fshape)
#
# Once we have the means for u,v, we can caluclate the
# variance-like parts of the equation. We have to multiply
# the mean^2 by the # of elements being summed-over
# to account for the mean being summed that many times.
#
p1sd = np.sum(pixels1**2) - p1_mean**2 * np.product(s)
p2sd = np.sum(pixels2**2) - p2_mean**2 * np.product(s)
#
# There's always chance of roundoff error for a zero value
# resulting in a negative sd, so limit the sds here
#
sd = np.sqrt(np.maximum(p1sd * p2sd, 0))
with warnings.catch_warnings():
warnings.simplefilter("ignore")
corrnorm = old_div(corr12, sd)
#
# There's not much information for points where the standard
# deviation is less than 1/100 of the maximum. We exclude these
# from consideration.
#
corrnorm[(unit < old_div(np.product(s), 2)) &
(sd < old_div(np.mean(sd), 100))] = 0
# Also exclude possibilites with few observed pixels.
corrnorm[unit < old_div(np.product(s), 4)] = 0
return corrnorm
def fft(self, scalar):
""" Performs fft of scalar field """
self.scalar_input_test(scalar)
return fft.rfftn(scalar, threads=self.num_threads)
def fft(self, vector):
""" Performs fft of vector field """
self.vector_input_test(vector)
return fft.rfftn(vector, axes=(1, 2), threads=self.num_threads)
