def run(self):
max_shape = self._find_max_shape()
# compute FT, assuming they are the same size
fft1 = cp.asarray(self.image1, dtype=cp.complex64)
fft2 = cp.asarray(self.image2, dtype=cp.complex64)
plan = get_fft_plan(fft1, value_type="C2C")
fft1 = fftn(fft1, overwrite_x=True, plan=plan)
fft2 = fftn(fft2, overwrite_x=True, plan=plan)
print(f"shape: {fft1.shape}, dtype: {fft1.dtype}")
@cp.fuse
def normalize(fft_image):
re, im = cp.real(fft_image), cp.imag(fft_image)
length = cp.sqrt(re * re + im * im)
return fft_image / length
fft1 = normalize(fft1)
fft2 = cp.conj(normalize(fft2))
# phase correlation spectrum
pcm = fft1 * fft2
pcm = ifftn(pcm, overwrite_x=True, plan=plan)
pcm = cp.real(pcm)
from skimage.morphology import disk
from skimage.filters import median
pcm = cp.asnumpy(pcm)
pcm = median(pcm, disk(3))
pcm = cp.asarray(pcm)
peak_list = self._extract_correlation_peaks(pcm)
def test_fftn_multiple_plan_error(self, dtype):
import cupy
import cupyx.scipy.fftpack as fftpack
x = testing.shaped_random(self.shape, cupy, dtype)
# hack: avoid testing the cases when getting a cuFFT plan is impossible
if _default_fft_func(x, s=self.s, axes=self.axes) is not _fftn:
return
plan = fftpack.get_fft_plan(x, shape=self.s, axes=self.axes)
with pytest.raises(RuntimeError) as ex, plan:
fftpack.fftn(x, shape=self.s, axes=self.axes, plan=plan)
assert 'Use the cuFFT plan either as' in str(ex.value)
def wiener_deconvolution(image, psf, snr=30, add_pad=0):
""" A GPU accelerated implementation of a linear Wiener filter. Some effort is made
to allow processing even relatively large images, but some kind of block-based processing
(as in the RL implementation) may be required in some cases."""
assert isinstance(image, Image)
assert isinstance(psf, Image)
image_s = Image(image.copy(), image.spacing)
orig_shape = image.shape
if image.ndim != psf.ndim:
raise ValueError("Image and psf dimensions do not match")
if psf.spacing != image.spacing:
psf = imops.zoom_to_spacing(psf, image.spacing)
if add_pad != 0:
new_shape = list(i + 2 * add_pad for i in image_s.shape)
image_s = imops.zero_pad_to_shape(image_s, new_shape)
if psf.shape != image_s.shape:
psf = imops.zero_pad_to_shape(psf, image_s.shape)
psf /= psf.max()
psf = fftshift(psf)
psf_dev = cp.asarray(psf.astype(np.complex64))
with get_fft_plan(psf_dev):
psf_dev = fftn(psf_dev, overwrite_x=True)
below = cp.asnumpy(psf_dev)
psf_abs = cp.abs(psf_dev)**2
psf_abs /= (psf_abs + snr)
above = cp.asnumpy(psf_abs)
psf_abs = None
psf_dev = None
image_dev = cp.asarray(image_s.astype(np.complex64))
with get_fft_plan(image_dev):
image_dev = fftn(image_dev, overwrite_x=True)
wiener_dev = cp.asarray(arrayops.safe_divide(above, below))
image_dev *= wiener_dev
result = cp.asnumpy(cp.abs(ifftn(image_dev, overwrite_x=True)).real)
result = Image(result, image.spacing)
return imops.remove_zero_padding(result, orig_shape)
def transform(self, data):
if self.f_contiguous:
_dat = data.T
else:
_dat = data
_dat[:] = cufftp.fftn(_dat, axes=self._ax, plan=self._fftplan)[:]
"""The transform is done inplace"""
def __get_fourier_psfs(self):
"""
Pre-calculates the PSFs during image fusion process.
"""
print("Pre-calculating PSFs")
padded_block_size = tuple(self.block_size + 2 * self.options.block_pad)
memmap_shape = (self.n_views, ) + padded_block_size
if self.options.disable_fft_psf_memmap:
self.psfs_fft = numpy.zeros(memmap_shape, dtype=numpy.complex64)
self.adj_psfs_fft = numpy.zeros(memmap_shape,
dtype=numpy.complex64)
else:
psfs_fft_f = os.path.join(self.memmap_directory, "psf_fft_f.dat")
self.psfs_fft = numpy.memmap(psfs_fft_f,
dtype='complex64',
mode='w+',
shape=memmap_shape)
adj_psfs_fft_f = os.path.join(self.memmap_directory,
"adj_psf_fft_f.dat")
self.adj_psfs_fft = numpy.memmap(adj_psfs_fft_f,
dtype='complex64',
mode='w+',
shape=memmap_shape)
for idx in range(self.n_views):
self.psfs_fft[idx] = ops_array.expand_to_shape(
self.psfs[idx], padded_block_size).astype(numpy.complex64)
self.adj_psfs_fft[idx] = ops_array.expand_to_shape(
self.adj_psfs[idx], padded_block_size).astype(numpy.complex64)
self.psfs_fft[idx] = numpy.fft.fftshift(self.psfs_fft[idx])
self.adj_psfs_fft[idx] = numpy.fft.fftshift(self.adj_psfs_fft[idx])
self.psfs_fft[idx] = cp.asnumpy(
fftpack.fftn(cp.asarray(self.psfs_fft[idx]),
plan=self._fft_plan))
self.adj_psfs_fft[idx] = cp.asnumpy(
fftpack.fftn(cp.asarray(self.adj_psfs_fft[idx]),
plan=self._fft_plan))
def _fft_convolve(self, h_data, h_kernel):
"""
Calculate a convolution on GPU, using FFTs.
:param h_data: a Numpy array with the data to convolve
:param h_kernel: a Numpy array with the convolution kernel. The kernel
should already be in Fourier domain (To avoid repeating the transform at
every iteration.)
"""
#todo: See whether to add back streams. I removed them on Cupy refactor.
d_data = cp.asarray(h_data)
d_data = fftpack.fftn(d_data, overwrite_x=True, plan=self._fft_plan)
d_kernel = cp.asarray(h_kernel)
d_data *= d_kernel
d_data = fftpack.ifftn(d_data, overwrite_x=True, plan=self._fft_plan)
return cp.asnumpy(d_data)
def phasecorr_gpu(X, cfRefImg, lcorr):
''' not being used - no speed up - may be faster with cuda.jit'''
nimg, Ly, Lx = X.shape
ly, lx = cfRefImg.shape[-2:]
lyhalf = int(np.floor(ly / 2))
lxhalf = int(np.floor(lx / 2))
# put on GPU
ref_gpu = cp.asarray(cfRefImg)
x_gpu = cp.asarray(X)
# phasecorrelation
x_gpu = fftn(x_gpu, axes=(1, 2),
overwrite_x=True) * np.sqrt(Ly - 1) * np.sqrt(Lx - 1)
for t in range(x_gpu.shape[0]):
tmp = x_gpu[t, :, :]
tmp = cp.multiply(tmp, ref_gpu)
tmp = cp.divide(tmp, cp.absolute(tmp) + 1e-5)
x_gpu[t, :, :] = tmp
x_gpu = ifftn(x_gpu, axes=(1, 2),
overwrite_x=True) * np.sqrt(Ly - 1) * np.sqrt(Lx - 1)
x_gpu = cp.fft.fftshift(cp.real(x_gpu), axes=(1, 2))
# get max index
x_gpu = x_gpu[cp.ix_(np.arange(0, nimg, 1, int),
np.arange(lyhalf - lcorr, lyhalf + lcorr + 1, 1, int),
np.arange(lxhalf - lcorr, lxhalf + lcorr + 1, 1,
int))]
ix = cp.argmax(cp.reshape(x_gpu, (nimg, -1)), axis=1)
cmax = x_gpu[np.arange(0, nimg, 1, int), ix]
ymax, xmax = cp.unravel_index(ix, (2 * lcorr + 1, 2 * lcorr + 1))
cmax = cp.asnumpy(cmax).flatten()
ymax = cp.asnumpy(ymax)
xmax = cp.asnumpy(xmax)
ymax, xmax = ymax - lcorr, xmax - lcorr
return ymax, xmax, cmax
def soft_mask(v, voxel_size, num_subunit_residues,
helical_repeat_distance=None, repeats_to_include=0,
filter_resolution=20, expansion_factor=1.2,
expansion_radius=0, print_progress=True, return_mask=False):
full_expansion_radius = expansion_radius + filter_resolution/2
# avg AA mol wt. in g/mol, density in g/cm3
avg_aa_molwt = 110
protein_density = 1.4
# 2
print helical_repeat_distance
v_thresh = np.zeros(v.shape)
v_thresh[:] = v[:]
sz = np.array(v.shape).astype(int)
total_molwt = num_subunit_residues*avg_aa_molwt/6.023e23
if helical_repeat_distance != None:
total_molwt = total_molwt * sz[2]*voxel_size / helical_repeat_distance
total_vol = np.prod(sz) * voxel_size**3 # vol in A3
mol_vol = total_molwt/protein_density / (1.0e-24) # vol in A3
mol_vol_frac = mol_vol/total_vol
target_vol_frac = mol_vol_frac*expansion_factor
thresh = find_binary_threshold(v_thresh, target_vol_frac)
true_frac = (0.0 + np.sum(v_thresh >= thresh)) / v_thresh.size
if repeats_to_include != 0:
zdim = np.round(repeats_to_include * helical_repeat_distance/voxel_size)
else:
zdim = sz[2]
if zdim > sz[2] - 4*np.ceil(filter_resolution/voxel_size):
zdim = sz[2] - 4*np.ceil(filter_resolution/voxel_size)
zdim = zdim.astype(int)
v_thresh[:,:,0:np.floor(sz[2]/2).astype(int) - np.floor(zdim/2).astype(int)] = 0
v_thresh[:,:,np.floor(sz[2]/2).astype(int) - np.floor(zdim/2).astype(int) + 1 + zdim - 1:] = 0
v_thresh[v_thresh < thresh] = 0
if print_progress:
print 'Target volume fraction: {}'.format(target_vol_frac)
print 'Achieved volume fraction: {}'.format(true_frac)
print 'Designated threshold: {}'.format(thresh)
progress_bar = tqdm(total=5)
v_thresh = fftpack.fftn(v_thresh)
progress_bar.update(1)
# 3
cosmask_filter = np.fft.fftshift(spherical_cosmask(sz, 0, np.ceil(filter_resolution/voxel_size)))
cosmask_filter = fftpack.fftn(cosmask_filter) / np.sum(cosmask_filter)
progress_bar.update(1)
v_thresh = v_thresh * cosmask_filter
v_thresh = fftpack.ifftn(v_thresh)
progress_bar.update(1)
v_thresh = np.real(v_thresh)
v_thresh[np.abs(v_thresh) < 10*np.finfo(type(v_thresh.ravel()[0])).eps] = 0
v_thresh[v_thresh != 0] = 1
# The extent of blurring is equal to the diameter of the cosmask sphere;
# if we want this to equal the expected falloff for filter_resolution,
# we therefore need to divide filter_res by 4 to get the
# desired radius for spherical_cosmask.
v_thresh = fftpack.fftn(v_thresh)
progress_bar.update(1)
v_thresh = v_thresh * cosmask_filter
v_thresh = fftpack.ifftn(v_thresh)
progress_bar.update(1)
v_thresh = np.real(v_thresh)
v_thresh[np.abs(v_thresh) < 10*np.finfo(type(v_thresh.ravel()[0])).eps] = 0
if return_mask:
v[:,:,:] = v_thresh
else:
v *= v_thresh
return v_thresh
def fftconvolve(in1, in2, mode="full", axes=None):
"""Convolve two N-dimensional arrays using FFT.
Convolve `in1` and `in2` using the fast Fourier transform method, with
the output size determined by the `mode` argument.
This is generally much faster than `convolve` for large arrays (n > ~500),
but can be slower when only a few output values are needed, and can only
output float arrays (int or object array inputs will be cast to float).
As of v0.19, `convolve` automatically chooses this method or the direct
method based on an estimation of which is faster.
Parameters
----------
in1 : array_like
First input.
in2 : array_like
Second input. Should have the same number of dimensions as `in1`.
mode : str {'full', 'valid', 'same'}, optional
A string indicating the size of the output:
``full``
The output is the full discrete linear convolution
of the inputs. (Default)
``valid``
The output consists only of those elements that do not
rely on the zero-padding. In 'valid' mode, either `in1` or `in2`
must be at least as large as the other in every dimension.
``same``
The output is the same size as `in1`, centered
with respect to the 'full' output.
axis : tuple, optional
axes : int or array_like of ints or None, optional
Axes over which to compute the convolution.
The default is over all axes.
Returns
-------
out : array
An N-dimensional array containing a subset of the discrete linear
convolution of `in1` with `in2`.
Examples
--------
Autocorrelation of white noise is an impulse.
>>> import cusignal
>>> import cupy as cp
>>> import numpy as np
>>> sig = cp.random.randn(1000)
>>> autocorr = cusignal.fftconvolve(sig, sig[::-1], mode='full')
>>> import matplotlib.pyplot as plt
>>> fig, (ax_orig, ax_mag) = plt.subplots(2, 1)
>>> ax_orig.plot(cp.asnumpy(sig))
>>> ax_orig.set_title('White noise')
>>> ax_mag.plot(np.arange(-len(sig)+1,len(sig)), autocorr)
>>> ax_mag.set_title('Autocorrelation')
>>> fig.tight_layout()
>>> fig.show()
Gaussian blur implemented using FFT convolution. Notice the dark borders
around the image, due to the zero-padding beyond its boundaries.
The `convolve2d` function allows for other types of image boundaries,
but is far slower.
>>> from scipy import misc
>>> face = misc.face(gray=True)
>>> kernel = cp.outer(cusignal.gaussian(70, 8), cusignal.gaussian(70, 8))
>>> blurred = cusignal.fftconvolve(face, kernel, mode='same')
>>> fig, (ax_orig, ax_kernel, ax_blurred) = plt.subplots(3, 1,
... figsize=(6, 15))
>>> ax_orig.imshow(face, cmap='gray')
>>> ax_orig.set_title('Original')
>>> ax_orig.set_axis_off()
>>> ax_kernel.imshow(cp.asnumpy(kernel), cmap='gray')
>>> ax_kernel.set_title('Gaussian kernel')
>>> ax_kernel.set_axis_off()
>>> ax_blurred.imshow(cp.asnumpy(blurred), cmap='gray')
>>> ax_blurred.set_title('Blurred')
>>> ax_blurred.set_axis_off()
>>> fig.show()
"""
in1 = cp.ascontiguousarray(in1)
in2 = cp.ascontiguousarray(in2)
noaxes = axes is None
if in1.ndim == in2.ndim == 0: # scalar inputs
return in1 * in2
elif 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 cp.array([])
_, axes = _init_nd_shape_and_axes_sorted(in1, shape=None, axes=axes)
# axes needs to be numpy type for proper execution in FFT
axes = cp.asnumpy(axes)
if not noaxes and not axes.size:
raise ValueError("when provided, axes cannot be empty")
if noaxes:
other_axes = np.array([], dtype=cp.intc)
else:
other_axes = np.setdiff1d(np.arange(in1.ndim), axes)
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
if not np.all((s1[other_axes] == s2[other_axes])
| (s1[other_axes] == 1)
| (s2[other_axes] == 1)):
raise ValueError("incompatible shapes for in1 and in2:"
" {0} and {1}".format(in1.shape, in2.shape))
complex_result = np.issubdtype(in1.dtype,
np.complexfloating) or np.issubdtype(
in2.dtype, cp.complexfloating)
shape = np.maximum(s1, s2)
shape[axes] = s1[axes] + s2[axes] - 1
# Check that input sizes are compatible with 'valid' mode
if _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 = [next_fast_len(d) for d in shape[axes]]
fslice = tuple([slice(sz) for sz in shape])
if not complex_result:
if (str(fshape), str(axes), "R2C") in _cupy_fft_cache:
rplan = _cupy_fft_cache[(str(fshape), str(axes), "R2C")]
else:
rplan = _cupy_fft_cache[(str(fshape), str(axes),
"R2C")] = fftpack.get_fft_plan(
in1,
fshape,
axes=axes,
value_type="R2C")
try:
with rplan:
sp1 = cp.fft.rfftn(in1, fshape, axes=axes)
sp2 = cp.fft.rfftn(in2, fshape, axes=axes)
except Exception:
sp1 = cp.fft.rfftn(in1, fshape, axes=axes)
sp2 = cp.fft.rfftn(in2, fshape, axes=axes)
ret = cp.fft.irfftn(sp1 * sp2, fshape, axes=axes)[fslice].copy()
else:
# Need to move to cupyx.scipy.fft with CuPy v8
if (str(fshape), str(axes)) in _cupy_fft_cache:
plan = _cupy_fft_cache[(str(fshape), str(axes))]
else:
plan = _cupy_fft_cache[(str(fshape),
str(axes))] = fftpack.get_fft_plan(
in1, fshape, axes=axes)
try:
with plan:
sp1 = fftpack.fftn(in1, fshape, axes=axes)
sp2 = fftpack.fftn(in2, fshape, axes=axes)
ret = fftpack.ifftn(sp1 * sp2, axes=axes)[fslice].copy()
except Exception:
sp1 = fftpack.fftn(in1, fshape, axes=axes)
sp2 = fftpack.fftn(in2, fshape, axes=axes)
ret = fftpack.ifftn(sp1 * sp2, axes=axes)[fslice].copy()
if mode == "full":
return ret
elif mode == "same":
return _centered(ret, s1)
elif mode == "valid":
shape_valid = shape.copy()
shape_valid[axes] = s1[axes] - s2[axes] + 1
return _centered(ret, shape_valid)
else:
raise ValueError("acceptable mode flags are \
'valid',"
" 'same', or 'full'")