def calc_dirspecs(i_dev: int, q_data: mp.Queue, n_data: int, q_out: mp.Queue):
""" create directional spectrogram.
pnm means SHD signal (SFT of multichannel signal)
anm means MC-SHD signal (SHD signal after mode compensation -- bnkr equalization)
_time or _t means time-domain signal
_spec means STFT data
_cp means cupy array
:param i_dev: GPU Device No.
:param q_data:
:param n_data:
:param q_out:
:return: None
"""
# Ready CUDA
cp.cuda.Device(i_dev).use()
win_cp = cp.array(win)
Ys_cp = cp.array(Ys)
sftdata_cp = SFTData(
**
{k: cp.array(v)
for k, v in asdict(sftdata).items() if v is not None})
for _ in range(n_data):
idx, i_speech, f_speech, i_loc, data, data_room = q_data.get()
data_cp = cp.array(data) # n,
data_room_cp = cp.array(data_room) # N_MIC x n
# Free-field
# n_hrm, n
anm_time_cp = cp.outer(Ys_cp[i_loc].conj(),
data_cp) # complex coefficients
if use_dv: # real coefficients
anm_time_cp = (sftdata_cp.T_real @ anm_time_cp).real
anm_spec_cp = stft(anm_time_cp, win_cp) # n_hrm x F x T
# F x T x 4
dirspec_free_cp = cp.empty((hp.n_freq, anm_spec_cp.shape[2], 4),
dtype=cp.float32)
if use_dv:
calc_direction_vec(anm_spec_cp, out=dirspec_free_cp[..., :3])
else:
calc_intensity(anm_spec_cp,
sftdata_cp.recur_coeffs,
out=dirspec_free_cp[..., :3])
cp.abs(anm_spec_cp[0], out=dirspec_free_cp[..., 3])
phase_free_cp = cp.angle(anm_spec_cp[0]) # F x T
# Room
pnm_time_cp = sftdata_cp.Yenc @ data_room_cp # n_hrm x n
# n_hrm x F x T
if use_dv: # real coefficients
# bnkr equalization in frequency domain
anm_time_cp = filter_overlap_add(pnm_time_cp,
sftdata_cp.bnkr_inv[...,
0], win_cp)
anm_t_real_cp = (sftdata_cp.T_real @ anm_time_cp).real
anm_spec_cp = stft(anm_t_real_cp, win_cp)
else: # complex coefficients
pnm_spec_cp = stft(pnm_time_cp, win_cp)
anm_spec_cp = pnm_spec_cp * sftdata_cp.bnkr_inv[:, :hp.n_freq]
# F x T x 4
dirspec_room_cp = cp.empty((hp.n_freq, anm_spec_cp.shape[2], 4),
dtype=cp.float32)
if use_dv:
calc_direction_vec(anm_spec_cp, out=dirspec_room_cp[..., :3])
else:
calc_intensity(anm_spec_cp,
sftdata_cp.recur_coeffs,
out=dirspec_room_cp[..., :3])
cp.abs(anm_spec_cp[0], out=dirspec_room_cp[..., 3])
phase_room_cp = cp.angle(anm_spec_cp[0]) # F x T
# Save (F x T x C)
dict_result = dict(
path_speech=str(f_speech),
dirspec_free=cp.asnumpy(dirspec_free_cp),
dirspec_room=cp.asnumpy(dirspec_room_cp),
phase_free_cp=cp.asnumpy(phase_free_cp)[..., np.newaxis],
phase_room_cp=cp.asnumpy(phase_room_cp)[..., np.newaxis],
)
q_out.put((idx, i_speech, i_loc, dict_result))
# Initial guess
h = cp.zeros(tomoshape, dtype='complex64', order='C') + 1
psi = cp.zeros(tomoshape, dtype='complex64', order='C') + 1
e = cp.zeros([3, *obj.shape], dtype='complex64', order='C')
phi = cp.zeros([3, *obj.shape], dtype='complex64', order='C')
lamd = cp.zeros(tomoshape, dtype='complex64', order='C')
mu = cp.zeros([3, *obj.shape], dtype='complex64', order='C')
u = cp.zeros(obj.shape, dtype='complex64', order='C')
# ADMM
u, psi, lagr = slv.admm(data, h, e, psi, phi, lamd, mu, u, alpha,
piter, titer, NITER, model)
u.real -= u[0].real
# Save result
name = 'reg'+str(alpha)+'noise'+str(noise)+'maxint' + \
str(maxint)+'prbshift'+str(prbshift)+'ntheta'+str(ntheta)+str(model)+str(piter)+str(titer)+str(NITER)
dxchange.write_tiff(u.imag.get(), 'beta/beta' + name)
dxchange.write_tiff(u.real.get(), 'delta/delta' + name)
dxchange.write_tiff(u[u.shape[0] // 2].imag.get(),
'betap/beta' + name)
dxchange.write_tiff(u[u.shape[0] // 2].real.get(),
'deltap/delta' + name)
dxchange.write_tiff(cp.angle(psi).get(), 'psi/psiangle' + name)
dxchange.write_tiff(cp.abs(psi).get(), 'psi/psiamp' + name)
if not os.path.exists('lagr'):
os.makedirs('lagr')
np.save('lagr/lagr' + name, lagr.get())
def _refineCarrier_cupy(self, band0, band1, kx_in, ky_in):
band0 = cp.asarray(band0)
band1 = cp.asarray(band1)
pxc0 = np.int(np.round(kx_in / self._dk) + self.N // 2)
pyc0 = np.int(np.round(ky_in / self._dk) + self.N // 2)
otf_exclude_min_radius = self.eta / 2
otf_exclude_max_radius = 1.5
# kr = cp.sqrt(cp.asarray(self._kx) ** 2 + cp.asarray(self._ky) ** 2)
kr = cp.asarray(self._kr, dtype=np.double)
m = (kr < 2)
otf = cp.fft.fftshift(self._tfm_cupy(kr, m) + (1 - m) * 0.0001)
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), pxc0 -
(self.N // 2)),
order=0))
band0_common = cp.fft.ifft2(
cp.fft.fft2(band0) / otf * otf_mask_for_band_common_freq)
band1_common = cp.fft.ifft2(
cp.fft.fft2(band1) / otf * otf_mask_for_band_common_freq)
band = band0_common * band1_common
mag = 25 * self.N / 256
ixfz, Kx, Ky = self._zoomf_cupy(band, self.N, np.single(self._k[pxc0]),
np.single(self._k[pyc0]), mag,
self._dk * self.N)
pyc, pxc = self._findPeak_cupy(abs(ixfz))
if self.debug:
plt.figure()
plt.title('Zoon Find carrier')
plt.imshow(abs(ixfz.get()))
kx = Kx[pxc]
ky = Ky[pyc]
xx = cp.arange(-self.N / 2 * self._dx,
self.N / 2 * self._dx,
self._dx,
dtype=np.double)
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)
scaling = 1 / cp.sum(band0_common * cp.conjugate(band0_common))
cross_corr_result = cp.sum(band0_common * band1_common * cp.outer(
phase_shift_to_ypeak, phase_shift_to_xpeak)) * scaling
ampl = cp.abs(cross_corr_result) * 2
phase = cp.angle(cross_corr_result)
return kx.get(), ky.get(), phase.get(), ampl.get()
psiP1 = cp.zeros((Nx, Ny, Nz))
psi0 = cp.zeros((Nx, Ny, Nz))
psiM1 = cp.sqrt(Tf) * 1 / cp.sqrt(3) * cp.sqrt((2 - eta))
psiM2 = cp.zeros((Nx, Ny, Nz))
Psi = [psiP2, psiP1, psi0, psiM1, psiM2] # Full 5x1 wavefunction
# Spin rotation on wavefunction:
alpha_angle = 0
beta_angle = 0.01
gamma_angle = 0
Psi = sm.rotation(Psi, alpha_angle, beta_angle, gamma_angle)
N = [dx * dy * cp.sum(cp.abs(wfn)**2)
for wfn in Psi] # Atom number of each component
theta_fix = [cp.angle(wfn) for wfn in Psi]
Psi = [cp.fft.fftn(wfn) for wfn in Psi] # Transforming wfn to Fourier space
# Store parameters in dictionary for saving
parameters = {
"c0": c0,
"c2": c2,
"c4": c4,
"q": q,
"p": p,
"omega_trap": omega_trap,
"omega_rot": omega_rot,
"alpha, beta, gamma": [alpha_angle, beta_angle, gamma_angle]
}
# Create dataset and save initial state
vort_pos_2 = iter(position_file['positions/pos_2'])
theta_1 = get_phase(N_vort, vort_pos_1, Nx, Ny, X, Y, len_x, len_y)
theta_2 = get_phase(N_vort, vort_pos_2, Nx, Ny, X, Y, len_x, len_y)
# Generate initial wavefunctions:
psi_1 = cp.sqrt(n0) * cp.exp(1j * theta_1)
psi_2 = cp.sqrt(n0) * cp.exp(1j * theta_2)
psi_1_k = cp.fft.fft2(psi_1)
psi_2_k = cp.fft.fft2(psi_2)
# Getting atom number for renormalisation in imaginary time evolution
atom_num_1 = dx * dy * cp.sum(cp.abs(psi_1)**2)
atom_num_2 = dx * dy * cp.sum(cp.abs(psi_2)**2)
# Phase of initial state to allow fixing of phase during imaginary time evolution
theta_fix_1 = cp.angle(psi_1)
theta_fix_2 = cp.angle(psi_2)
# ------------------------------------------------------------------------------------------------------------------
# Imaginary time evolution
# ------------------------------------------------------------------------------------------------------------------
for i in range(2000):
# Kinetic step:
sm.kinetic_evolution(psi_1_k, -1j * dt, Kx, Ky)
sm.kinetic_evolution(psi_2_k, -1j * dt, Kx, Ky)
psi_1 = cp.fft.ifft2(psi_1_k)
psi_2 = cp.fft.ifft2(psi_2_k)
# Potential step:
psi_1, psi_2 = sm.potential_evolution(psi_1, psi_2, -1j * dt, g1, g2,
def remove_global_phase(states):
gp = exp(-1j * angle(states[0]))
c = swapaxes(states, 0, 1)
states = multiply(gp.T, c).T
return states
def find_Bragg_disks_CUDA(
datacube,
probe,
corrPower=1,
sigma=2,
edgeBoundary=20,
minRelativeIntensity=0.005,
minAbsoluteIntensity=0.0,
relativeToPeak=0,
minPeakSpacing=60,
maxNumPeaks=70,
subpixel="multicorr",
upsample_factor=16,
filter_function=None,
name="braggpeaks_raw",
_qt_progress_bar=None,
batching=True,
):
"""
Finds the Bragg disks in all diffraction patterns of datacube by cross, hybrid, or
phase correlation with probe. When hist = True, returns histogram of intensities in
the entire datacube.
Args:
DP (ndarray): a diffraction pattern
probe (ndarray): the vacuum probe template, in real space.
corrPower (float between 0 and 1, inclusive): the cross correlation power. A
value of 1 corresponds to a cross correaltion, and 0 corresponds to a
phase correlation, with intermediate values giving various hybrids.
sigma (float): the standard deviation for the gaussian smoothing applied to
the cross correlation
edgeBoundary (int): minimum acceptable distance from the DP edge, in pixels
minRelativeIntensity (float): the minimum acceptable correlation peak intensity,
relative to the intensity of the brightest peak
relativeToPeak (int): specifies the peak against which the minimum relative
intensity is measured -- 0=brightest maximum. 1=next brightest, etc.
minPeakSpacing (float): the minimum acceptable spacing between detected peaks
maxNumPeaks (int): the maximum number of peaks to return
subpixel (str): Whether to use subpixel fitting, and which algorithm to use.
Must be in ('none','poly','multicorr').
* 'none': performs no subpixel fitting
* 'poly': polynomial interpolation of correlogram peaks (default)
* 'multicorr': uses the multicorr algorithm with DFT upsampling
upsample_factor (int): upsampling factor for subpixel fitting (only used when
subpixel='multicorr')
global_threshold (bool): if True, applies global threshold based on
minGlobalIntensity and metric
minGlobalThreshold (float): the minimum allowed peak intensity, relative to the
selected metric (0-1), except in the case of 'manual' metric, in which the
threshold value based on the minimum intensity that you want thresholder
out should be set.
metric (string): the metric used to compare intensities. 'average' compares peak
intensity relative to the average of the maximum intensity in each
diffraction pattern. 'max' compares peak intensity relative to the maximum
intensity value out of all the diffraction patterns. 'median' compares peak
intensity relative to the median of the maximum intensity peaks in each
diffraction pattern. 'manual' Allows the user to threshold based on a
predetermined intensity value manually determined. In this case,
minIntensity should be an int.
name (str): name for the returned PointListArray
filter_function (callable): filtering function to apply to each diffraction
pattern before peakfinding. Must be a function of only one argument (the
diffraction pattern) and return the filtered diffraction pattern. The
shape of the returned DP must match the shape of the probe kernel (but does
not need to match the shape of the input diffraction pattern, e.g. the filter
can be used to bin the diffraction pattern). If using distributed disk
detection, the function must be able to be pickled with by dill.
_qt_progress_bar (QProgressBar instance): used only by the GUI.
batching (bool): Whether to batch the FFT cross correlation steps.
Returns:
(PointListArray): the Bragg peak positions and correlation intensities
"""
# Make the peaks PointListArray
coords = [("qx", float), ("qy", float), ("intensity", float)]
peaks = PointListArray(coordinates=coords, shape=(datacube.R_Nx, datacube.R_Ny))
# check that the filtered DP is the right size for the probe kernel:
if filter_function:
assert callable(filter_function), "filter_function must be callable"
DP = (
datacube.data[0, 0, :, :]
if filter_function is None
else filter_function(datacube.data[0, 0, :, :])
)
assert np.all(
DP.shape == probe.shape
), "Probe kernel shape must match filtered DP shape"
# Get the probe kernel FT as a cupy array
probe_kernel_FT = cp.conj(cp.fft.fft2(cp.array(probe)))
bytes_per_pattern = probe_kernel_FT.nbytes
# get the maximal array kernel
# if probe_kernel_FT.dtype == 'float64':
# get_maximal_points = kernels['maximal_pts_float64']
# elif probe_kernel_FT.dtype == 'float32':
# get_maximal_points = kernels['maximal_pts_float32']
# else:
# raise TypeError("Maximal kernel only valid for float32 and float64 types...")
get_maximal_points = kernels["maximal_pts_float64"]
if get_maximal_points.max_threads_per_block < DP.shape[1]:
# naive blocks/threads will not work, figure out an OK distribution
blocks = ((np.prod(DP.shape) // get_maximal_points.max_threads_per_block + 1),)
threads = get_maximal_points.max_threads_per_block
else:
blocks = (DP.shape[0],)
threads = (DP.shape[1],)
if _qt_progress_bar is not None:
from PyQt5.QtWidgets import QApplication
t0 = time()
if batching:
# compute the batch size based on available VRAM:
max_num_bytes = cp.cuda.Device().mem_info[0]
# use a fudge factor to leave room for the fourier transformed data
# I have set this at 10, which results in underutilization of
# VRAM, because this yielded better performance in my testing
batch_size = max_num_bytes // (bytes_per_pattern * 10)
num_batches = datacube.R_N // batch_size + 1
print(f"Using {num_batches} batches of {batch_size} patterns each...")
for batch_idx in tqdmnd(
range(num_batches), desc="Finding Bragg disks in batches", unit="batch"
):
# the final batch may be smaller than the other ones:
probes_remaining = datacube.R_N - (batch_idx * batch_size)
this_batch_size = (
probes_remaining if probes_remaining < batch_size else batch_size
)
# allocate array for batch of DPs
# potential optimization: allocate this only once
batched_subcube = cp.zeros(
(this_batch_size, datacube.Q_Nx, datacube.Q_Ny), dtype=cp.float64
)
# fill in diffraction patterns, with filtering
for subbatch_idx in range(this_batch_size):
patt_idx = batch_idx * batch_size + subbatch_idx
rx, ry = np.unravel_index(patt_idx, (datacube.R_Nx, datacube.R_Ny))
batched_subcube[subbatch_idx, :, :] = cp.array(
datacube.data[rx, ry, :, :]
if filter_function is None
else filter_function(datacube.data[rx, ry, :, :]),
dtype=cp.float64,
)
# Perform the FFT and multiplication by probe_kernel on the batched array
batched_crosscorr = (
cufft.fft2(batched_subcube, overwrite_x=True)
* probe_kernel_FT[None, :, :]
)
# Iterate over the patterns in the batch and do the Bragg disk stuff
for subbatch_idx in range(this_batch_size):
patt_idx = batch_idx * batch_size + subbatch_idx
rx, ry = np.unravel_index(patt_idx, (datacube.R_Nx, datacube.R_Ny))
subFFT = batched_crosscorr[subbatch_idx]
ccc = cp.abs(subFFT) ** corrPower * cp.exp(1j * cp.angle(subFFT))
cc = cp.maximum(cp.real(cp.fft.ifft2(ccc)), 0)
_find_Bragg_disks_single_DP_FK_CUDA(
None,
None,
ccc=ccc,
cc=cc,
corrPower=corrPower,
sigma=sigma,
edgeBoundary=edgeBoundary,
minRelativeIntensity=minRelativeIntensity,
minAbsoluteIntensity=minAbsoluteIntensity,
relativeToPeak=relativeToPeak,
minPeakSpacing=minPeakSpacing,
maxNumPeaks=maxNumPeaks,
subpixel=subpixel,
upsample_factor=upsample_factor,
filter_function=filter_function,
peaks=peaks.get_pointlist(rx, ry),
get_maximal_points=get_maximal_points,
blocks=blocks,
threads=threads,
)
else:
# Loop over all diffraction patterns
for (Rx, Ry) in tqdmnd(
datacube.R_Nx,
datacube.R_Ny,
desc="Finding Bragg Disks",
unit="DP",
unit_scale=True,
):
if _qt_progress_bar is not None:
_qt_progress_bar.setValue(Rx * datacube.R_Ny + Ry + 1)
QApplication.processEvents()
DP = datacube.data[Rx, Ry, :, :]
_find_Bragg_disks_single_DP_FK_CUDA(
DP,
probe_kernel_FT,
corrPower=corrPower,
sigma=sigma,
edgeBoundary=edgeBoundary,
minRelativeIntensity=minRelativeIntensity,
minAbsoluteIntensity=minAbsoluteIntensity,
relativeToPeak=relativeToPeak,
minPeakSpacing=minPeakSpacing,
maxNumPeaks=maxNumPeaks,
subpixel=subpixel,
upsample_factor=upsample_factor,
filter_function=filter_function,
peaks=peaks.get_pointlist(Rx, Ry),
get_maximal_points=get_maximal_points,
blocks=blocks,
threads=threads,
)
t = time() - t0
print(
f"Analyzed {datacube.R_N} diffraction patterns in {t//3600}h {t % 3600 // 60}m {t % 60:.2f}s\n(avg. speed {datacube.R_N/t:0.4f} patterns per second)".format()
)
peaks.name = name
return peaks
def _estimate_fractional_delay(self, iq_0, iq_1, integer_delay, rate, fc):
'''
Returns fractional sampling time correction between channels in seconds.
First corrects integer sample lag to make estimating the fractional lag
tractable, then finds the slope of the phase of the cross-correlation by
linear regression to estimate the fractional sample delay.
Parameters
----------
iq_0, iq_1 : cusignal mapped, pinned array
GPU memory containing SDR IQ data from channels
integer_delay : int
estimated value from _estimate_integer_delay(), sec
rate : float
SDR sample rate, samples per second
fc : float
SDR center tuning frequency, Hz
Returns
-------
frac_delay, float
The fractional lag between the argument signals in seconds
'''
N = 8192
f0 = cp.fft.fft(cp.array(iq_0), n=N)
f1 = cp.fft.fft(cp.array(iq_1), n=N)
freqs = cp.fft.fftfreq(N, d=1 / rate) + fc
# Yes, there is a double negative, for mathematical clarity. Most eqns
# have -i, and we are "undoing" the delay, so -delay.
rot = cp.exp(-2j * cp.pi * freqs * -integer_delay)
# Integer sample correction as a phase rotation in frequency space
xcorr = f0 * cp.conj(f1 * rot)
# Prepare to fit residual phase gradient:
phases = cp.angle(xcorr)
# Due to receiver bandpass shape, edge frequencies have less power => less certain phase
# Assign weights accordingly
weights = cp.abs(xcorr)
weights /= cp.max(weights)
# Fit phase slope across band
# https://scipython.com/book/chapter-8-scipy/examples/weighted-and-non-weighted-least-squares-fitting/
def model(x, m, b):
# From "Reliable fitting of phase data without unwrapping by wrapping the fit model," Kramer et. al. 2012
return ((m * x + b) + np.pi) % (2. * np.pi) - np.pi
# initial guesses
p0 = [0, 0]
# fitting
popt, pcov = optimize.curve_fit(
model,
cp.asnumpy(freqs),
cp.asnumpy(phases),
p0,
sigma=1. / cp.asnumpy(weights),
absolute_sigma=False # not real sigmas, just weights
)
slope, intc = popt
frac_delay = slope
if np.abs(frac_delay) > 1 / rate:
fig = plt.figure(100)
ax = plt.axes()
ax.scatter(cp.asnumpy(freqs),
cp.asnumpy(phases),
alpha=0.1,
label='Calibration data: phase')
ax.plot(cp.asnumpy(freqs),
model(cp.asnumpy(freqs), slope, intc),
color='red',
label='Fit: phase slope')
ax.set_xlabel('Frequency (Hz)')
ax.set_ylabel('Phase (rad)')
ax.legend(loc='best', framealpha=0.4)
fig.show()
self.logger.warning(
'1st-pass delay calibration failed: fractional' +
' sample time correction, |{}| > 1/sample rate, {} '.format(
frac_delay, 1 / rate))
return frac_delay
def grad(
self,
data,
psi,
scan,
prb,
first,
dpsi,
gradpsi,
testpsi,
piter,
model='gaussian',
recover_prb=False,
):
"""Conjugate gradients for ptychography.
Parameters
----------
model : str gaussian or poisson
The noise model to use for the gradient.
piter : int
The number of gradient steps to take.
recover_prb : bool
Whether to recover the probe or assume the given probe is correct.
"""
assert prb.ndim == 3, "prb needs 3 dimensions, not %d" % prb.ndim
print("# congujate gradient parameters\n"
"iteration, step size object, step size probe, function min"
) # csv column headers
gammaprb=0
for i in range(piter):
# 1) object retrieval subproblem with fixed probe
# forward operator
self.fpsi.fill(0+0j)
self.fpsi = self.fwd_ptycho(self.fpsi, psi, scan, prb)
gradpsi0 = gradpsi
# take gradient
if model == 'gaussian':
gradpsi = self.adj_ptycho(
self.fpsi - cp.sqrt(data) * cp.exp(1j * cp.angle(self.fpsi)),
scan,
prb,
) / (cp.max(cp.abs(prb))**2)
elif model == 'poisson':
gradpsi = self.adj_ptycho(
self.fpsi - data * self.fpsi / (cp.abs(self.fpsi)**2 + 1e-32),
scan,
prb,
) / (cp.max(cp.abs(prb))**2)
gammapsi = 0.25
if (recover_prb):
# 2) probe retrieval subproblem with fixed object
# forward operator
fprb = self.fwd_ptycho(psi, scan, prb)
# take gradient
if model == 'gaussian':
gradprb = self.adj_ptycho_prb(
fprb - cp.sqrt(data) * cp.exp(1j * cp.angle(fprb)),
scan,
psi,
) / cp.max(cp.abs(psi))**2 / self.nscan
elif model == 'poisson':
gradprb = self.adj_ptycho_prb(
fprb - data * fprb / (cp.abs(fprb)**2 + 1e-32),
scan,
psi,
) / cp.max(cp.abs(psi))**2 / self.nscan
# Dai-Yuan direction
if (i == 0):
dprb = -gradprb
else:
dprb = -gradprb + (
cp.linalg.norm(gradprb)**2 /
(cp.sum(cp.conj(dprb) * (gradprb - gradprb0))) * dprb)
gradprb0 = gradprb
# line search
fdprb = self.fwd_ptycho(psi, scan, dprb)
gammaprb = self.line_search(minf, fprb, fdprb)
# update prb
prb = prb + gammaprb * dprb
# # check convergence
# if (np.mod(i, 8) == 0):
# fpsi = self.fwd_ptycho(psi, scan, prb)
# print("%4d, %.3e, %.3e, %.7e" %
# (i, gammapsi, gammaprb, minf(fpsi)))
return {
'psi': psi,
'prb': prb,
'dpsi': dpsi,
'gradpsi': gradpsi,
'testpsi': testpsi,
'gammapsi': gammapsi,
}
def angle(arr):
return cp.angle(arr)
if loading_vortex_pos:
with h5py.File('vortex_positions/vortex_pos_uniform.hdf5',
'r') as data:
vort_pos = iter(data['positions'][...])
else:
vort_pos = include.phase_imprinting.get_positions(
N_vort, 2 * xi, len_x, len_y) # Generator of vortex positions
theta = include.phase_imprinting.get_phase(N_vort, vort_pos, X,
Y) # Phase imprinting
# Construct wavefunction and related:
psi = cp.sqrt(n_0) * cp.exp(1j * theta)
atom_num = dx * dy * cp.sum(cp.abs(psi)**2)
theta_fix = cp.angle(psi)
psi_k = cp.fft.fft2(psi)
# ------------------------------------------------------------------------------------------------------------------
# Imaginary time evolution
# ------------------------------------------------------------------------------------------------------------------
for i in range(2000):
# Kinetic energy:
psi_k *= cp.exp(-0.25 * dt * (Kx**2 + Ky**2))
# Backward FFT:
psi = cp.fft.ifft2(psi_k)
# Interaction term:
psi *= cp.exp(-dt * (c0 * cp.abs(psi)**2))
def WISHrun(self, y0: np.ndarray, SLM: np.ndarray, delta3: float, delta4: float, N_os: int, N_iter: int,\
N_batch: int, plot: bool=True):
"""
Runs the WISH algorithm using a Gerchberg Saxton loop for phase retrieval.
:param y0: Target modulated amplitudes in the sensor plane
:param SLM: SLM modulation patterns
:param delta3: Apparent sampling size of the SLM as seen from the sensor plane
:param delta4: Sampling size of the sensor plane
:param N_os: Number of observations per image
:param N_iter: Maximal number of Gerchberg Saxton iterations
:param N_batch: Number of batches (modulations)
:param plot: If True, plots the advance of the retrieval every 10 iterations
:return u4_est, idx_converge: Estimated field of size (N,N) and the convergence indices to check convergence
speed
"""
wvl = self.wavelength
z3 = self.z
## parameters
N = y0.shape[0]
k = 2 * np.pi / wvl
#u3_batch = np.zeros((N, N, N_os), dtype=complex) # store all U3 gpu
#u4 = np.zeros((N, N, N_os), dtype=complex) # gpu
#y = np.zeros((N, N, N_os), dtype=complex) # store all U3 gpu
u3_batch = cp.zeros((N, N, N_os),
dtype=cp.complex64) # store all U3 gpu
u4 = cp.zeros((N, N, N_os), dtype=cp.complex64) # gpu
y = cp.zeros((N, N, N_os), dtype=cp.complex64) # store all U3 gpu
## initilize a3
k = 2 * np.pi / wvl
xx = cp.linspace(0, N - 1, N,
dtype=cp.float) - (N / 2) * cp.ones(N, dtype=cp.float)
yy = cp.linspace(0, N - 1, N,
dtype=cp.float) - (N / 2) * cp.ones(N, dtype=cp.float)
X, Y = float(delta4) * cp.meshgrid(
xx, yy)[0], float(delta4) * cp.meshgrid(xx, yy)[1]
R = cp.sqrt(X**2 + Y**2)
Q = cp.exp(1j * (k / (2 * z3)) * R**2)
for ii in range(N_os):
#SLM_batch = SLM[:,:, ii]
SLM_batch = cp.asarray(SLM[:, :, ii])
y0_batch = y0[:, :, ii]
#u3_batch[:,:, ii] = self.frt(y0_batch, delta4, -z3) * np.conj(SLM_batch) #y0_batch gpu
#u3_batch[:,:, ii] = self.frt_gpu(cp.asarray(y0_batch), delta4, -z3) * cp.conj(SLM_batch) #y0_batch gpu
u3_batch[:, :, ii] = self.frt_gpu_s(
cp.asarray(y0_batch) / Q, delta4, -z3) * cp.conj(
SLM_batch) #y0_batch gpu
#u3 = np.mean(u3_batch, 2) # average it
u3 = cp.mean(u3_batch, 2)
## Recon run : GS loop
idx_converge = np.empty(N_iter)
for jj in range(N_iter):
sys.stdout.write(f"\rGS iteration {jj+1}")
sys.stdout.flush()
#u3_collect = np.zeros(u3.shape, dtype=complex)
u3_collect = cp.zeros(u3.shape, dtype=cp.complex64)
idx_converge0 = np.empty(N_batch)
for idx_batch in range(N_batch):
# put the correct batch into the GPU (no GPU for now)
#SLM_batch = SLM[:,:, int(N_os * idx_batch): int(N_os * (idx_batch+1))]
#y0_batch = y0[:,:, int(N_os * idx_batch): int(N_os * (idx_batch+1))]
SLM_batch = cp.asarray(
SLM[:, :,
int(N_os * idx_batch):int(N_os * (idx_batch + 1))])
y0_batch = cp.asarray(
y0[:, :,
int(N_os * idx_batch):int(N_os * (idx_batch + 1))])
for _ in range(N_os):
#u4[:,:,_] = self.frt(u3 * SLM_batch[:,:,_], delta3, z3) # U4 is the field on the sensor
u4[:, :,
_] = self.frt_gpu_s(u3 * SLM_batch[:, :, _], delta3,
z3) # U4 is the field on the sensor
y[:, :,
_] = y0_batch[:, :, _] * cp.exp(1j * cp.angle(
u4[:, :, _])) # force the amplitude of y to be y0
#u3_batch[:,:,_] = self.frt(y[:,:,_], delta4, -z3) * np.conj(SLM_batch[:,:,_])
u3_batch[:, :, _] = self.frt_gpu_s(
y[:, :, _], delta4, -z3) * cp.conj(SLM_batch[:, :, _])
#u3_collect = u3_collect + np.mean(u3_batch, 2) # collect(add) U3 from each batch
u3_collect = u3_collect + cp.mean(
u3_batch, 2) # collect(add) U3 from each batch
#idx_converge0[idx_batch] = np.mean(np.mean(np.mean(y0_batch,1),0)/np.sum(np.sum(np.abs(np.abs(u4)-y0_batch),1),0))
#idx_converge0[idx_batch] = cp.asnumpy(cp.mean(cp.mean(cp.mean(y0_batch,1),0)/cp.sum(cp.sum(cp.abs(cp.abs(u4)-y0_batch),1),0)))
# convergence index matrix for each batch
idx_converge0[idx_batch] = cp.linalg.norm(
cp.abs(u4) - y0_batch) / cp.linalg.norm(y0_batch)
u3 = (u3_collect / N_batch) # average over batches
idx_converge[jj] = np.mean(idx_converge0) # sum over batches
sys.stdout.write(f" (convergence index : {idx_converge[jj]})")
#u4_est = self.frt(u3, delta3, z3)
u4_est = cp.asnumpy(self.frt_gpu_s(u3, delta3, z3) * Q)
if jj % 10 == 0 and plot:
plt.close('all')
fig = plt.figure(0)
fig.suptitle(f'Iteration {jj}')
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)
im = ax1.imshow(np.abs(u4_est), cmap='viridis')
ax1.set_title('Amplitude')
ax2.imshow(np.angle(u4_est), cmap='viridis')
ax2.set_title('Phase')
fig1 = plt.figure(1)
ax = fig1.gca()
ax.plot(np.arange(0, jj, 1), idx_converge[0:jj], marker='o')
ax.set_xlabel('Iterations')
ax.set_ylabel('Convergence estimator')
ax.set_title('Convergence curve')
plt.show()
time.sleep(2)
# exit if the matrix doesn 't change much
if jj > 1:
if cp.abs(idx_converge[jj] -
idx_converge[jj - 1]) / idx_converge[jj] < 1e-4:
print('\nConverged. Exit the GS loop ...')
#idx_converge = idx_converge[0:jj]
idx_converge = cp.asnumpy(idx_converge[0:jj])
break
return u4_est, idx_converge
# Load probe
probe_amp = dxchange.read_tiff('data/probe_amp.tiff')
probe_angle = dxchange.read_tiff('data/probe_angle.tiff')
prb = probe_amp * np.exp(1j * probe_angle)
# Load scan positions
scan = np.load('data/scan.npy')
plt.plot(scan[1], scan[0], 'r.')
plt.savefig(f'data/scan.png')
# copy to gpu
psi = cp.array(psi)
prb = cp.array(prb)
scan = cp.array(scan)
with ptychodistrib.SolverPtycho(nz, n, nscan, ndet, nprb, 1) as pslv:
# data = ||FQpsi||^2
data = cp.abs(pslv.fwd_ptycho(psi, prb, scan))**2
# gradient solver
psi = cp.ones_like(psi)
niter = 32
psi = pslv.grad_ptycho(data, psi, prb, scan, None, -1, niter)
dxchange.write_tiff(cp.angle(psi).get(),
'rec/object_angle.tiff',
overwrite=True)
dxchange.write_tiff(cp.abs(psi).get(),
'rec/object_amp.tiff',
overwrite=True)
def reconstruct_alt(imgs,
discs,
hres_size,
row,
n_iters=1,
o_f_init=None,
del_1=1000,
del_2=1,
round_values=True,
plot_per_frame=False,
show_interval=None,
subtract_bg=False,
out_path=None):
"""The main reconstruction algorithm. Adapted from Tian et. al."""
# Put input images on GPU, estimate background noise
imgs = [cp.array(img) for img in imgs]
bgs = get_bg(imgs) if subtract_bg else cp.zeros(len(imgs))
IMAGESIZE = imgs[0].shape[0]
CUTOFF_FREQ_px = get_cutoff(row)
FRAMES = len(imgs)
orig = IMAGESIZE // 2 - 1 # Low-res origin
lres_size = (IMAGESIZE, IMAGESIZE)
m1, n1 = lres_size
m, n = hres_size
losses = [] # Reconstruction Loss
convs = [] # Inverse Convergence index
# Initial high-res guess
if lres_size == hres_size: # Initialize with ones
# Use old algorithm
F = lambda x: cp.fft.fftshift(cp.fft.fft2(x))
Ft = lambda x: cp.fft.ifft2(cp.fft.ifftshift(x))
o = cp.ones(hres_size)
o_f = F(o)
elif o_f_init is not None: # Initialize with given initialization
F = lambda x: cp.fft.fftshift(cp.fft.fft2(cp.fft.ifftshift(x)))
Ft = lambda x: cp.fft.fftshift(cp.fft.ifft2(cp.fft.ifftshift(x)))
o = cp.zeros_like(o_f_init)
o_f = o_f_init
else: # Intialize with resized first frame from imgs
F = lambda x: cp.fft.fftshift(cp.fft.fft2(cp.fft.ifftshift(x)))
Ft = lambda x: cp.fft.fftshift(cp.fft.ifft2(cp.fft.ifftshift(x)))
o = cp.sqrt(
cp.array(cv2.resize(cp.asnumpy(imgs[0] - bgs[0]), hres_size)))
o_f = Ft(o)
# Pupil Function
p = cp.zeros(lres_size)
p = cp.array(cv2.circle(cp.asnumpy(p), (orig, orig), CUTOFF_FREQ_px, 1,
-1))
ctf = p.copy() # Ideal Pupil, for filtering later on
# Main Loop
log = tqdm(
total=n_iters,
desc=f'Starting...',
bar_format=
'{percentage:3.0f}% [{elapsed}<{remaining} ({rate_inv_fmt})]{bar}{desc}',
leave=False,
ascii=True)
for j in range(n_iters):
conv = [] # Convergence Index
for i in range(FRAMES):
if discs[i] == 0: # Empty frame
continue
# Get k0x, k0y and hence, shifting values
k0x, k0y = discs[i]
# Construct auxillary functions for the set of LEDs (= 1, here)
if hres_size == lres_size:
shift_x, shift_y = [
-round(k0x - orig), -round(k0y - orig)
] if round_values else [-(k0x - orig), -(k0y - orig)]
if not round_values:
o_f_i = FourierShift2D(o_f,
[shift_x, shift_y]) # O_i(k - k_m)
else:
o_f_i = cp.roll(o_f, int(shift_y), axis=0)
o_f_i = cp.roll(o_f_i, int(shift_x), axis=1)
yl, xl = 0, 0 # To reduce code later on
else: # Output size larger than individual frames
_orig = hres_size[0] // 2 - 1
del_x, del_y = k0x - orig, k0y - orig
x, y = round(_orig - del_x), round(_orig - del_y)
yl = int(y - m1 // 2)
xl = int(x - n1 // 2)
assert xl > 0 and yl > 0, 'Both should be > 0'
o_f_i = o_f[yl:yl + n1, xl:xl + m1].copy()
psi_k = o_f_i * p * ctf #DEBUG: REPLACE * ctf with * p
# Plot outputs after each frame, for debugging
if plot_per_frame:
o_i = Ft(o_f_i * p)
plt.figure(figsize=(10, 2))
plt.subplot(161)
plt.imshow(cp.asnumpy(correct(abs(o_i))))
plt.title(f'$I_{{l}}({i})$')
opts() #DEBUG
plt.subplot(162)
plt.imshow(
cp.asnumpy(
cv2.convertScaleAbs(
cp.asnumpy(20 * cp.log(1 + abs(o_f_i * p))))))
plt.title(f'$S_{{l}}({i})$')
opts() #DEBUG
# Impose intensity constraint and update auxillary function
psi_r = F(psi_k) #DEBUG: CHANGE BACK TO F
# Low-res estimate obtained from our reconstruction
I_l = abs(psi_r) if lres_size != hres_size else abs(psi_r)
# Subtract background noise and clip values to avoid NaN
I_hat = cp.clip(imgs[i] - bgs[i], a_min=0)
phi_r = cp.sqrt(I_hat / (cp.abs(psi_r)**2)) * psi_r
phi_k = Ft(phi_r) #DEBUG: CHANGE BACK TO Ft
# Update object and pupil estimates
if hres_size == lres_size:
if not round_values:
p_i = FourierShift2D(p, [-shift_x, -shift_y]) # P_i(k+k_m)
else:
p_i = cp.roll(p, int(-shift_y), axis=0)
p_i = cp.roll(p_i, int(-shift_x), axis=1)
if not round_values:
phi_k_i = FourierShift2D(
phi_k, [-shift_x, -shift_y]) # Phi_m_i(k+k_m)
else:
phi_k_i = cp.roll(phi_k, int(-shift_y), axis=0)
phi_k_i = cp.roll(phi_k_i, int(-shift_x), axis=1)
else: # Output size larger than individual frames
p_i = p.copy()
phi_k_i = phi_k.copy()
## O_{i+1}(k)
temp = o_f[yl:yl + n1, xl:xl + m1].copy() + ( cp.abs(p_i) * cp.conj(p_i) * (phi_k_i - o_f[yl:yl + n1, xl:xl + m1].copy() * p_i) ) / \
( cp.abs(p).max() * (cp.abs(p_i) ** 2 + del_1) )
## P_{i+1}(k)
p = p + ( cp.abs(o_f_i) * cp.conj(o_f_i) * (phi_k - o_f_i * p) ) / \
( cp.abs(o_f[yl:yl + n1, xl:xl + m1].copy()).max() * (cp.abs(o_f_i) ** 2 + del_2) )
o_f[yl:yl + n1, xl:xl + m1] = temp.copy()
###### Using F here instead of Ft to get upright image
o = F(o_f) if lres_size != hres_size else Ft(o_f)
######
if plot_per_frame:
plt.subplot(163)
plt.imshow(cp.asnumpy(cp.mod(ctf * cp.angle(p), 2 * cp.pi)))
plt.title(f'P({i})')
opts() #DEBUG
plt.subplot(164)
plt.imshow(cp.asnumpy(correct(abs(o))))
plt.title(f'$I_{{h}}({i})$')
opts() #DEBUG
plt.subplot(165)
plt.imshow(cp.asnumpy(correct(cp.angle(o))))
plt.title(f'$\\theta(I_{{h}}({i}))$')
opts() #DEBUG
plt.subplot(166)
plt.imshow(cp.asnumpy(show(cp.asnumpy(o_f))))
plt.title(f'$S_{{h}}({i})$')
opts()
plt.show() #DEBUG
c = inv_conv_idx(I_l, imgs[i])
conv.append(c)
if not plot_per_frame and (show_interval is not None
and j % show_interval == 0):
o_i = Ft(o_f_i * p) #DEBUG
plt.figure(figsize=(10, 2))
plt.subplot(161)
plt.imshow(cp.asnumpy(correct(abs(o_i))))
plt.title(f'$I_{{l}}({i})$')
opts() #DEBUG
plt.subplot(162)
plt.imshow(
cp.asnumpy(
cv2.convertScaleAbs(
cp.asnumpy(20 * cp.log(1 + abs(o_f_i * p))))))
plt.title(f'$S_{{l}}({i})$')
opts() #DEBUG
plt.subplot(163)
plt.imshow(cp.asnumpy(cp.mod(ctf * cp.angle(p), 2 * cp.pi)))
plt.title(f'P({i})')
opts() #DEBUG
plt.subplot(164)
plt.imshow(cp.asnumpy(correct(abs(o))))
plt.title(f'$I_{{h}}({i})$')
opts() #DEBUG
plt.subplot(165)
plt.imshow(cp.asnumpy(correct(cp.angle(o))))
plt.title(f'$\\theta(I_{{h}}({i}))$')
opts() #DEBUG
plt.subplot(166)
plt.imshow(
cp.asnumpy(
cv2.convertScaleAbs(cp.asnumpy(20 *
cp.log(1 + abs(o_f))))))
plt.title(f'$S_{{h}}({i})$')
opts()
plt.show() #DEBUG
loss = metric_norm(imgs, o_f_i, p)
losses.append(loss)
conv = float(sum(conv) / len(conv))
convs.append(conv)
log.set_description_str(
f'[Iteration {j + 1}] Convergence Loss: {cp.asnumpy(conv):e}')
log.update(1)
scale = 7
plt.figure(figsize=(3 * scale, 4 * scale))
plt.subplot(421)
plt.plot(cp.asnumpy(cp.arange(len(losses))),
cp.asnumpy(cp.clip(cp.array(losses), a_min=None, a_max=1e4)),
'b-')
plt.title('Loss Curve')
plt.ylabel('Loss Value')
plt.xlabel('Iteration')
plt.subplot(422)
plt.plot(cp.asnumpy(cp.arange(len(convs))),
cp.asnumpy(cp.clip(cp.array(convs), a_min=None, a_max=1e14)),
'b-')
plt.title('Convergence Index Curve')
plt.ylabel('Convergence Index')
plt.xlabel('Iteration')
amp = cp.array(cv2.resize(
read_tiff(row.AMPLITUDE.values[0])[0], hres_size))
phase = cp.array(cv2.resize(read_tiff(row.PHASE.values[0])[0], hres_size))
plt.subplot(434)
plt.title(f'amplitude (Scaled up from {lres_size})')
plt.imshow(cp.asnumpy(to_uint8(amp)))
opts()
plt.subplot(435)
plt.title(f'phase (Scaled up from {lres_size})')
plt.imshow(cp.asnumpy(to_uint8(phase)))
plt.subplot(436)
rec = abs(cp.sqrt(amp) * cp.exp(1j * phase))
plt.title(f'Ground Truth (Scaled up from {lres_size})')
plt.imshow(cp.asnumpy(to_uint8(rec)))
plt.subplot(437)
plt.title('Reconstruction Amplitude')
amp = abs(o)
if lres_size == hres_size:
amp = correct(amp)
plt.imshow(cp.asnumpy(to_uint8((amp))))
plt.subplot(438)
plt.title('Reconstruction Phase')
phase = cp.angle(o)
if lres_size == hres_size:
phase = correct(phase)
plt.imshow(cp.asnumpy(to_uint8(phase)))
plt.subplot(439)
plt.title('Reconstructed Image')
rec = abs(cp.sqrt(amp) * cp.exp(1j * phase))
plt.imshow(cp.asnumpy(to_uint8(rec)))
plt.subplot(427)
plt.title(f'Recovered Pupil')
p_show = cp.mod(ctf * cp.angle(p), 2 * cp.pi)
p_show = (p_show / p_show.max() * 255).astype(np.uint8)
plt.imshow(cp.asnumpy(p_show), cmap='nipy_spectral')
plt.subplot(428)
plt.title(f'Raw frames\' mean (Scaled up from {lres_size})')
plt.imshow(cv2.resize(cp.asnumpy(cp.array(imgs).mean(axis=0)), hres_size))
if out_path is None:
plt.show()
else:
plt.savefig(out_path, bbox_inches='tight')
plt.close('all')
# Ignore early noise and print where the error is lowest
if n_iters > 10:
it = cp.argmin(cp.array(convs[10:])) + 11
if out_path is not None:
print(f'Convergence index lowest at {it}th iteration.')
else:
it = cp.argmin(cp.array(convs)) + 1
if out_path is not None:
print(f'Convergence index lowest at {it}th iteration.')
if lres_size == hres_size:
o = correct(o)
return o, p, it
def order(self):
o = cp.mean(cp.exp(self.phase * np.complex(0, 1)))
return abs(o), cp.angle(o)
def reconstruct_v2(
imgs,
discs,
row,
hres_size,
n_iters=1,
do_fil=False,
denoise=False,
crop_our_way=True,
plot=True,
adaptive_noise=1,
adaptive_pupil=1,
adaptive_img=1,
alpha=1,
# delta_img=0.1,
# delta_pupil=1e-6,
delta_img=10,
delta_pupil=1e-4,
eps=1e-9,
calibrate_freq_pos=False,
out_path=None):
"""
Adapted From Aidukas et al. (2018)
"""
# Define basic parameters
dim_segment_interp = hres_size[0]
dim_segment = imgs[0].shape[0]
orig = dim_segment // 2 - 1
lres_size = (dim_segment, dim_segment)
clahe_p = cv2.createCLAHE(clipLimit=3,
tileGridSize=(10, 10)) # For contrast Enhancment
clahe_a = cv2.createCLAHE(clipLimit=1.5,
tileGridSize=(10, 10)) # For contrast Enhancment
# Initialize high-res estimate (freq. domain)
high_res_freq_estimate = o_f4(imgs, hres_size, row=row, do_fil=do_fil)
# Pupil and aperture
CUTOFF_FREQ_px = get_cutoff(row)
pupil = cp.zeros((dim_segment, dim_segment))
pupil = cp.array(
cv2.circle(cp.asnumpy(pupil), (orig, orig), int(CUTOFF_FREQ_px), 1,
-1))
aperture = pupil.copy()
# Progress bar
log = tqdm(
total=n_iters,
desc=f'Working...',
bar_format='{percentage:3.0f}% [{elapsed}<{remaining} ({rate_inv_fmt})]'
'{bar}{desc}',
leave=False,
)
convs = []
# Main Loop
for iteration_number in range(n_iters):
conv = []
for i in range(len(imgs)): # Iterate over all the images
freq_pos = discs[i]
if freq_pos == 0: # Skip empty frames
continue
if not crop_our_way: # Crop with formula used by original author
x1 = int((dim_segment_interp - dim_segment) / 2 + freq_pos[1])
x2 = int((dim_segment_interp + dim_segment) / 2 + freq_pos[1])
y1 = int((dim_segment_interp - dim_segment) / 2 + freq_pos[0])
y2 = int((dim_segment_interp + dim_segment) / 2 + freq_pos[0])
low_res_freq_estimate = cp.copy(high_res_freq_estimate[x1:x2,
y1:y2])
else: # Use our own (Earlier doesn't seem to work)
# High-res origin
_orig = hres_size[0] // 2 - 1
k0x, k0y = freq_pos # Wavevectors
del_x, del_y = k0x - orig, k0y - orig # Distance from origin in low-res
x, y = round(_orig + del_x), round(
_orig + del_y) # Distance from origin in high-res
xl = int(y - dim_segment // 2)
yl = int(x - dim_segment // 2)
x1, x2 = xl, xl + dim_segment
y1, y2 = yl, yl + dim_segment
# Low-res estimate (Fourier domain)
low_res_freq_estimate = cp.copy(high_res_freq_estimate[x1:x2,
y1:y2])
# show_f(low_res_freq_estimate, title='low_res_freq_estimate')
# Filter with pupil and aperture
low_res_freq_estimate_filtered = pupil * aperture * low_res_freq_estimate
# show_f(low_res_freq_estimate_filtered, title='low_res_freq_estimate_filtered')
# Back to real domain
low_res_estimate_filtered = cp_inverse_fourier(
low_res_freq_estimate_filtered)
I_l = abs(low_res_estimate_filtered)
c = inv_conv_idx(I_l, imgs[i])
conv.append(c)
# Measured intensity values
experimental_amp = cp.copy(imgs[i])
# Skipping sparse sampling (setting bayer = 1)
bayer = 1
if denoise: # Denoise measurements (not required for our case)
noise = cp.abs(
cp.mean(experimental_amp) -
cp.mean(cp.abs(low_res_estimate_filtered)**2 *
bayer)) * adaptive_noise
denoised_image = (cp.abs(experimental_amp) - noise) * (
(cp.abs(experimental_amp) - noise) > 0)
denoised_image = cp.sqrt(denoised_image) * (bayer) + cp.abs(
low_res_estimate_filtered) * cp.abs(1 - bayer)
else:
denoised_image = cp.sqrt(cp.abs(experimental_amp))
# Update low-res estimate with measured intensity values
low_res_updated = cp.abs(denoised_image) \
* low_res_estimate_filtered / (cp.abs(low_res_estimate_filtered)+ eps)
low_res_freq_updated = cp_forward_fourier(low_res_updated)
# show_f(low_res_freq_updated, 'low_res_freq_updated')
# Update high-res estimate (Fourier domain)
temp = aperture * (low_res_freq_updated -
low_res_freq_estimate_filtered)
Omax = cp.max(cp.abs(high_res_freq_estimate))
high_res_freq_estimate[x1:x2, y1:y2] = low_res_freq_estimate + alpha*adaptive_pupil**(iteration_number+1) \
* cp.abs(pupil) * cp.conj(pupil) * temp \
/ (cp.max(cp.abs(pupil)) * (cp.abs(pupil)**2 + delta_img))
# show_f(high_res_freq_estimate, 'high_res_freq_estimate')
# Update Pupil
pupil = pupil + alpha*adaptive_pupil**(iteration_number+1) \
* aperture * cp.abs(low_res_freq_estimate) * cp.conj(low_res_freq_estimate) * temp \
/ (Omax * (cp.abs(low_res_freq_estimate)**2 + delta_pupil))
# Spectral Correlation Calibration
if calibrate_freq_pos:
# print(discs[i]) #DEBUG
discs[i] = spectral_correlation_calibration_GPU(high_res_freq_estimate, denoised_image,\
pupil, freq_pos, dim_segment_interp, dim_segment)
# print(discs[i]) #DEBUG
# Out of inner loop
if plot:
conv = float(sum(conv) / len(conv))
convs.append(conv)
# Update progeress bar
log.update()
# Get high-res estimate
high_res_estimate = inverse_fourier(high_res_freq_estimate)
if plot:
log.set_description_str('Drawing figure...')
# Create Figure
scale = 7
plt.figure(figsize=(2 * scale, 4 * scale))
amp = cp.array(
cv2.resize(read_tiff(row.AMPLITUDE.values[0])[0],
hres_size,
interpolation=cv2.INTER_NEAREST))
phase = cp.array(
cv2.resize(read_tiff(row.PHASE.values[0])[0],
hres_size,
interpolation=cv2.INTER_NEAREST))
plt.subplot(421)
plt.title(f'Intensity (Scaled up from {lres_size})')
plt.imshow(cp.asnumpy(to_uint8(amp)), interpolation='none')
plt.colorbar()
plt.subplot(422)
plt.title(f'Phase (Scaled up from {lres_size})')
plt.imshow(cp.asnumpy(to_uint8(phase)), interpolation='none')
plt.colorbar()
plt.subplot(423)
plt.title('Reconstruction Intensity')
amp = cp.abs(cp.rot90(high_res_estimate, 0))
plt.imshow(cp.asnumpy(to_uint8((amp))), interpolation='none')
plt.colorbar()
plt.subplot(424)
plt.title('Reconstruction Phase')
phase = cp.angle(cp.rot90(high_res_estimate, 0))
plt.imshow(cp.asnumpy(to_uint8(phase)), interpolation='none')
plt.colorbar()
plt.subplot(425)
# plt.title('Reconstruction Intensity (Contrast Corrected)')
# amp = cp.abs(cp.rot90(high_res_estimate, 0))
# # plt.imshow(cv2.equalizeHist(cp.asnumpy(to_uint8(phase))))
# plt.imshow(clahe_a.apply(cp.asnumpy(to_uint8(amp))))
plt.title(f'Mean of Measured Intensity (Scaled up from {lres_size})')
plt.imshow(cv2.resize(cp.asnumpy(to_uint8(
cp.array(imgs).mean(axis=0))),
hres_size,
interpolation=cv2.INTER_NEAREST),
interpolation='none')
plt.colorbar()
plt.subplot(426)
plt.title('Reconstruction Phase (Contrast Corrected)')
phase = cp.angle(cp.rot90(high_res_estimate, 0))
# plt.imshow(cv2.equalizeHist(cp.asnumpy(to_uint8(phase))))
plt.imshow(clahe_p.apply(cp.asnumpy(to_uint8(phase))),
interpolation='none')
plt.colorbar()
plt.subplot(427)
plt.title(f'Recovered Pupil Intensity')
p_show = cp.asnumpy(cp.abs(pupil))
plt.imshow(p_show, interpolation='none')
plt.colorbar()
plt.subplot(428)
plt.title(f'Recovered Pupil Phase')
p_show = cp.asnumpy(cp.angle(pupil))
plt.imshow(p_show, interpolation='none')
plt.colorbar()
# plt.subplot(4,3,12)
# plt.title(f'Raw frames\' mean (Scaled up from {lres_size})')
# plt.imshow(cv2.resize(cp.asnumpy(cp.array(imgs).mean(axis=0)), hres_size, interpolation=cv2.INTER_NEAREST))
if (out_path is None) and plot:
plt.show()
else: # Save figure to the specified path
plt.savefig(out_path, bbox_inches='tight')
plt.close('all')
log.set_description_str('Done.')
log.close()
return high_res_estimate, pupil
pantie = strs[0] + '.png'
fmin = np.float64(strs[1])
fmax = np.float64(strs[2][:-4])
print('Process: ' + pantie)
fpantie = Image.open('./converted/Fourier/' + fname)
fimg_origin = np.fft.ifftshift(np.exp(np.array(fpantie,dtype=np.float64) / np.max(np.array(fpantie)) * fmax + fmin))
mask = (np.array(Image.open('./dream/' + pantie))[:, :, 3] > 0).astype(np.float32)[:, :, None]
[r, c, d] = fimg_origin.shape
img = np.fft.ifft2(fimg_origin * np.exp(2j * np.pi * np.random.rand(r, c, d)), axes=(0, 1)) # Initialize the phase
if args.hio:
mask_inv = 1 - mask
previous_img = np.copy(img)
for i in range(args.itr):
if args.hio:
fimg = np.fft.fft2((np.abs(img) * mask + (np.abs(previous_img) - args.beta * np.abs(img)) * mask_inv) * np.exp(1j * np.angle(img)), axes=(0, 1))
previous_img = np.copy(img)
else: # Error reduction
fimg = np.fft.fft2(np.abs(img) * mask * np.exp(1j * np.angle(img)), axes=(0, 1))
img = np.fft.ifft2(fimg_origin * np.exp(1j * np.angle(fimg)), axes=(0, 1))
if np.max(np.array(fpantie)) == 255:
img = (norm_img(np.abs(img)) * 255).astype(np.uint8)
else:
img = (norm_img(np.abs(img)) * 65535).astype(np.uint16)
if args.gpu:
img = np.asnumpy(img)
if args.hio:
outname = './converted/Fourier_inv/' + pantie[:-4] + '_hio.png'
else:
outname = './converted/Fourier_inv/' + pantie
restored_pantie = Image.fromarray(img)
