def FFTev(equivalences, bandcoeff, allvec, dims):
"""Rebuild a single band from the interpolation coefficients.
Args:
equivalences: list of k-point equivalence classes in direct coordinates
bandcoeff: interpolation coefficients for the band
allvec: Cartesian coordinates of all k points
dims: upper bound on the dimensions of the k-point grid
Returns:
A 2-tuple (eb, vg) with the energies and group velocities of the
reconstructed bands.
"""
egrid = np.zeros(dims, dtype=complex)
vgrid = np.zeros([3] + list(dims), dtype=complex)
i = 0
for coeff, equiv in zip(bandcoeff, equivalences):
j = i + len(equiv)
c = coeff / len(equiv)
egrid[equiv[:, 0], equiv[:, 1], equiv[:, 2]] = c
vgrid[:, equiv[:, 0], equiv[:, 1], equiv[:, 2]] = allvec[i:j, :].T * c
i = j
vgrid *= 1j
eb = np.prod(dims) * npf.ifftn(egrid).real.flatten()
vg = [
np.prod(dims) * npf.ifftn(vgrid[0]).real.flatten(),
np.prod(dims) * npf.ifftn(vgrid[1]).real.flatten(),
np.prod(dims) * npf.ifftn(vgrid[2]).real.flatten()
]
return eb, np.array(vg)
def potential(self, q, steps):
hx = steps[0]
hy = steps[1]
hz = steps[2]
Nx = q.shape[0]
Ny = q.shape[1]
Nz = q.shape[2]
out = np.zeros((2*Nx-1, 2*Ny-1, 2*Nz-1))
out[:Nx, :Ny, :Nz] = q
K1 = self.sym_kernel(q.shape, steps)
K2 = np.zeros((2*Nx-1, 2*Ny-1, 2*Nz-1))
K2[0:Nx, 0:Ny, 0:Nz] = K1
K2[0:Nx, 0:Ny, Nz:2*Nz-1] = K2[0:Nx, 0:Ny, Nz-1:0:-1] #z-mirror
K2[0:Nx, Ny:2*Ny-1,:] = K2[0:Nx, Ny-1:0:-1, :] #y-mirror
K2[Nx:2*Nx-1, :, :] = K2[Nx-1:0:-1, :, :] #x-mirror
t0 = time()
if pyfftw_flag:
nthread = multiprocessing.cpu_count()
K2_fft = pyfftw.builders.fftn(K2, axes=None, overwrite_input=False, planner_effort='FFTW_ESTIMATE',
threads=nthread, auto_align_input=False, auto_contiguous=False, avoid_copy=True)
out_fft = pyfftw.builders.fftn(out, axes=None, overwrite_input=False, planner_effort='FFTW_ESTIMATE',
threads=nthread, auto_align_input=False, auto_contiguous=False, avoid_copy=True)
out_ifft = pyfftw.builders.ifftn(out_fft()*K2_fft(), axes=None, overwrite_input=False, planner_effort='FFTW_ESTIMATE',
threads=nthread, auto_align_input=False, auto_contiguous=False, avoid_copy=True)
out = np.real(out_ifft())
else:
out = np.real(ifftn(fftn(out)*fftn(K2)))
t1 = time()
if self.debug: print( 'fft time:', t1-t0, ' sec')
out[:Nx, :Ny, :Nz] = out[:Nx,:Ny,:Nz]/(4*pi*epsilon_0*hx*hy*hz)
return out[:Nx, :Ny, :Nz]
def cryo_conv_vol(x, kernel_f):
n = x.shape[0]
n_ker = kernel_f.shape[0]
if np.any(np.array(x.shape) != n):
raise ValueError('Volume in `x` must be cubic')
if np.any(np.array(kernel_f.shape) != n_ker):
raise ValueError('Convolution kernel in `kernel_f` must be cubic')
is_singleton = len(x.shape) == 3
shifted_kernel_f = np.fft.ifftshift(np.fft.ifftshift(np.fft.ifftshift(kernel_f, 0), 1), 2)
if is_singleton:
x = numpy_fft.fftn(x, [n_ker] * 3)
else:
x = numpy_fft.fft(x, n=n_ker, axis=0)
x = numpy_fft.fft(x, n=n_ker, axis=1)
x = numpy_fft.fft(x, n=n_ker, axis=2)
x *= shifted_kernel_f
if is_singleton:
x = numpy_fft.ifftn(x)
x = x[:n, :n, :n]
else:
x = numpy_fft.ifft(x, axis=0)
x = numpy_fft.ifft(x, axis=1)
x = numpy_fft.ifft(x, axis=2)
x = x.real
return x
def uifftn(inarray, dim=None):
"""N-dimensional unitary inverse Fourier transform.
Parameters
----------
inarray : ndarray
The array to transform.
dim : int, optional
The last axis along which to compute the transform. All
axes by default.
Returns
-------
outarray : ndarray (same shape than inarray)
The unitary inverse N-D Fourier transform of ``inarray``.
Examples
--------
>>> input = np.ones((3, 3, 3))
>>> output = uifftn(input)
>>> np.allclose(np.sum(input) / np.sqrt(input.size), output[0, 0, 0])
True
>>> output.shape
(3, 3, 3)
"""
if dim is None:
dim = inarray.ndim
outarray = ifftn(inarray, axes=range(-dim, 0))
return outarray * np.sqrt(np.prod(inarray.shape[-dim:]))
def convolve(arr1, arr2, dx=None, axes=None):
"""
Performs a centred convolution of input arrays
Parameters
----------
arr1, arr2 : `numpy.ndarray`
Arrays to be convolved. If dimensions are not equal then 1s are appended
to the lower dimensional array. Otherwise, arrays must be broadcastable.
dx : float > 0, list of float, or `None` , optional
Grid spacing of input arrays. Output is scaled by
`dx**max(arr1.ndim, arr2.ndim)`. default=`None` applies no scaling
axes : tuple of ints or `None`, optional
Choice of axes to convolve. default=`None` convolves all axes
"""
if arr2.ndim > arr1.ndim:
arr1, arr2 = arr2, arr1
if axes is None:
axes = range(arr2.ndim)
arr2 = arr2.reshape(arr2.shape + (1, ) * (arr1.ndim - arr2.ndim))
if dx is None:
dx = 1
elif isscalar(dx):
dx = dx**(len(axes) if axes is not None else arr1.ndim)
else:
dx = prod(dx)
arr1 = fftn(arr1, axes=axes)
arr2 = fftn(ifftshift(arr2), axes=axes)
out = ifftn(arr1 * arr2, axes=axes) * dx
return require(out, requirements="CA")
def fft2d(arr2d,
mode,
numpy_fft=pyfftw.interfaces.numpy_fft,
only_real=False,
batch=False):
'''
we apply an alterating +1/-1 multiplicative before we go to/from Fourier space.
Later we apply this again to the transform.
TODO: look into pyfftw.interfaces.numpy_fft.irfftn
'''
assert (arr2d.ndim == 2 and not batch) or (batch and arr2d.ndim == 3)
n1, n2 = arr2d.shape[-2:]
assert n1 == n2
arr2d = neg_pos_2d(arr2d.reshape(-1, n1, n1).copy())
if mode == 'f':
arr2d_f = numpy_fft.fftn(arr2d, axes=(-2, -1))
arr2d_f /= n1
elif mode == 'i':
arr2d_f = numpy_fft.ifftn(arr2d, axes=(-2, -1))
arr2d_f *= n1
if only_real: arr2d_f = arr2d_f.real
arr2d_f = neg_pos_2d(arr2d_f.copy())
if not batch: arr2d_f = arr2d_f.reshape(n1, n2)
return (arr2d_f)
def ifftn(a, axes=None):
if _use_fftw:
from pyfftw.interfaces.numpy_fft import ifftn
_init_pyfftw()
return ifftn(a, axes=axes, **_fft_extra_args)
else:
return np.fft.ifftn(a, axes=axes)
def easy_ifft(data, axes=None):
"""utility method that includes fft shifting"""
return ifftshift(
ifftn(
fftshift(
data, axes=axes
), axes=axes
), axes=axes)
def removePhaseRamps( img ): #... by centering the Bragg peak in the array
fimg = fftshift( fftn( fftshift( img ) ) )
intens = np.absolute( fimg )**2
maxHere = np.where( intens==intens.max() )
for n in [ 0, 1, 2 ]:
fimg = np.roll( fimg, fimg.shape[n]//2-maxHere[n], axis=n )
imgout = fftshift( ifftn( fftshift( fimg ) ) )
return imgout
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 realign_image_1d(arr, shift, angle=0):
# if both shifts are integers, do circular shift; otherwise perform Fourier shift.
if np.abs(np.array(shift) - np.round(shift)) < 0.01:
temp = np.roll(arr, int(shift))
temp = temp.astype('float32')
else:
temp = fourier_shift(fftn(arr), shift)
temp = ifftn(temp)
temp = np.abs(temp).astype('float32')
return temp
def cifftn(data, axes):
""" Centered inverse fast fourier transform, n-dimensional.
:param data: Complex input data.
:param axes: Axes along which to shift.
:return: Inverse fourier transformed data.
"""
return fft.ifftshift(fft.ifftn(fft.fftshift(data, axes=axes),
axes=axes,
norm='ortho'),
axes=axes)
def backward(self, fg):
"""Fourier backward transform a function.
Args:
fg (np.ndarray): function in G space (3D array)
Returns:
function in R space (with same grid size)
"""
assert fg.ndim == 3 and fg.shape == (self.n1, self.n2, self.n3)
fr = self.N * ifftn(fg)
return fr
def ifftw(E):
if np.any(np.array(E.shape[2:]) > 1):
assert(E.ndim-2 == self.data_shape.size and np.all(E.shape[2:] == self.data_shape))
if np.all(E.shape == ifft_vec_object.input_shape):
result_array = self.__empty_word_aligned(E.shape, dtype=np.complex128)
ifft_vec_object(E, result_array)
else:
log.debug('Inverse Fourier Transform: Array shape not standard, falling back to default interface.')
result_array = ft.ifftn(E, axes=ft_axes)
return result_array
else:
return E.copy()
def iDFT(fy, axes=None):
"""
Discrete inverse Fourier transform
Parameters
----------
fy : `numpy.ndarray`
Array defining a function evaluated on a mesh.
axes : `int` or list of `int` , optional
Specification of which axes to transform. default=`None` transforms all.
Returns
-------
fy : `numpy.ndarray`
The Fourier transform of `fx` evaluated on a mesh
"""
NDIM = fy.ndim
if axes is None:
axes = [NDIM + i for i in range(-ndim, 0)]
elif not hasattr(axes, '__iter__'):
axes = (axes, )
axes = array(axes)
axes.sort()
FT = fy.astype('complex' +
('128' if fy.real.dtype.itemsize == 8 else '64'),
copy=True)
if NDIM != 3:
# This is not typically a bottle-neck in <3D
for i in axes:
sz = [1] * NDIM
sz[axes[i]] = -1
FT *= exp(+xmin[i] * Y[i].reshape(sz) *
1j) / (dx[i] if dx[i] != 0 else 1)
else:
F = [
exp(+xmin[i] * Y[i] * 1j) / (dx[i] if dx[i] != 0 else 1)
for i in range(FT.ndim)
]
apply_phase_3D(FT, *F)
# Equivalently: FT = ifftshift(FT, axes=axes)
FT = ifftn(FT, axes=axes, overwrite_input=True)
fftshift_phase(FT) # removes need for ifftshift
return FT
def ifftw_inplace(E):
strides_not_identical = np.any(E.strides != fftw_vec_array.strides)
if strides_not_identical:
log.debug('In-place Fourier Transform: strides not identical.')
if not pyfftw.is_n_byte_aligned(E, pyfftw.simd_alignment) or strides_not_identical:
log.debug('In-place Inverse Fourier Transform: Input/Output array not %d-byte word aligned, aligning now.'
% pyfftw.simd_alignment)
E = self.__word_align(E)
if np.any(np.array(E.shape[2:]) > 1):
assert(E.ndim-2 == self.data_shape.size and np.all(E.shape[2:] == self.data_shape))
if np.all(E.shape == ifft_vec_object.input_shape):
ifft_vec_object(E, E) # E should be in a SIMD-word-aligned memory zone
else:
log.debug('Inverse Fourier Transform: Array shape not standard, falling back to default interface.')
E = ft.ifftn(E, axes=ft_axes)
return E
def fftdeconvolve(image, psf):
"""
De-convolution by directly dividing the DFT... may not be numerically
desirable for many applications. Noise could be an issue.
Use scipy.fftpack for now; will re-write for anfft later...
Taken from this post on stackoverflow.com:
http://stackoverflow.com/questions/17473917/is-there-a-equivalent-of-scipy-signal-deconvolve-for-2d-arrays
"""
if not _pyfftw:
raise NotImplementedError
image = image.astype('float')
psf = psf.astype('float')
# image_fft = fftpack.fftshift(fftpack.fftn(image))
# psf_fft = fftpack.fftshift(fftpack.fftn(psf))
image_fft = fftshift(fftn(image))
psf_fft = fftshift(fftn(psf))
kernel = fftshift(ifftn(ifftshift(image_fft / psf_fft)))
return kernel
def ftgr(fg, grid, real=False):
"""Fourier transform function fg from G space to R space.
Args:
fg (np.ndarray): G space function. shape == (grid.n1, grid.n2, grid.n3).
grid (FFTGrid): FFT grid on which fg is defined.
real (bool): if True, fg contains only iG1 >= 0 and thus has the shape
of (grid.n1 // 2 + 1, grid.n2, grid.n3), after FT R space function is real.
Returns:
R space function.
"""
if real:
assert fg.shape == (grid.n1h, grid.n2, grid.n3)
fgzyx = fg.T
return grid.N * irfftn(fgzyx, s=(grid.n1, grid.n2, grid.n3)).T
else:
assert fg.shape == (grid.n1, grid.n2, grid.n3)
return grid.N * ifftn(fg)
def fft3d(
arr3d, mode, neg_pos_3d=None, numpy_fft=pyfftw.interfaces.numpy_fft, only_real=False
):
if neg_pos_3d is None:
neg_pos_3d = make_neg_pos_3d(arr3d)
if arr3d.shape[0] % 4 != 0:
neg_pos_3d *= -1
if mode == "f":
arr3d_f = numpy_fft.fftn(neg_pos_3d * arr3d)
arr3d_f /= np.sqrt(np.prod(arr3d_f.shape))
elif mode == "i":
arr3d_f = numpy_fft.ifftn(neg_pos_3d * arr3d)
arr3d_f *= np.sqrt(np.prod(arr3d_f.shape))
if only_real:
arr3d_f = arr3d_f.real
arr3d_f *= neg_pos_3d
return arr3d_f
def potential(self, q, steps):
hx = steps[0]
hy = steps[1]
hz = steps[2]
Nx = q.shape[0]
Ny = q.shape[1]
Nz = q.shape[2]
out = np.zeros((2*Nx-1, 2*Ny-1, 2*Nz-1))
out[:Nx, :Ny, :Nz] = q
K1 = self.sym_kernel(q.shape, steps)
K2 = np.zeros((2*Nx-1, 2*Ny-1, 2*Nz-1))
K2[0:Nx, 0:Ny, 0:Nz] = K1
K2[0:Nx, 0:Ny, Nz:2*Nz-1] = K2[0:Nx, 0:Ny, Nz-1:0:-1] #z-mirror
K2[0:Nx, Ny:2*Ny-1,:] = K2[0:Nx, Ny-1:0:-1, :] #y-mirror
K2[Nx:2*Nx-1, :, :] = K2[Nx-1:0:-1, :, :] #x-mirror
t0 = time()
out = np.real(ifftn(fftn(out)*fftn(K2)))
t1 = time()
if self.debug: print( 'fft time:', t1-t0, ' sec')
out[:Nx, :Ny, :Nz] = out[:Nx,:Ny,:Nz]/(4*pi*epsilon_0*hx*hy*hz)
return out[:Nx, :Ny, :Nz]
def plan_ifft(A, n=None, axis=None, norm=None, **_):
"""
Plans an ifft for repeated use. Parameters are the same as for `pyfftw`'s `ifftn`
which are, where possible, the same as the `numpy` equivalents.
Note that some functionality is only possible when using the `pyfftw` backend.
Parameters
----------
A : `numpy.ndarray`, of dimension `d`
Array of same shape to be input for the ifft
n : iterable or `None`, `len(n) == d`, optional
The output shape of ifft (default=`None` is same as `A.shape`)
axis : `int`, iterable length `d`, or `None`, optional
The axis (or axes) to transform (default=`None` is all axes)
overwrite : `bool`, optional
Whether the input array can be overwritten during computation
(default=False)
planner : {0, 1, 2, 3}, optional
Amount of effort put into optimising Fourier transform where 0 is low
and 3 is high (default=`1`).
threads : `int`, `None`
Number of threads to use (default=`None` is all threads)
auto_align_input : `bool`, optional
If `True` then may re-align input (default=`True`)
auto_contiguous : `bool`, optional
If `True` then may re-order input (default=`True`)
avoid_copy : `bool`, optional
If `True` then may over-write initial input (default=`False`)
norm : {None, 'ortho'}, optional
Indicate whether ifft is normalised (default=`None`)
Returns
-------
plan : function
Returns the inverse Fourier transform of `B`, `plan() == ifftn(B)`
B : `numpy.ndarray`, `A.shape`
Array which should be modified inplace for ifft to be computed. If
possible, `B is A`.
"""
return lambda: ifftn(A, n, axis, norm), A
def interp(self, fr, n1, n2, n3):
"""Fourier interpolate a function to a smoother grid.
Args:
fr: function to be interpolated
n1, n2, n3 (int): new grid size
Returns:
interpolated function (3D array of size n1 by n2 by n3)
"""
assert fr.ndim == 3 and fr.shape == (self.n1, self.n2, self.n3)
assert n1 <= self.n1 and n2 <= self.n2 and n3 <= n3
if (n1, n2, n3) == (self.n1, self.n2, self.n3):
return fr
fg = fftn(fr)
sfg = fftshift(fg)
newsfg = sfg[(self.n1-n1-1)//2+1:(self.n1-n1-1)//2+1+n1,
(self.n2-n2-1)//2+1:(self.n2-n2-1)//2+1+n2,
(self.n3-n3-1)//2+1:(self.n3-n3-1)//2+1+n3,
]
newfg = ifftshift(newsfg)
newfr = (float(n1*n2*n3)/float(self.N)) * ifftn(newfg)
return newfr
def FFTc(equivalences, bandcoeff, allvec, dims):
"""Rebuild the curvature of a band from the interpolation coefficients.
Args:
equivalences: list of k-point equivalence classes in direct coordinates
bandcoeff: interpolation coefficients for the band
allvec: Cartesian coordinates of all k points
dims: upper bound on the dimensions of the k-point grid
Returns:
An array used to rebuild the curvatures in fft_worker.
"""
cgrid = np.zeros([6] + list(dims), dtype=complex)
i = 0
ii1, ii2 = np.triu_indices(3)
for coeff, equiv in zip(bandcoeff, equivalences):
j = i + len(equiv)
c = coeff / len(equiv)
eq = allvec[i:j, :]
cgrid[:, equiv[:, 0], equiv[:, 1],
equiv[:, 2]] = -(eq[:, ii1] * eq[:, ii2]).T * c
i = j
cg = [npf.ifftn(cgrid[i]).real.flatten() for i in range(6)]
return np.prod(dims) * np.array(cg)
def plot_fixed_t32(self,
t32_time,
*,
part='real',
signal='rephasing',
ft=True,
savefig=True,
omega_0=1):
""""""
self.load_eigen_params()
t32_index, t32_time = self.get_closest_index_and_value(
t32_time, self.t32_array)
dt21 = self.t21_array[1] - self.t21_array[0]
dt43 = self.t43_array[1] - self.t43_array[0]
if signal == 'rephasing':
sig = self.rephasing_signal[:, t32_index, :]
elif signal == 'nonrephasing':
sig = self.nonrephasing_signal[:, t32_index, :]
if ft:
w21 = fftshift(fftfreq(self.t21_array.size, d=dt21)) * 2 * np.pi
w21 += self.ground_to_excited_transition + self.center
w21 *= omega_0
w43 = fftshift(fftfreq(self.t43_array.size, d=dt43)) * 2 * np.pi
w43 += self.ground_to_excited_transition + self.center
w43 *= omega_0
X, Y = np.meshgrid(w21, w43, indexing='ij')
if signal == 'nonrephasing':
ifft_t21_norm = self.t21_array.size * dt21
ifft_t43_norm = self.t43_array.size * dt43
sig = fftshift(ifftn(sig, axes=(0, 1)),
axes=(0, 1)) * ifft_t21_norm * ifft_t43_norm
elif signal == 'rephasing':
fft_t21_norm = dt21
ifft_t43_norm = self.t43_array.size * dt43
sig = fftshift(ifft(sig, axis=1), axes=(1)) * ifft_t43_norm
sig = fftshift(fft(sig, axis=0), axes=(0)) * fft_t21_norm
if omega_0 == 1:
xlab = '$\omega_{21}$ ($\omega_0$)'
ylab = '$\omega_{43}$ ($\omega_0$)'
else:
xlab = '$\omega_{21}$ (cm$^{-1}$)'
ylab = '$\omega_{43}$ (cm$^{-1}$)'
else:
X = self.T21[:, 0, :]
Y = self.T43[:, 0, :]
xlab = r'$t_{21}$ ($\omega_0^{-1}$)'
ylab = r'$t_{43}$ ($\omega_0^{-1}$)'
if part == 'real':
sig = sig.real
if part == 'imag':
sig = sig.imag
plt.figure()
plt.contour(X, Y, sig, 12, colors='k')
plt.contourf(X, Y, sig, 12)
plt.title(part + ' ' + signal + r' at $t_{32}$' +
' = {}'.format(t32_time))
plt.xlabel(xlab)
plt.ylabel(ylab)
plt.xlim([17000, 21000])
plt.ylim([17000, 21000])
plt.colorbar()
if savefig:
fig_name = os.path.join(
self.base_path,
part + '_' + signal + '_t_32_{}.png'.format(t32_time))
plt.savefig(fig_name)
plt.show()
def find_offsets(frames2reg_fft, template_fft):
"""
Computes the convolution of the template with the frames by taking
advantage of their FFTs for faster computation that an ordinary
convolution ( O(N*lg(N)) vs O(N^2) )
< http://ccrma.stanford.edu/~jos/ReviewFourier/FFT_Convolution_vs_Direct.html >.
Once computed the maximum of the convolution is found to determine the
best overlap of each frame with the template, which provides the needed
offset. Some corrections are performed to make reasonable offsets.
Notes:
Adapted from code provided by Wenzhi Sun with speed improvements
provided by Uri Dubin.
Args:
frames2reg(numpy.ndarray): image stack to register (time
is the first dimension uses
C-order tyx or tzyx).
template_fft(numpy.ndarray): what to register the image
stack against (single frame
using C-order yx or zyx).
Returns:
(numpy.ndarray): an array containing the
translations to apply to each
frame.
Examples:
>>> a = numpy.zeros((5, 3, 4)); a[:,0] = 1; a[2,0] = 0; a[2,2] = 1
>>> a
array([[[ 1., 1., 1., 1.],
[ 0., 0., 0., 0.],
[ 0., 0., 0., 0.]],
<BLANKLINE>
[[ 1., 1., 1., 1.],
[ 0., 0., 0., 0.],
[ 0., 0., 0., 0.]],
<BLANKLINE>
[[ 0., 0., 0., 0.],
[ 0., 0., 0., 0.],
[ 1., 1., 1., 1.]],
<BLANKLINE>
[[ 1., 1., 1., 1.],
[ 0., 0., 0., 0.],
[ 0., 0., 0., 0.]],
<BLANKLINE>
[[ 1., 1., 1., 1.],
[ 0., 0., 0., 0.],
[ 0., 0., 0., 0.]]])
>>> af = numpy.fft.fftn(a, axes=range(1, a.ndim)); af
array([[[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ]],
<BLANKLINE>
[[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ]],
<BLANKLINE>
[[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[-2.+3.46410162j, 0.+0.j , 0.+0.j , 0.+0.j ],
[-2.-3.46410162j, 0.+0.j , 0.+0.j , 0.+0.j ]],
<BLANKLINE>
[[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ]],
<BLANKLINE>
[[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ],
[ 4.+0.j , 0.+0.j , 0.+0.j , 0.+0.j ]]])
>>> tf = numpy.fft.fftn(a.mean(axis=0)); tf
array([[ 4.0+0.j , 0.0+0.j , 0.0+0.j , 0.0+0.j ],
[ 2.8+0.69282032j, 0.0+0.j , 0.0+0.j , 0.0+0.j ],
[ 2.8-0.69282032j, 0.0+0.j , 0.0+0.j , 0.0+0.j ]])
>>> find_offsets(af, tf)
array([[ 0, 0],
[ 0, 0],
[-2, 0],
[ 0, 0],
[ 0, 0]])
"""
# If there is only one frame, add a singleton axis to indicate this.
frames2reg_fft_added_singleton = (frames2reg_fft.ndim == template_fft.ndim)
if frames2reg_fft_added_singleton:
frames2reg_fft = frames2reg_fft[None]
# Compute the product of the two FFTs (i.e. the convolution of the regular
# versions).
frames2reg_template_conv_fft = frames2reg_fft * template_fft.conj()[None]
# Find the FFT inverse (over all spatial dimensions) to return to the
# convolution.
frames2reg_template_conv = fft.ifftn(
frames2reg_template_conv_fft, axes=range(1, frames2reg_fft.ndim))
# Find where the convolution is maximal. Will have the most things in
# common between the template and frames.
frames2reg_template_conv_max, frames2reg_template_conv_max_indices = xnumpy.max_abs(
frames2reg_template_conv,
axis=range(1, frames2reg_fft.ndim),
return_indices=True
)
# First index is just the frame, which will be in sequential order. We
# don't need this so we drop it.
frames2reg_template_conv_max_indices = frames2reg_template_conv_max_indices[1:]
# Convert indices into an array for easy manipulation.
frames2reg_template_conv_max_indices = numpy.array(
frames2reg_template_conv_max_indices
).T.copy()
# Shift will have to be in the opposite direction to bring everything to
# the center.
numpy.negative(
frames2reg_template_conv_max_indices,
out=frames2reg_template_conv_max_indices
)
return(frames2reg_template_conv_max_indices)
def _ModProject(self):
self._cImage = ifftn(self._modulus *
np.exp(1j * np.angle(fftn(self._cImage))))
return
def _register_translation(src_image, target_image, upsample_factor=1,
space="real"):
"""
*****************************************
From skimage.feature.register_translation
*****************************************
Efficient subpixel image translation registration by cross-correlation.
This code gives the same precision as the FFT upsampled cross-correlation
in a fraction of the computation time and with reduced memory requirements.
It obtains an initial estimate of the cross-correlation peak by an FFT and
then refines the shift estimation by upsampling the DFT only in a small
neighborhood of that estimate by means of a matrix-multiply DFT.
Parameters
----------
src_image : ndarray
Reference image.
target_image : ndarray
Image to register. Must be same dimensionality as ``src_image``.
upsample_factor : int, optional
Upsampling factor. Images will be registered to within
``1 / upsample_factor`` of a pixel. For example
``upsample_factor == 20`` means the images will be registered
within 1/20th of a pixel. Default: 1 (no upsampling).
space : string, one of "real" or "fourier"
Defines how the algorithm interprets input data. "real" means data
will be FFT'd to compute the correlation, while "fourier" data will
bypass FFT of input data. Case insensitive. Default: "real".
Returns
-------
shifts : ndarray
Shift vector (in pixels) required to register ``target_image`` with
``src_image``. Axis ordering is consistent with numpy (e.g. Z, Y, X)
error : float
Translation invariant normalized RMS error between ``src_image`` and
``target_image``.
phasediff : float
Global phase difference between the two images (should be
zero if images are non-negative).
References
----------
.. [1] Manuel Guizar-Sicairos, Samuel T. Thurman, and James R. Fienup,
"Efficient subpixel image registration algorithms,"
Optics Letters 33, 156-158 (2008).
"""
# images must be the same shape
if src_image.shape != target_image.shape:
raise ValueError("Error: images must be same size for "
"register_translation")
# only 2D data makes sense right now
if src_image.ndim != 2 and upsample_factor > 1:
raise NotImplementedError("Error: register_translation only supports "
"subpixel registration for 2D images")
# assume complex data is already in Fourier space
if space.lower() == 'fourier':
src_freq = src_image
target_freq = target_image
# real data needs to be fft'd.
elif space.lower() == 'real':
src_image = np.array(src_image, dtype=np.complex128, copy=False)
target_image = np.array(target_image, dtype=np.complex128, copy=False)
src_freq = fftn(src_image)
target_freq = fftn(target_image)
else:
raise ValueError("Error: register_translation only knows the \"real\" "
"and \"fourier\" values for the ``space`` argument.")
# Whole-pixel shift - Compute cross-correlation by an IFFT
shape = src_freq.shape
image_product = src_freq * target_freq.conj()
cross_correlation = ifftn(image_product)
# Locate maximum
maxima = np.unravel_index(np.argmax(np.abs(cross_correlation)),
cross_correlation.shape)
midpoints = np.array([np.fix(axis_size / 2) for axis_size in shape])
shifts = np.array(maxima, dtype=np.float64)
shifts[shifts > midpoints] -= np.array(shape)[shifts > midpoints]
if upsample_factor == 1:
src_amp = np.sum(np.abs(src_freq) ** 2) / src_freq.size
target_amp = np.sum(np.abs(target_freq) ** 2) / target_freq.size
# CCmax = cross_correlation.max()
# If upsampling > 1, then refine estimate with matrix multiply DFT
else:
# Initial shift estimate in upsampled grid
shifts = np.round(shifts * upsample_factor) / upsample_factor
upsampled_region_size = np.ceil(upsample_factor * 1.5)
# Center of output array at dftshift + 1
dftshift = np.fix(upsampled_region_size / 2.0)
upsample_factor = np.array(upsample_factor, dtype=np.float64)
normalization = (src_freq.size * upsample_factor ** 2)
# Matrix multiply DFT around the current shift estimate
sample_region_offset = dftshift - shifts * upsample_factor
cross_correlation = _upsampled_dft(image_product.conj(),
upsampled_region_size,
upsample_factor,
sample_region_offset).conj()
cross_correlation /= normalization
# Locate maximum and map back to original pixel grid
maxima = np.array(np.unravel_index(
np.argmax(np.abs(cross_correlation)),
cross_correlation.shape),
dtype=np.float64)
maxima -= dftshift
shifts = shifts + maxima / upsample_factor
# CCmax = cross_correlation.max()
src_amp = _upsampled_dft(src_freq * src_freq.conj(),
1, upsample_factor)[0, 0]
src_amp /= normalization
target_amp = _upsampled_dft(target_freq * target_freq.conj(),
1, upsample_factor)[0, 0]
target_amp /= normalization
# If its only one row or column the shift along that dimension has no
# effect. We set to zero.
for dim in range(src_freq.ndim):
if shape[dim] == 1:
shifts[dim] = 0
return shifts
def ifft(*args, **kwargs):
return fftw.ifftn(*args, **kwargs, threads=2)
def __init__(self, nb_dims, data_shape, sample_pitch=None):
self.__nb_dims = nb_dims
self.__data_shape = data_shape
if sample_pitch is None:
sample_pitch = np.ones(self.data_shape.size)
self.__sample_pitch = sample_pitch
self.__cutoff = len(self.matrix_shape)
self.__total_dims = len(self.matrix_shape) + len(self.data_shape)
self.__eye = np.eye(nb_dims).reshape((nb_dims, nb_dims, *np.ones(len(self.data_shape), dtype=int)))
ft_axes = range(self.__cutoff, self.__total_dims) # Don't Fourier transform the matrix dimensions
if 'pyfftw' not in globals():
if 'pyfftw.interfaces.numpy_fft' not in sys.modules:
log.debug('Module pyfftw not imported, using stock FFT.')
else:
log.debug('Module pyfftw not imported, but module pyfftw.interfaces.numpy_fft is.'
' Using the FFTW-NumPy interface.')
self.__ft = lambda E: ft.fftn(E, axes=ft_axes)
self.__ift = lambda E: ft.ifftn(E, axes=ft_axes)
self.__empty_word_aligned = lambda shape, dtype: np.empty(shape=shape, dtype=dtype)
self.__zeros_word_aligned = lambda shape, dtype: np.zeros(shape=shape, dtype=dtype)
self.__word_align = lambda a: a
else:
log.debug('Module pyfftw imported, using FFTW instead of stock FFT.')
ftflags = ('FFTW_DESTROY_INPUT', 'FFTW_ESTIMATE', )
# ftflags = ('FFTW_DESTROY_INPUT', 'FFTW_PATIENT', )
# ftflags = ('FFTW_DESTROY_INPUT', 'FFTW_MEASURE', )
# ftflags = ('FFTW_DESTROY_INPUT', 'FFTW_EXHAUSTIVE', ) # very slow, little gain in general
nb_threads = multiprocessing.cpu_count()
log.debug("Number of threads available for FFTW: %d" % nb_threads)
self.__empty_word_aligned = \
lambda shape, dtype: pyfftw.empty_aligned(shape=shape, dtype=dtype, n=pyfftw.simd_alignment)
self.__zeros_word_aligned = \
lambda shape, dtype: pyfftw.zeros_aligned(shape=shape, dtype=dtype, n=pyfftw.simd_alignment)
self.__word_align = lambda a: utils.word_align(a, word_length=pyfftw.simd_alignment)
log.debug('Allocating FFTW''s operating memory.')
fftw_vec_array = self.__empty_word_aligned([self.nb_dims, 1, *self.data_shape], dtype=np.complex128)
log.debug('Initializing FFTW''s forward Fourier transform.')
fft_vec_object = pyfftw.FFTW(fftw_vec_array, fftw_vec_array, axes=ft_axes, flags=ftflags, direction='FFTW_FORWARD',
planning_timelimit=None, threads=nb_threads)
log.debug('Initializing FFTW''s backward Fourier transform.')
ifft_vec_object = pyfftw.FFTW(fftw_vec_array, fftw_vec_array, axes=ft_axes, flags=ftflags, direction='FFTW_BACKWARD',
planning_timelimit=None, threads=nb_threads)
log.debug('FFTW''s wisdoms generated.')
# Redefine the default method
def fftw(E):
if np.any(np.array(E.shape[2:]) > 1):
assert(E.ndim-2 == self.data_shape.size and np.all(E.shape[2:] == self.data_shape))
if np.all(E.shape == fft_vec_object.input_shape):
result_array = self.__empty_word_aligned(E.shape, dtype=np.complex128)
fft_vec_object(E, result_array)
else:
log.debug('Fourier Transform: Array shape not standard, falling back to default interface.')
result_array = ft.fftn(E, axes=ft_axes)
return result_array
else:
return E.copy()
def fftw_inplace(E):
if np.any(np.array(E.shape[2:]) > 1):
strides_not_identical = np.any(E.strides != fftw_vec_array.strides)
if strides_not_identical:
log.debug('In-place Fourier Transform: strides not identical.')
E = self.__word_align(E.copy())
if not pyfftw.is_n_byte_aligned(E, pyfftw.simd_alignment):
log.debug('In-place Fourier Transform: Input/Output array not %d-byte word aligned, aligning now.'
% pyfftw.simd_alignment)
E = self.__word_align(E)
assert(E.ndim-2 == self.data_shape.size and np.all(E.shape[2:] == self.data_shape))
if np.all(E.shape == fft_vec_object.input_shape) \
and pyfftw.is_n_byte_aligned(E, pyfftw.simd_alignment):
fft_vec_object(E, E) # E should be in SIMD-word-aligned memory zone
else:
log.debug('Fourier Transform: Array shape not standard, falling back to default interface.')
E = ft.fftn(E, axes=ft_axes)
return E
# Redefine the default method
def ifftw(E):
if np.any(np.array(E.shape[2:]) > 1):
assert(E.ndim-2 == self.data_shape.size and np.all(E.shape[2:] == self.data_shape))
if np.all(E.shape == ifft_vec_object.input_shape):
result_array = self.__empty_word_aligned(E.shape, dtype=np.complex128)
ifft_vec_object(E, result_array)
else:
log.debug('Inverse Fourier Transform: Array shape not standard, falling back to default interface.')
result_array = ft.ifftn(E, axes=ft_axes)
return result_array
else:
return E.copy()
def ifftw_inplace(E):
strides_not_identical = np.any(E.strides != fftw_vec_array.strides)
if strides_not_identical:
log.debug('In-place Fourier Transform: strides not identical.')
if not pyfftw.is_n_byte_aligned(E, pyfftw.simd_alignment) or strides_not_identical:
log.debug('In-place Inverse Fourier Transform: Input/Output array not %d-byte word aligned, aligning now.'
% pyfftw.simd_alignment)
E = self.__word_align(E)
if np.any(np.array(E.shape[2:]) > 1):
assert(E.ndim-2 == self.data_shape.size and np.all(E.shape[2:] == self.data_shape))
if np.all(E.shape == ifft_vec_object.input_shape):
ifft_vec_object(E, E) # E should be in a SIMD-word-aligned memory zone
else:
log.debug('Inverse Fourier Transform: Array shape not standard, falling back to default interface.')
E = ft.ifftn(E, axes=ft_axes)
return E
self.__ft = fftw
self.__ift = ifftw
def ifftnc(x, axes):
tmp = fft.fftshift(x, axes=axes)
tmp = fft.ifftn(tmp, axes=axes)
return fft.ifftshift(tmp, axes=axes)
def ift(psi):
"""Go from spectral space to physical space."""
return ifftn(psi, axes=(0,1))
def precondFunc(x):
return np.real(
np_fft.ifftn(self.invGsq *
np_fft.fftn(np.reshape(x, S))).flatten())
