mirror of
https://github.com/c-sooyoung/fold_slice.git
synced 2026-09-17 21:39:08 +09:00
385 lines
18 KiB
Matlab
385 lines
18 KiB
Matlab
%LSQML least squared maximum likelihood solver
|
|
%
|
|
%[self, cache, fourier_error] = LSQML(self,par,cache,fourier_error,iter)
|
|
%
|
|
% ** self structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
|
|
% ** par structure containing parameters for the engines
|
|
% ** cache structure with precalculated values to avoid unnecessary overhead
|
|
% ** fourier_error array [Npos,1] containing evolution of reconstruction error
|
|
% ** iter current iteration number
|
|
% returns:
|
|
% ++ self self-like structure with final reconstruction
|
|
% ++ cache structure with precalculated values to avoid unnecessary overhead
|
|
% ++ fourier_error array [Npos,1] containing evolution of reconstruction error
|
|
|
|
function [self, cache, fourier_error] = LSQML(self,par,cache,fourier_error,iter)
|
|
import engines.GPU.shared.*
|
|
import math.*
|
|
import utils.*
|
|
import plotting.*
|
|
import engines.GPU.GPU_wrapper.*
|
|
import engines.GPU.LSQML.*
|
|
|
|
assert( ~(par.Nscans > 1 && par.object_modes > 1), 'Multiobject + multiscan not supported')
|
|
|
|
% define some useful variables
|
|
object_modes = size(self.object,1);
|
|
object_upd_sum = cell(object_modes,par.Nlayers);
|
|
obj_illum_sum = cell(object_modes,par.Nlayers);
|
|
obj_proj = cell(par.object_modes,1);
|
|
apply_subpx_shift(); % reset persistent values
|
|
|
|
for ll = 1:object_modes
|
|
for layer = 1:par.Nlayers
|
|
obj_illum_sum{ll,layer} = Gzeros(self.Np_o);
|
|
object_upd_sum{ll,layer} = Gzeros(self.Np_o, true);
|
|
end
|
|
end
|
|
probe_update_sum = Gzeros(size(self.probe{1}));
|
|
probe_amp_corr = [0,0];
|
|
%%%%%%%%%%%%%%%%%% LSQ-ML algorithm %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
|
beta_probe = ones(self.Npos,par.Nlayers);
|
|
beta_object = ones(self.Npos,par.Nlayers);
|
|
|
|
%for ML is more useful to get close/overlapping positions
|
|
% use already precalculated indices
|
|
if is_method(par, 'MLs')
|
|
% load sparse indices
|
|
rand_ind = randi(length( cache.preloaded_indices_sparse));
|
|
indices = cache.preloaded_indices_sparse{rand_ind}.indices;
|
|
scan_ids = cache.preloaded_indices_sparse{rand_ind}.scan_ids;
|
|
else % load compact indices
|
|
indices = cache.preloaded_indices_compact{1}.indices;
|
|
scan_ids = cache.preloaded_indices_compact{1}.scan_ids;
|
|
end
|
|
Nind = length(indices);
|
|
|
|
for jj = 1:Nind
|
|
layer_ids{jj} = 1:par.Nlayers;
|
|
end
|
|
|
|
% shuffle order but keep same over iterations
|
|
if is_method(par, 'MLs')
|
|
% shuffle order
|
|
ind_range = randperm(Nind);
|
|
else
|
|
% MLc, call groups in given order, they are sorted by size to make
|
|
% execution on GPU more effecient + stable convergence
|
|
ind_range = 1:Nind;
|
|
end
|
|
|
|
%% apply updated in parallel over sets indices{jj}
|
|
for jj = ind_range
|
|
% list of positions solved in current subiteration
|
|
g_ind = indices{jj};
|
|
|
|
for ll = 1:max(par.probe_modes, par.object_modes)
|
|
% generate indices of the used probes
|
|
% single probe only
|
|
if par.share_probe % share incoherent modes
|
|
p_ind{ll} = 1;
|
|
else
|
|
if all(scan_ids{jj} == scan_ids{jj}(1))
|
|
p_ind{ll} = scan_ids{jj}(1);
|
|
else
|
|
% scan positions from multiple scans are processed in a single bunch
|
|
p_ind{ll} = scan_ids{jj};
|
|
end
|
|
end
|
|
end
|
|
|
|
% estimate forward model, ie wavefront behind the sample
|
|
[self, probe, obj_proj, psi] = get_forward_model(self, obj_proj, par,cache, g_ind, p_ind, scan_ids{jj}, layer_ids{jj});
|
|
|
|
%% load data to GPU
|
|
modF = get_modulus(self, cache, g_ind,true,jj);
|
|
mask = get_mask(self, cache.mask_indices, g_ind, par.damped_mask);
|
|
noise = get_noise(self, par, g_ind);
|
|
|
|
% get intensity (modulus) on detector including different corrections
|
|
[aPsi, aPsi2, cache, self] = get_reciprocal_model(self, psi, modF, mask,iter, g_ind, par, cache);
|
|
|
|
|
|
%%%%%%%%%%%%%%%%%%%%%%% LINEAR MODEL CORRECTIONS END %%%%%%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%
|
|
if iter > 0 && (verbose >= -1 || par.number_iterations > par.plot_results_every) && ...
|
|
((mod(iter,min(20, 2^(floor(2+iter/50)))) == 0 || iter < 10) ...
|
|
|| verbose()> 2) || any(iter == max(1,[1,par.number_iterations])) % || par.accelerated_gradients_start-1 >= iter
|
|
% calculate only sometimes to make it faster
|
|
[fourier_error(iter,g_ind)] = get_fourier_error(modF, aPsi, noise,mask, par.likelihood);
|
|
|
|
if any(~isfinite(fourier_error(iter,g_ind)))
|
|
if par.accelerated_gradients_start < par.number_iterations || par.momentum > 0
|
|
error('Convergence failed, error contains NaNs, quitting ... \n%s', 'If repeated, try to set eng.accelerated_gradients_start = inf or eng.momentum = 0; ')
|
|
else
|
|
error('Convergence failed, error contains NaNs, quitting ... \n%s', 'If repeated, try to set eng.beta_LSQ = 0.5 or less ')
|
|
end
|
|
end
|
|
end
|
|
%if iter>0
|
|
% disp(any(~isfinite(fourier_error(iter,:))))
|
|
%end
|
|
if strcmp(par.likelihood, 'poisson') || iter == 0
|
|
[chi,R] = modulus_constraint(modF,aPsi,psi, mask, noise, par, 1);
|
|
if iter == 0
|
|
% in the first iteration only find optimal scale for the probe
|
|
probe_amp_corr(1) = probe_amp_corr(1) + Ggather(sum(modF(:).^2));
|
|
probe_amp_corr(2) = probe_amp_corr(2) + Ggather(sum(aPsi(:).^2));
|
|
continue
|
|
end
|
|
else
|
|
chi = modulus_constraint(modF,aPsi,psi, mask, noise, par, 1);
|
|
end
|
|
|
|
if ~strcmp(par.likelihood, 'poisson') % soft memory cleanup
|
|
mask = []; aPsi = []; noise = []; modF = []; R=[];
|
|
end
|
|
|
|
if iter > par.estimate_NF_distance
|
|
% update estimation of the nearfield propagation distance c
|
|
[self, cache] = gradient_NF_propagation_solver(self,psi(:,end),chi, cache, g_ind);
|
|
end
|
|
if ~strcmp(par.likelihood, 'poisson') % soft memory cleanup
|
|
psi = [];
|
|
end
|
|
|
|
if strcmp(par.likelihood, 'poisson') % calculate only for the first mode
|
|
%ll = 1;
|
|
%% automatically find optimal step-size, note that for Gauss it is 1 !!
|
|
beta_xi = gradient_descent_xi_solver(self,modF, aPsi2, R,mask, g_ind, mean(cache.beta_xi(g_ind)), cache);
|
|
R = [];
|
|
for ll = 1:max(par.probe_modes, par.object_modes)
|
|
chi{ll} = chi{ll} .* beta_xi ;
|
|
end
|
|
end
|
|
|
|
%%%%%%%%%%%% LSQ optimization code (probe & object updates) %%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
|
for layer = par.Nlayers:-1:1
|
|
for ll = 1:max(par.probe_modes, par.object_modes)
|
|
object_reconstruct = iter >= par.object_change_start && (par.apply_multimodal_update || is_used(par, 'fly_scan') || ll <= par.object_modes );
|
|
probe_reconstruct = iter >= par.probe_change_start;
|
|
|
|
llo = min(par.object_modes, ll);
|
|
llp = min(par.probe_modes, ll);
|
|
|
|
% propagate to the previous layer
|
|
chi{ll} = back_fourier_proj(chi{ll}, self.modes{layer},g_ind);
|
|
|
|
if layer ~= par.Nlayers
|
|
% if only single layer is used, reuse obj_proj already
|
|
% loaded, but avoid storing obj_proj for each layer, rather load it again
|
|
obj_proj{llo} = get_views(self.object, obj_proj{llo},layer_ids{jj}(layer),llo, g_ind, cache, scan_ids{jj},[]);
|
|
end
|
|
|
|
% get update directions for each scan positions
|
|
if ( probe_reconstruct || layer > 1) && object_reconstruct
|
|
[probe_update, object_update_proj] = Gfun(@get_update_both, chi{ll}, obj_proj{llo}, probe{llp,layer});
|
|
elseif ( probe_reconstruct || layer > 1)
|
|
probe_update = Gfun(@get_update, chi{ll}, obj_proj{llo});
|
|
object_update_proj = 0;
|
|
else
|
|
probe_update = 0; m_probe_update = 0;
|
|
object_update_proj = Gfun(@get_update, chi{ll}, probe{llp,layer});
|
|
end
|
|
|
|
% refine single optimal probe update direction (use overlap constraint)
|
|
if probe_reconstruct || layer > 1
|
|
[self,m_probe_update, probe_update, cache] = refine_probe_update(self, obj_proj{llo}, probe_update, chi{ll},layer,ll,p_ind{ll},g_ind, par, cache);
|
|
end
|
|
if layer == 1
|
|
probe_update = [] ; % soft memory clean
|
|
end
|
|
|
|
% refine single optimal object update direction (use overlap constraint)
|
|
if object_reconstruct
|
|
[object_upd_sum,object_update_proj, cache] = refine_object_update(self, ...
|
|
object_update_proj,object_upd_sum,layer_ids{jj}(layer),scan_ids{jj},g_ind, par, cache);
|
|
end
|
|
|
|
|
|
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
|
%%%%% calculate the optimal step %%%%%%%%%%%%%%%%%%%%%%
|
|
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
|
if ll == 1 % && layer == par.Nlayers
|
|
if par.beta_LSQ > 0 && object_reconstruct && probe_reconstruct && par.Nlayers == 1 % for Nlayers > 1 seems to be better to better using gradient_projection_solver in order to keep behaviour the same for each layer !!
|
|
% calculate the optimal step using LSQ method
|
|
[beta_probe(g_ind,layer), beta_object(g_ind,layer)] = ...
|
|
get_optimal_LSQ_step(self,chi{ll},object_update_proj,m_probe_update,obj_proj{llo},probe{llp,layer},p_ind{ll} , par);
|
|
|
|
elseif par.beta_LSQ > 0 && (object_reconstruct || probe_reconstruct )
|
|
% computationally cheaper method that assumes only
|
|
% diagonal terms of the AA matrix
|
|
[beta_probe(g_ind,layer),beta_object(g_ind,layer)] = ...
|
|
gradient_projection_solver(self,chi{ll},obj_proj{llo},probe{llp,layer},...
|
|
object_update_proj, m_probe_update,p_ind{ll}, par, cache);
|
|
|
|
elseif par.beta_LSQ == 0
|
|
% use constant step
|
|
beta_probe(g_ind,layer) = par.beta_probe;
|
|
beta_object(g_ind,layer) = par.beta_object;
|
|
end
|
|
|
|
beta_object(g_ind,layer) = beta_object(g_ind,layer) / par.Nlayers;
|
|
|
|
if (is_used(par, 'fly_scan') || par.apply_multimodal_update)
|
|
beta_object(g_ind,layer) = beta_object(g_ind,layer) / par.probe_modes;
|
|
end
|
|
end
|
|
object_update_proj = []; % soft memory clean
|
|
beta_object = max(0, beta_object);
|
|
beta_probe = max(0, beta_probe);
|
|
|
|
if any(jj == ind_range(end)) && ll == 1
|
|
% show the optimal steps calculation
|
|
% if verbose() > 2
|
|
% if mod(iter, 5) == 0
|
|
% plotting.smart_figure(121122)
|
|
% plot(squeeze(real(beta_object)))
|
|
% hold all
|
|
% plot(squeeze(real(beta_probe)))
|
|
% hold off
|
|
% legend({'O', 'P'})
|
|
% title('Update step for each scan position')
|
|
% xlabel('Scan positions #')
|
|
% ylabel('Step')
|
|
% drawnow
|
|
% end
|
|
% end
|
|
verbose(1,'Average step p%i: %3.3g o%i: %3.3g layer %i', llp,mean(beta_probe(g_ind,layer)),llo,mean(beta_object(g_ind,layer)), layer);
|
|
end
|
|
|
|
%%%%%%%%%%%%%%% apply update with the optimal LSQ step %%%%%%%%%%%%%%%%%
|
|
if probe_reconstruct && layer == 1 && ll <= max(par.probe_modes) % multilayer extension -> update probe only from the first layer
|
|
self.probe{ll} = update_probe(self.probe{ll}, m_probe_update, par, p_ind{ll}, g_ind, beta_probe, Nind); % finally update also the probe
|
|
end
|
|
|
|
if object_reconstruct && is_method(par, 'MLs')
|
|
self.object = update_object(self, self.object, object_upd_sum, layer_ids{jj}(layer), llo, {g_ind}, scan_ids(jj), par, cache, beta_object);
|
|
end
|
|
|
|
if ll == 1 && layer == ceil(par.Nlayers/2) % assume that sample in center is best constrained
|
|
%%%%%%%%%%%%% update other parameters of the ptychography model
|
|
if iter >= par.probe_position_search || iter >= par.detector_rotation_search || iter >= par.detector_scale_search
|
|
% find optimal position shift that minimize chi{1} in current iteration
|
|
[pos_update, pos_rotate_upd,probe_scale_upd, cache] = gradient_position_solver(self, chi{1}, obj_proj{1},probe{1,layer}, g_ind, iter, cache, par);
|
|
self.modes{1}.probe_scale_upd(end+1)=self.modes{1}.probe_scale_upd(end)+mean(probe_scale_upd);
|
|
self.modes{1}.probe_positions(g_ind,:)=self.modes{1}.probe_positions(g_ind,:)+pos_update;
|
|
self.modes{1}.probe_rotation_all(g_ind)=self.modes{1}.probe_rotation_all(g_ind)+squeeze(pos_rotate_upd);
|
|
end
|
|
|
|
if iter > par.probe_fourier_shift_search
|
|
% search position corrections in the Fourier space, use
|
|
% only informatiom from the first mode, has to be after
|
|
% the probes updated , SEARCH ONLY FOR THE FIRST MODE
|
|
self.modes{1} = gradient_fourier_position_solver(chi{1}, obj_proj{1},probe{1,layer},self.modes{1}, g_ind);
|
|
end
|
|
end
|
|
|
|
if par.Nlayers > 1
|
|
% get update direction for the next layer
|
|
chi{ll} = probe_update; % .* median(beta_probe(g_ind,layer));
|
|
else
|
|
chi{ll} = []; % soft memory clean
|
|
end
|
|
end
|
|
end
|
|
|
|
if iter > par.estimate_NF_distance
|
|
%% correct propagation distance if updated
|
|
for i = 1:par.Nmodes
|
|
ASM = exp( self.modes{i}.distances(end)* cache.ASM_difference);
|
|
self.modes{i}.ASM_factor = ASM;
|
|
self.modes{i}.cASM_factor = conj(ASM);
|
|
end
|
|
end
|
|
% to be used for momentum calculation, use only the last layer
|
|
if par.momentum
|
|
uniq_p_ind = unique(p_ind{ll});
|
|
probe_update_sum(:,:,uniq_p_ind,1) = probe_update_sum(:,:,uniq_p_ind,1) + m_probe_update / Nind;
|
|
end
|
|
|
|
end
|
|
%if iter>0
|
|
% disp(size(fourier_error(iter,:)))
|
|
% disp(any(~isfinite(fourier_error(iter,:))))
|
|
%end
|
|
if iter == 0
|
|
% apply initial correction for the probe intensity and return
|
|
probe_amp_corr = sqrt(probe_amp_corr(1) / probe_amp_corr(2)); %% calculate ratio between modF^2 and aPsi^2
|
|
for ii = 1:par.probe_modes
|
|
self.probe{ii} = self.probe{ii}*probe_amp_corr;
|
|
end
|
|
verbose(2,'Probe amplitude corrected by %.3g',probe_amp_corr)
|
|
return
|
|
end
|
|
|
|
% applying single update emulates behaviour of the original ML method ->
|
|
% provides better noise robustness
|
|
% advantage is that less memory is needed and no linesearch is required
|
|
|
|
object_reconstruct = iter >= par.object_change_start; % && (par.apply_multimodal_update || is_used(par, 'fly_scan') || ll <= par.object_modes );
|
|
probe_reconstruct = iter >= par.probe_change_start;
|
|
% if true, caclulate momentum and use is for acceleration
|
|
momentum_acceleration = isfield(par, 'momentum') && par.momentum && par.number_iterations < par.accelerated_gradients_start;
|
|
if object_reconstruct && is_method(par, 'MLc')
|
|
for ll = 1:par.object_modes
|
|
for layer = 1:par.Nlayers
|
|
[self.object, object_upd_sum] = update_object(self, self.object, object_upd_sum, layer, ll, indices, scan_ids, par, cache, beta_object(:,layer));
|
|
end
|
|
end
|
|
%% apply momentum update on the object
|
|
if momentum_acceleration
|
|
[self, cache] = add_momentum_object(self, cache, par, object_upd_sum, iter, fourier_error, beta_object);
|
|
end
|
|
end
|
|
|
|
if probe_reconstruct && momentum_acceleration
|
|
%% apply momentum update on the probe
|
|
[self, cache] = add_momentum_probe(self, cache, par, {probe_update_sum}, iter, fourier_error, beta_probe);
|
|
end
|
|
|
|
%% FLY-SCAN: join all subprobes
|
|
if iter >= par.probe_change_start
|
|
if is_used(par,'fly_scan')
|
|
probe_new = 0;
|
|
for ll = 1:par.Nmodes
|
|
probe_new = probe_new + self.probe{ll}/par.Nmodes;
|
|
end
|
|
for ll = 1:par.Nmodes
|
|
switch par.flyscan_intensity
|
|
case 'varying' % allow variation of the modes intensity
|
|
aa = sum2(self.probe{ll} .* conj(probe_new));
|
|
bb = sum2(abs(probe_new).^2);
|
|
%proj(ll,1,:) = real(aa./ bb); %seems no need to have proj
|
|
%self.probe{ll} = proj(ll,1,:) .* probe_new;
|
|
self.probe{ll} = real(aa./bb) .* probe_new;
|
|
case 'constant' % assume constant intensity
|
|
self.probe{ll} = probe_new;
|
|
end
|
|
end
|
|
end
|
|
end
|
|
% if iter >= par.probe_change_start
|
|
% for ll = 1:par.probe_modes
|
|
% self.probe{ll} = apply_probe_contraints(self.probe{ll}, self.modes{min(ll,end)});
|
|
% end
|
|
% end
|
|
|
|
|
|
end
|
|
|
|
|
|
%% merged CUDA kernels for faster calculations
|
|
|
|
function update = get_update(chi, proj)
|
|
update = chi .* conj(proj);
|
|
end
|
|
|
|
function [update_1, update_2] = get_update_both(chi, proj_1, proj_2)
|
|
update_1 = chi .* conj(proj_1);
|
|
update_2 = chi .* conj(proj_2);
|
|
end
|
|
|
|
|
|
|
|
|