def _get_nearplane_steps(diff, dOP, dPO, A1, A4, recover_psi, recover_probe):
# (22) Use least-squares to find the optimal step sizes simultaneously
if recover_psi and recover_probe:
b1 = cp.sum((dOP.conj() * diff).real, axis=(-2, -1))
b2 = cp.sum((dPO.conj() * diff).real, axis=(-2, -1))
A2 = cp.sum((dOP * dPO.conj()), axis=(-2, -1))
A3 = A2.conj()
determinant = A1 * A4 - A2 * A3
x1 = -cp.conj(A2 * b2 - A4 * b1) / determinant
x2 = cp.conj(A1 * b2 - A3 * b1) / determinant
elif recover_psi:
b1 = cp.sum((dOP.conj() * diff).real, axis=(-2, -1))
x1 = b1 / A1
elif recover_probe:
b2 = cp.sum((dPO.conj() * diff).real, axis=(-2, -1))
x2 = b2 / A4
if recover_psi:
step = 0.9 * cp.maximum(0, x1[..., None, None].real)
# (27b) Object update
weighted_step_psi = cp.mean(step, keepdims=True, axis=-5)
if recover_probe:
step = 0.9 * cp.maximum(0, x2[..., None, None].real)
weighted_step_probe = cp.mean(step, axis=-5, keepdims=True)
else:
weighted_step_probe = None
return weighted_step_psi, weighted_step_probe
def cg_tomo(self, xi0, xi1, K, init, rho, tau, titer):
# minimization functional
def minf(KRu, gu):
return rho * cp.linalg.norm(KRu - xi0)**2 + tau * cp.linalg.norm(
gu - xi1)**2
u = init.copy()
gamma = 2 # init gamma as a large value
for i in range(titer):
KRu = K * self.fwd_tomo_batch(u)
gu = self.fwd_reg(u)
grad = rho*self.adj_tomo_batch(cp.conj(K)*(KRu-xi0)) + \
tau*self.adj_reg(gu-xi1)
# Dai-Yuan direction
if i == 0:
d = -grad
else:
d = -grad+cp.linalg.norm(grad)**2 / \
((cp.sum(cp.conj(d)*(grad-grad0))))*d
grad0 = grad
# line search
gamma = self.line_search(minf, gamma, KRu, gu,
K * self.fwd_tomo_batch(d),
self.fwd_reg(d))
# update step
u = u + gamma * d
return u
def take_lagr(self, psi, phi, data, h, e, lamd, mu, alpha, rho, tau,
model):
lagr = cp.zeros(7, dtype="float32")
# Lagrangian ptycho part by angles partitions
for k in range(0, self.ptheta):
ids = np.arange(k * self.tomoshapep[0],
(k + 1) * self.tomoshapep[0])
self.cl_ptycho.setobj(self.scan[:, ids].data.ptr,
self.prb.data.ptr)
fpsi = self.fwd_ptycho(psi[ids])
datap = cp.array(data[ids])
if (model == 'poisson'):
lagr[0] += cp.sum(
cp.abs(fpsi)**2 - 2 * datap * self.mlog(cp.abs(fpsi)) -
(datap - 2 * datap * self.mlog(cp.sqrt(datap))))
if (model == 'gaussian'):
lagr[0] += cp.linalg.norm(cp.abs(fpsi) - cp.sqrt(datap))**2
lagr[1] = alpha * cp.sum(
np.sqrt(cp.real(cp.sum(phi * cp.conj(phi), 0))))
lagr[2] = 2 * cp.sum(cp.real(cp.conj(lamd) * (h - psi)))
lagr[3] = rho * cp.linalg.norm(h - psi)**2
lagr[4] = 2 * cp.sum(np.real(cp.conj(mu) * (e - phi)))
lagr[5] = tau * cp.linalg.norm(e - phi)**2
lagr[6] = cp.sum(lagr[0:5])
return lagr
def grad_ptycho(self, data, psi, prb, scan, zlamd, rho, niter):
"""Gradient solver for the ptychography problem |||FQpsi|-sqrt(data)||^2_2 + rho||psi-zlamd||^2_2"""
# minimization functional
def minf(fpsi, psi):
f = cp.linalg.norm(cp.abs(fpsi) - cp.sqrt(data))**2
if(rho > 0):
f += rho*cp.linalg.norm(psi-zlamd)**2
return f
for i in range(niter):
# compute the gradient
fpsi = self.fwd_ptycho(psi, prb, scan)
gradpsi = self.adj_ptycho(
fpsi - cp.sqrt(data)*fpsi/(cp.abs(fpsi)+1e-32), prb, scan)
# normalization coefficient for skipping the line search procedure
afpsi = self.adj_ptycho(fpsi, prb, scan)
norm_coeff = cp.real(cp.sum(psi*cp.conj(afpsi)) /
(cp.sum(afpsi*cp.conj(afpsi))+1e-32))
if(rho > 0):
gradpsi += rho*(psi-zlamd)
gradpsi *= min(1/rho, norm_coeff)/2
else:
gradpsi *= norm_coeff/2
# update psi
psi = psi - 0.5*gradpsi
# check convergence
# print(f'{i}) {minf(fpsi, psi).get():.2e} ')
return psi
def _probability(self, x):
"""Unnormalized probability of one configuration P(x)
Parameters
----------
x : numpy array, shape (n_features,)
One configuration
Returns
-------
probability : float
"""
w2 = np.reshape(self.w,
(self.n_features, self.d, self.D, self.D, self.mu))
tmp = w2[0, x[0], 0, :, :]
tmp2 = np.einsum('ij,kj->ik', tmp,
np.conj(tmp)).reshape(self.D * self.D)
for i in xrange(1, self.n_features - 1):
tmp = np.einsum('imj,klj->ikml', w2[i, x[i], :, :, :],
np.conj(w2[i, x[i], :, :, :])).reshape(
(self.D * self.D, self.D * self.D))
tmp2 = np.dot(tmp2, tmp)
tmp = np.einsum(
'ij,kj->ik', w2[self.n_features - 1, x[self.n_features - 1], :,
0, :],
np.conj(w2[self.n_features - 1, x[self.n_features - 1], :,
0, :])).reshape(self.D * self.D)
probability = np.abs(np.inner(tmp2, tmp))
return probability
def calc_Qpsi(fz, fp, Wfn):
Qpsi = [2 * fz * Wfn[0] + fp * Wfn[1],
cp.conj(fp) * Wfn[0] + fz * Wfn[1] + cp.sqrt(3 / 2) * fp * Wfn[2],
cp.sqrt(3 / 2) * (cp.conj(fp) * Wfn[1] + fp * Wfn[3]),
cp.sqrt(3 / 2) * cp.conj(fp) * Wfn[2] - fz * Wfn[3] + fp * Wfn[4],
cp.conj(fp) * Wfn[3] - 2 * fz * Wfn[4]]
return Qpsi
def sampling(self):
self.reset_estimators()
# for running average
L_run = 0.0
# Clear LRU
self.lru.clear()
# do it sweep times
it = 0
while self.move_rate < self.sweeps:
# for i in range(self.sweeps):
#make a move
self.move()
# calculate estimators and expectation values for SR
key = str(self.S_i) + str(self.S_j)
local_L = self.search_dictionary(key)
devs = cp.conj(
cp.asarray(self.nqs.nqs_derivative(self.S_i, self.S_j)))
self.L += local_L
self.Ok = cp.add(self.Ok, devs)
self.Okk = cp.add(
self.Okk,
cp.dot(devs, cp.conj(devs.T)) +
cp.dot(devs, cp.conj(devs.T)).T)
# self.Okk = cp.add(self.Okk, cp.dot(devs, cp.conj(devs.T)));
local_L_gpu = cp.asarray(local_L)
self.LOk = cp.add(
self.LOk,
local_L_gpu * devs + cp.conj(local_L_gpu) * cp.conj(devs))
# self.LOk = cp.add(self.LOk, local_L_gpu*devs);
# calculate batched running variance
delta = np.real(local_L) - L_run
L_run += delta / (it + 1)
delta2 = np.real(local_L) - L_run
self.L_sq += delta * delta2
it += 1
# Finalizing the blocked running variance calculation
self.L_sq = np.sqrt(self.L_sq) / (it - 1)
if self.gpu:
return it, self.L, cp.asnumpy(self.Ok), cp.asnumpy(
self.Okk), cp.asnumpy(
self.LOk), self.L_sq, self.move_rate, len(self.lru)
else:
return it, self.L, self.Ok, self.Okk, self.LOk, self.L_sq, self.move_rate, len(
self.lru)
def _derivativenorm(self):
"""Compute the derivative of the norm
Returns
-------
derivative : numpy array, shape (m_parameters,)
"""
w2 = np.reshape(self.w,
(self.n_features, self.d, self.D, self.D, self.mu))
derivative = np.zeros(
(self.n_features, self.d, self.D, self.D, self.mu),
dtype=np.complex128)
tmp = np.zeros((self.n_features, self.D * self.D), dtype=np.complex128)
tmp2 = np.zeros((self.n_features, self.D * self.D),
dtype=np.complex128)
tmp[0, :] = np.einsum('ijk,ilk->jl', w2[0, :, 0, :, :],
np.conj(w2[0, :,
0, :, :])).reshape(self.D * self.D)
for i in xrange(1, self.n_features - 1):
newtmp = np.einsum('pimj,pklj->ikml', w2[i, :, :, :, :],
np.conj(w2[i, :, :, :, :])).reshape(
(self.D * self.D, self.D * self.D))
tmp[i, :] = np.dot(tmp[i - 1, :], newtmp)
newtmp = np.einsum('ijk,ilk->jl', w2[self.n_features - 1, :, :, 0, :],
np.conj(w2[self.n_features - 1, :, :,
0, :])).reshape(self.D * self.D)
mpscontracted = np.inner(tmp[self.n_features - 2, :], newtmp)
tmp[self.n_features - 1, :] = mpscontracted
tmp2[self.n_features - 1, :] = newtmp
for i in xrange(self.n_features - 2, -1, -1):
newtmp = np.einsum('pimj,pklj->ikml', w2[i, :, :, :, :],
np.conj(w2[i, :, :, :, :])).reshape(
(self.D * self.D, self.D * self.D))
tmp2[i, :] = np.dot(newtmp, tmp2[i + 1, :])
newtmp = np.einsum('ijk,ilk->jl', w2[0, :, 0, :, :],
np.conj(w2[0, :, 0, :, :])).reshape(self.D * self.D)
tmp2[0, :] = np.inner(newtmp, tmp2[1, :])
for j in xrange(self.d):
derivative[0, j, 0, :, :] = 2 * np.einsum(
'ij,il->lj', w2[0, j, 0, :, :], tmp2[1, :].reshape(
self.D, self.D))
derivative[self.n_features-1,j,:,0,:]=\
2*np.einsum('ij,il->lj',w2[self.n_features-1,j,:,0,:],
tmp[self.n_features-2,:].reshape(self.D,self.D))
for i in xrange(1, self.n_features - 1):
temp1 = tmp[i - 1, :].reshape(self.D, self.D)
temp2 = tmp2[i + 1, :].reshape(self.D, self.D)
for j in xrange(self.d):
derivative[i, j, :, :, :] = 2 * np.einsum(
'ikm,ij,kl->jlm', w2[i, j, :, :, :], temp1, temp2)
return derivative.reshape(self.m_parameters)
def interaction_flow(wfn_plus, wfn_0, wfn_minus, C, S, Fz, F_perp, dt, V, p,
c0, n):
"""Solves the interaction part of the flow."""
new_wfn_plus = (C * wfn_plus - S *
(Fz * wfn_plus + cp.conj(F_perp) / cp.sqrt(2) * wfn_0)
) * cp.exp(-1j * dt * (V - p + c0 * n))
new_wfn_0 = (C * wfn_0 - S / cp.sqrt(2) *
(F_perp * wfn_plus + cp.conj(F_perp) * wfn_minus)) * cp.exp(
-1j * dt * (V + c0 * n))
new_wfn_minus = (C * wfn_minus - S *
(F_perp / cp.sqrt(2) * wfn_0 - Fz * wfn_minus)) * cp.exp(
-1j * dt * (V + p + c0 * n))
return new_wfn_plus, new_wfn_0, new_wfn_minus
def calc_spin_dens(wfn_plus, wfn_0, wfn_minus, dt, c2):
"""Calculates various quantities such as spin vectors, sin and cosine terms and the atomic density."""
spin_perp = cp.sqrt(2.) * (cp.conj(wfn_plus) * wfn_0 +
cp.conj(wfn_0) * wfn_minus)
spin_z = cp.abs(wfn_plus)**2 - cp.abs(wfn_minus)**2
F = cp.sqrt(cp.abs(spin_z)**2 +
cp.abs(spin_perp)**2) # Magnitude of spin vector
cos_term = cp.cos(c2 * F * dt)
sin_term = 1j * cp.sin(c2 * F * dt) / F
sin_term = cp.nan_to_num(sin_term) # Corrects division by 0
density = cp.abs(wfn_minus)**2 + cp.abs(wfn_0)**2 + cp.abs(wfn_plus)**2
return spin_perp, spin_z, cos_term, sin_term, density
def pure_state_overlap_cupy(wf1: cupy.ndarray, wf2: cupy.ndarray):
"""
Returns:
a complex scalar of array (depends on the shape of wf1 and wf2)
"""
ndim = wf1.ndim
if ndim != wf2.ndim:
raise ValueError("wf1:{0:s}\nwf2:{1:s}".format(str(ndim), str(wf2.ndim)))
if ndim == 1:
return transpose(conj(wf1)).dot(wf2)
elif ndim == 2:
wf1 = conj(wf1)
return _sum(wf1 * wf2, axis=0) # element-wise
else:
raise ValueError(str(wf1.shape))
def calculate_snrs(self, interferometer, waveform_polarizations):
name = interferometer.name
signal_ifo = xp.sum(
xp.vstack([
waveform_polarizations[mode] * float(
interferometer.antenna_response(
self.parameters["ra"],
self.parameters["dec"],
self.parameters["geocent_time"],
self.parameters["psi"],
mode,
)) for mode in waveform_polarizations
]),
axis=0,
)[interferometer.frequency_mask]
time_delay = (self.parameters["geocent_time"] -
interferometer.strain_data.start_time +
interferometer.time_delay_from_geocenter(
self.parameters["ra"],
self.parameters["dec"],
self.parameters["geocent_time"],
))
signal_ifo *= xp.exp(-2j * np.pi * time_delay * self.frequency_array)
d_inner_h = xp.sum(
xp.conj(signal_ifo) * self.strain[name] / self.psds[name])
h_inner_h = xp.sum(xp.abs(signal_ifo)**2 / self.psds[name])
return d_inner_h, h_inner_h
def getRandomHermitianMatrix(M):
ret = xp.diag(0j + randn(M))
for y in range(0, M - 1):
for x in range(y + 1, M):
ret[y, x] = randn_c()
ret[x, y] = xp.conj(ret[y, x])
return ret
def cg_tomo_reg(self, xi0, u, titer, tau, xi1, gpu=0, dbg=False):
"""CG solver for 1 slice partition"""
# minimization functional
def minf(Ru, gu):
f = cp.linalg.norm(Ru-xi0)**2+tau*cp.linalg.norm(gu-xi1)**2
return f
for i in range(titer):
Ru = self.fwd_tomo(u, gpu)
gu = self.fwd_reg(u)
grad = (self.adj_tomo(Ru-xi0, gpu) /
(self.ntheta * self.n/2)+tau*self.adj_reg(gu-xi1))/2/max(tau, 1) # normalized gradient
if i == 0:
d = -grad
else:
d = -grad+cp.linalg.norm(grad)**2 / \
(cp.sum(cp.conj(d)*(grad-grad0))+1e-32)*d
# line search
Rd = self.fwd_tomo(d, gpu)
gd = self.fwd_reg(d)
gamma = 0.5*self.line_search_ext(minf, 1, Ru, Rd, gu, gd)
grad0 = grad
# update step
u = u + gamma*d
# check convergence
if (dbg):
print("%4d, %.3e, %.7e" %
(i, gamma, minf(Ru, gu)))
return u
def cg_lam(self, data0, u0, theta0, titer, dbg=False):
"""CG solver for ||Lu-data||_2"""
u = cp.asarray(u0)
theta = cp.asarray(theta0)
data = cp.asarray(data0)
# minimization functional
def minf(Lu):
f = cp.linalg.norm(Lu - data)**2
return f
for i in range(titer):
Lu = self.fwd_lam(u, theta)
grad = self.adj_lam(Lu-data, theta) * 1 / \
self.ntheta/self.n0/self.n1/self.n2
if i == 0:
d = -grad
else:
d = -grad+cp.linalg.norm(grad)**2 / \
(cp.sum(cp.conj(d)*(grad-grad0))+1e-32)*d
# line search
Ld = self.fwd_lam(d, theta)
gamma = 0.5 * self.line_search(minf, 1, Lu, Ld)
grad0 = grad
# update step
u = u + gamma * d
# check convergence
if (dbg == True):
print("%4d, %.3e, %.7e" % (i, gamma, minf(Lu)))
if (isinstance(u0, np.ndarray)):
u = u.get()
return u
def cg_deform(self, data, psi, flow, titer, xi1=0, rho=0, gpu=0):
"""CG solver for deformation"""
# minimization functional
def minf(psi, Dpsi):
f = cp.linalg.norm(Dpsi-data)**2+rho*cp.linalg.norm(psi-xi1)**2
return f
for i in range(titer):
Dpsi = self.apply_flow_gpu(psi, flow, gpu)
grad = (self.apply_flow_gpu(Dpsi-data, -flow, gpu) +
rho*(psi-xi1))/max(rho, 1)
if i == 0:
d = -grad
else:
d = -grad+cp.linalg.norm(grad)**2 / \
(cp.sum(cp.conj(d)*(grad-grad0))+1e-32)*d
# line search
Td = self.apply_flow_gpu(d, flow, gpu)
gamma = 0.5*self.line_search(minf, 1, psi, Dpsi, d, Td)
if(gamma == 0):
break
grad0 = grad
# update step
psi = psi + gamma*d
# check convergence
# if (0):
# print("%4d, %.3e, %.7e" %
# (i, gamma, minf(psi, Dpsi+gamma*Td)))
return psi
def cg_tomo(self, xi0, u, titer):
"""CG solver for ||Ru-xi0||_2"""
# minimization functional
def minf(Ru):
f = cp.linalg.norm(Ru - xi0)**2
return f
for i in range(titer):
Ru = self.fwd_tomo(u)
grad = self.adj_tomo(Ru-xi0) / \
(self.ntheta * self.n/2)
if i == 0:
d = -grad
else:
d = -grad+cp.linalg.norm(grad)**2 / \
(cp.sum(cp.conj(d)*(grad-grad0))+1e-32)*d
# line search
Rd = self.fwd_tomo(d)
gamma = 0.5 * self.line_search(minf, 1, Ru, Rd)
grad0 = grad
# update step
u = u + gamma * d
# check convergence
if (np.mod(i, 1) == -1):
print("%4d, %.3e, %.7e" % (i, gamma, minf(Ru)))
return u
def calc_spin_vectors_cuda(psiP2, psiP1, psi0, psiM1, psiM2):
"""
:param psiP2: psi_+2 component
:param psiP1: psi_+1 component
:param psi0: psi_0 component
:param psiM1: psi_-1 component
:param psiM2: psi_-2 component
:return: fp, fz: perpendicular and longitudinal spin vectors
"""
fp = cp.sqrt(6) * (psiP1 * cp.conj(psi0) + psi0 * cp.conj(psiM1)) + \
2 * (psiM1 * cp.conj(psiM2) + psiP2 * cp.conj(psiP1))
fz = 2 * (cp.abs(psiP2)**2 -
cp.abs(psiM2)**2) + cp.abs(psiP1)**2 - cp.abs(psiM1)**2
return fp, fz
def cupy_signal(signal):
amp = cp.sqrt(cp.real(signal * cp.conj(signal)))
phase = cp.angle(signal)
real = cp.real(signal)
imag = cp.imag(signal)
return amp, phase, real, imag
def mvdr(x, sv):
"""
Minimum variance distortionless response (MVDR) beamformer weights
Parameters
----------
x : ndarray
Received signal, assume 2D array with size [num_sensors, num_samples]
sv: ndarray
Steering vector, assume 1D array with size [num_sensors, 1]
Note: Unlike MATLAB where input matrix x is of size MxN where N represents
the number of array elements, we assume row-major formatted data where each
row is assumed to be complex-valued data from a given sensor (i.e. NxM)
"""
if x.shape[0] > x.shape[1]:
raise ValueError('Matrix has more sensors than samples. Consider \
transposing and remember cuSignal is row-major, unlike MATLAB')
if x.shape[0] != sv.shape[0]:
raise ValueError('Steering Vector and input data do not align')
R = cp.cov(x)
R_inv = cp.linalg.inv(R)
svh = cp.transpose(cp.conj(sv))
wB = cp.matmul(R_inv, sv)
# wA is a 1x1 scalar
wA = cp.matmul(svh, wB)
w = wB / wA
return w
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 cg_ptycho(self, data, init, h, lamd, rho, piter, model):
# minimization functional
def minf(psi, fpsi):
if model == 'gaussian':
f = cp.linalg.norm(cp.abs(fpsi) - cp.sqrt(data))**2
elif model == 'poisson':
f = cp.sum(
cp.abs(fpsi)**2 - 2 * data * self.mlog(cp.abs(fpsi)))
f += rho * cp.linalg.norm(h - psi + lamd / rho)**2
return f
psi = init.copy()
gamma = 8 # init gamma as a large value
for i in range(piter):
fpsi = self.fwd_ptycho(psi)
if model == 'gaussian':
grad = self.adj_ptycho(fpsi - cp.sqrt(data) *
cp.exp(1j * cp.angle(fpsi)))
elif model == 'poisson':
grad = self.adj_ptycho(fpsi - data * fpsi /
(cp.abs(fpsi)**2 + 1e-32))
grad -= rho * (h - psi + lamd / rho)
# Dai-Yuan direction
if i == 0:
d = -grad
else:
d = -grad+cp.linalg.norm(grad)**2 / \
((cp.sum(cp.conj(d)*(grad-grad0))))*d
grad0 = grad
# line search
fd = self.fwd_ptycho(d)
gamma = self.line_search(minf, gamma, psi, fpsi, d, fd)
psi = psi + gamma * d
# print(i,minf(psi,fpsi))
return psi
def cg_shift_gpu(self, data, psi, flow, titer, xi1=0, rho=0, dbg=False):
"""CG solver for shift"""
# minimization functional
def minf(psi, Tpsi):
f = cp.linalg.norm(Tpsi -
data)**2 + rho * cp.linalg.norm(psi - xi1)**2
return f
for i in range(titer):
Tpsi = self.apply_shift(psi, flow)
#flow = self.registration_shift(data, psi, 1)
# Tpsi = self.apply_shift(psi, flow)
# print('a',np.linalg.norm(Tpsi-data))
grad = (self.apply_shift(Tpsi - data, -flow) + rho *
(psi - xi1)) / max(rho, 1)
if i == 0:
d = -grad
else:
d = -grad+cp.linalg.norm(grad)**2 / \
(np.sum(cp.conj(d)*(grad-grad0))+1e-32)*d
# line search
Td = self.apply_shift(d, flow)
gamma = 0.5 * self.line_search(minf, 1, psi, Tpsi, d, Td)
grad0 = grad
# update step
psi = psi + gamma * d
# check convergence
if (dbg):
print("%4d, %.3e, %.7e" %
(i, gamma, minf(psi, Tpsi + gamma * Td)))
return psi, flow
def _findCarrier_cupy(self, band0, band1, mask):
band0 = cp.asarray(band0)
band1 = cp.asarray(band1)
mask = cp.asarray(mask)
band = band0 * band1
ixf = abs(cp.fft.fftshift(cp.fft.fft2(cp.fft.fftshift(band))))
pyc0, pxc0 = self._findPeak_cupy(
(ixf - gaussian_filter_cupy(ixf, 20)) * mask)
ixfz, Kx, Ky = self._zoomf_cupy(band, self.N,
cp.asarray(self._kx)[pyc0, pxc0],
cp.asarray(self._ky)[pyc0, pxc0], 100,
self._dk * self.N)
pyc, pxc = self._findPeak_cupy(abs(ixfz))
kx = Kx[pxc]
ky = Ky[pyc]
otf_exclude_min_radius = 0.5
otf_exclude_max_radius = 1.5
kr = cp.sqrt(cp.asarray(self._kx)**2 + cp.asarray(self._ky)**2)
m = (kr < 2)
otf = cp.fft.fftshift(self._tfm_cupy(kr, m) + (1 - m))
otf_mask = (kr > otf_exclude_min_radius) & (kr <
otf_exclude_max_radius)
otf_mask_for_band_common_freq = cp.fft.fftshift(
otf_mask
& cupyx.scipy.ndimage.shift(otf_mask, (pyc0 -
(self.N // 2 + 1), pxc0 -
(self.N // 2 + 1)),
order=0))
band0_common = cp.fft.ifft2(
cp.fft.fft2(band0) / otf * otf_mask_for_band_common_freq)
xx = cp.arange(-self.N / 2 * self._dx,
self.N / 2 * self._dx,
self._dx,
dtype=np.single)
phase_shift_to_xpeak = cp.exp(-1j * kx * xx * 2 * pi * self.NA /
self.wavelength)
phase_shift_to_ypeak = cp.exp(-1j * ky * xx * 2 * pi * self.NA /
self.wavelength)
band1_common = cp.fft.ifft2(
cp.fft.fft2(band1) /
otf * otf_mask_for_band_common_freq) * cp.outer(
phase_shift_to_ypeak, phase_shift_to_xpeak)
scaling = 1 / cp.sum(band0_common * cp.conj(band0_common))
cross_corr_result = cp.sum(band0_common * band1_common) * scaling
ampl = cp.abs(cross_corr_result) * 2
phase = cp.angle(cross_corr_result)
return kx.get(), ky.get(), phase.get(), ampl.get()
def test_cuFloatComplex(self):
N = 100
block = 32
grid = (N + block - 1) // block
dtype = cupy.complex64
mod = cupy.RawModule(code=_test_cuComplex, translate_cucomplex=True)
a = cupy.random.random((N, )) + 1j * cupy.random.random((N, ))
a = a.astype(dtype)
b = cupy.random.random((N, )) + 1j * cupy.random.random((N, ))
b = b.astype(dtype)
c = cupy.random.random((N, )) + 1j * cupy.random.random((N, ))
c = c.astype(dtype)
out = cupy.zeros((N, ), dtype=dtype)
out_float = cupy.zeros((N, ), dtype=cupy.float32)
out_up = cupy.zeros((N, ), dtype=cupy.complex128)
ker = mod.get_function('test_addf')
ker((grid, ), (block, ), (a, b, out))
assert (out == a + b).all()
ker = mod.get_function('test_subf')
ker((grid, ), (block, ), (a, b, out))
assert (out == a - b).all()
ker = mod.get_function('test_mulf')
ker((grid, ), (block, ), (a, b, out))
assert cupy.allclose(out, a * b)
ker = mod.get_function('test_divf')
ker((grid, ), (block, ), (a, b, out))
assert (out == a / b).all()
ker = mod.get_function('test_conjf')
ker((grid, ), (block, ), (a, out))
assert (out == cupy.conj(a)).all()
ker = mod.get_function('test_absf')
ker((grid, ), (block, ), (a, out_float))
assert (out_float == cupy.abs(a)).all()
ker = mod.get_function('test_fmaf')
ker((grid, ), (block, ), (a, b, c, out))
assert cupy.allclose(out, a * b + c)
ker = mod.get_function('test_makef')
ker((grid, ), (block, ), (out, ))
# because of precision issue, the (A==B).all() semantics would fail
assert cupy.allclose(out, 1.8 - 1j * 8.7)
ker = mod.get_function('test_upcast')
ker((grid, ), (block, ), (a, out_up))
assert (out_up == a.astype(cupy.complex128)).all()
# NumPy scalars.
b = cupy.complex64(2 + 3j)
ker = mod.get_function('test_addf_scalar')
ker((grid, ), (block, ), (a, b, out))
assert (out == a + b).all()
def solve_reg(self, u, mu, tau, alpha):
z = self.fwd_reg(u) + mu / tau
# Soft-thresholding
za = cp.sqrt(cp.real(cp.sum(z * cp.conj(z), 0)))
z[:, za <= alpha / tau] = 0
z[:, za > alpha/tau] -= alpha/tau * \
z[:, za > alpha/tau]/(za[za > alpha/tau])
return z
def make_density_matrix_cupy(wf: cupy.ndarray):
if wf.ndim == 1:
return outer(wf, cupy.conj(wf))
if wf.ndim == 2:
wf_dim, num_wf = wf.shape
try:
ret = empty(shape=(num_wf, wf_dim, wf_dim), dtype=wf.dtype)
except OutOfMemoryError:
logger.critical(
"OOM when creating density matrix for wavefunction "
f"of shape {wf.shape}"
)
raise
for wf_idx in range(num_wf):
a_wf = wf[:, wf_idx]
ret[wf_idx, :, :] = outer(a_wf, conj(a_wf))
return ret
raise NotImplementedError(wf.shape)
def pulse_compression(x, template, normalize=False, window=None, nfft=None):
"""
Pulse Compression is used to increase the range resolution and SNR
by performing matched filtering of the transmitted pulse (template)
with the received signal (x)
Parameters
----------
x : ndarray
Received signal, assume 2D array with [num_pulses, sample_per_pulse]
template : ndarray
Transmitted signal, assume 1D array
normalize : bool
Normalize transmitted signal
window : array_like, callable, string, float, or tuple, optional
Specifies the window applied to the signal in the Fourier
domain.
nfft : int, size of FFT for pulse compression. Default is number of
samples per pulse
Returns
-------
compressedIQ : ndarray
Pulse compressed output
"""
[num_pulses, samples_per_pulse] = x.shape
if nfft is None:
nfft = samples_per_pulse
if window is not None:
Nx = len(template)
if callable(window):
W = window(cp.fft.fftfreq(Nx))
elif isinstance(window, cp.ndarray):
if window.shape != (Nx, ):
raise ValueError("window must have the same length as data")
W = window
else:
W = get_window(window, Nx, False)
template = cp.multiply(template, W)
if normalize is True:
template = cp.divide(template, cp.linalg.norm(template))
fft_x = cp.fft.fft(x, nfft)
fft_template = cp.conj(cp.tile(cp.fft.fft(template, nfft),
(num_pulses, 1)))
compressedIQ = cp.fft.ifft(cp.multiply(fft_x, fft_template), nfft)
return compressedIQ
def routine_gpu(data, sr, omega, morlet_frequency):
n_chans, n_ts = data.shape
data_gpu = cp.asarray(data)
win = cp.array(mne.time_frequency.morlet(sr, [morlet_frequency], omega)[0])
data_preprocessed = cp.zeros_like(data_gpu, dtype=cp.complex64)
surr_data = cp.zeros_like(data_preprocessed)
for i in range(n_chans):
data_preprocessed[i] = cusignal.fftconvolve(data_gpu[i], win, 'same')
data_preprocessed[i] /= cp.abs(data_preprocessed[i])
surr_data[i] = cp.roll(data_preprocessed[i],
np.random.randint(n_ts - 1))
plv = cp.inner(data_preprocessed, cp.conj(data_preprocessed)) / n_ts
plv_surr = cp.inner(surr_data, cp.conj(surr_data)) / n_ts
return cp.asnumpy(plv), cp.asnumpy(plv_surr)
def fsm1(InXog, alpha, nfft=8192, nl=128, start=10):
ntest = nfft * nl
x = InXog[start:start + ntest]
xafp = cp.fft.fft(x * cp.exp(1j * cp.pi * alpha * cp.arange(ntest)))
xafn = cp.fft.fft(x * cp.exp(1j * cp.pi * (-alpha) * cp.arange(ntest)))
xaf = xafp * cp.conj(xafn)
res = cp.reshape(xaf, (nfft, nl)) @ cp.hamming(nl)
return res / (ntest)
