def _sortbyabs(array, axis=0):
"""
Sort array along a given axis by absolute values.
Parameters
----------
array : (N, ..., M) ndarray
Array with input image data.
axis : int
Axis along which to sort.
Returns
-------
array : (N, ..., M) ndarray
Array sorted along a given axis by absolute values.
Notes
-----
Modified from: http://stackoverflow.com/a/11253931/4067734
"""
# Create auxiliary array for indexing
index = list(cp.ix_(*[cp.arange(i) for i in array.shape]))
# Get indices of abs sorted array
index[axis] = cp.abs(array).argsort(axis)
# Return abs sorted array
return array[tuple(index)]
def adj(self, coeffs, shifts=None):
"""Adjoint operator for GWP decomposition
Parameters
----------
coeffs : [K](Nangles,L0,L1) complex64
Decomposition coefficients for box levels 0:K, angles 0:Nangles,
defined box grids with sizes [L0[k],L1[k]], k=0:K
Returns
-------
f : [N,N] complex64
2D function in the space domain
"""
print('adj transform')
# build spectrum by using gwp coefficients
F = cp.zeros(self.fgridshape, dtype="complex64")
# loop over box orientations
for ang in range(0, self.nangles):
print('angle', ang)
g = cp.zeros(int(np.prod(self.boxshape[-1])), dtype='complex64')
for k in range(self.K):
fcoeffs = cp.array(coeffs[k][ang])
# normalize
fcoeffs /= (np.prod(self.boxshape[-1]))
if(shifts is not None):
# shift wrt region in rectangular grid
fcoeffs = util.checkerboard(cp.fft.fftn(
util.checkerboard(fcoeffs), norm='ortho'), inverse=True)
[py, px] = cp.meshgrid(
self.phasemanyy[k]*shifts[ang, k, 0], self.phasemanyx[k]*shifts[ang, k, 1], indexing='ij')
fcoeffs *= cp.exp(1j*(px+py))
fcoeffs = util.checkerboard(cp.fft.ifftn(
util.checkerboard(fcoeffs), norm='ortho'), inverse=True)
# shift wrt xi_cent
fcoeffs *= cp.exp(1j*self.phase[k])
fcoeffs = util.checkerboard(cp.fft.fftn(
util.checkerboard(fcoeffs), norm='ortho'), inverse=True)
# broadcast values to smaller boxes, multiply by the gwp kernel function
g[self.inds[k]] += self.gwpf[k]*fcoeffs.flatten()
g = g.reshape(self.boxshape[-1])
# 2) Interpolation to the global grid
# Conventional scattering operation from the global to local grid is replaced
# by an equivalent gathering operation.
# calculate global grid coordinates in the space domain
xr = self.coordinates_adj(ang)
# extract ids of the global spectrum subregion contatining the box
[idsy, idsx] = self.subregion_ids(ang)
# gathering to the global grid
F[cp.ix_(idsy, idsx)] += self.U.gather(g, xr,
g.shape).reshape(len(idsy), len(idsx))
# util.mplt(F)
# 3) apply 2D IFFT, compensate for the USFFT kernel function in the space domain
f = self.U.ifftcomp(
F.reshape(self.fgridshape)).get().astype('complex64')
# free cupy pools
# mempool.free_all_blocks()
# pinned_mempool.free_all_blocks()
return f
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 cp_regularize_conv(kernel, input_shape, clip_to, devices=-1):
A = cp.fft.fft2(kernel, input_shape, axes=(0, 1))
# complex SVD using real formulation https://www.osti.gov/servlets/purl/756121
A_resh = cp.reshape(A, (-1, A.shape[2], A.shape[3]))
R_resh = A_resh.real
I_resh = A_resh.imag
upper_batch = cp.concatenate((R_resh, I_resh), axis=2)
lower_batch = cp.concatenate((-I_resh, R_resh), axis=2)
K_batch = cp.concatenate((upper_batch, lower_batch), axis=1)
KTK = cp.matmul(cp.transpose(K_batch, axes=(0, 2, 1)), K_batch)
if isinstance(devices, int):
_, sKTKinv = batch_sqrtm(KTK, numIters=10, reg=1.)
elif isinstance(devices, dict):
if len(devices) == 1:
_, sKTKinv = batch_sqrtm(KTK, numIters=10, reg=1.)
else:
sKTKinv = batch_sqrtm_multi_gpu(KTK, devices, numIters=10, reg=1.)
K_tran_batch = cp.squeeze(cp.matmul(K_batch,
sKTKinv)) * clip_to #* np.sqrt(1./2.)
if len(K_tran_batch.shape) == 2:
K_tran_batch = K_tran_batch[cp.newaxis, :, :]
_, M, N = K_tran_batch.shape
A_tran = K_tran_batch[:, 0:M / 2,
0:N / 2] - 1j * K_tran_batch[:, 0:M / 2, -N / 2:]
A_tran = cp.reshape(A_tran, A.shape).astype(cp.complex64)
clipped_kernel = cp.fft.ifft2(A_tran, axes=(0, 1)).real
return clipped_kernel[cp.ix_(*[range(d) for d in kernel.shape])]
def fwd(self, f, shifts=None):
"""Forward operator for GWP decomposition
Parameters
----------
f : [N,N] complex64
2D function in the space domain
Returns
-------
coeffs : [K](Nangles,L3,L0,L1) complex64
Decomposition coefficients for box levels 0:K, angles 0:Nangles,
defined on box grids with sizes [L3[k],L0[k],L1[k]], k=0:K
"""
print('fwd transform')
# 1) Compensate for the USFFT kernel function in the space domain and apply 2D FFT
F = self.U.compfft(cp.array(f))
# allocate memory for coefficeints
coeffs = [None]*self.K
for k in range(self.K):
coeffs[k] = np.zeros(
[self.nangles, *self.boxshape[k]], dtype='complex64')
# loop over box orientations
for ang in range(0, self.nangles):
print('angle', ang)
# 2) Interpolation to the local box grid.
# Gathering operation from the global to local grid.
# extract ids of the global spectrum subregion contatining the box
[idsy, idsx] = self.subregion_ids(ang)
# calculate box coordinates in the space domain
xr = self.coordinates_fwd(ang)
# gather values to the box grid
g = self.U.gather(F[cp.ix_(idsy, idsx)], xr, F.shape)
# 3) IFFTs on each box.
# find coefficients on each box
for k in range(self.K):
# broadcast values to smaller boxes, multiply by the gwp kernel function
fcoeffs = self.gwpf[k]*g[self.inds[k]]
fcoeffs = fcoeffs.reshape(self.boxshape[k])
# shift wrt xi_cent
fcoeffs = util.checkerboard(cp.fft.ifftn(
util.checkerboard(fcoeffs), norm='ortho'), inverse=True)
fcoeffs *= cp.exp(-1j*self.phase[k])
if(shifts is not None):
# shift wrt region in rectangular grid
fcoeffs = util.checkerboard(cp.fft.fftn(
util.checkerboard(fcoeffs), norm='ortho'), inverse=True)
[py, px] = cp.meshgrid(
self.phasemanyy[k]*shifts[ang, k, 0], self.phasemanyx[k]*shifts[ang, k, 1], indexing='ij')
fcoeffs *= cp.exp(-1j*(px+py))
fcoeffs = util.checkerboard(cp.fft.ifftn(
util.checkerboard(fcoeffs), norm='ortho'), inverse=True)
# normalize
fcoeffs /= (np.prod(self.boxshape[-1]))
coeffs[k][ang] = fcoeffs.get()
# free cupy pools
# mempool.free_all_blocks()
# pinned_mempool.free_all_blocks()
return coeffs