mirror of
https://github.com/c-sooyoung/fold_slice.git
synced 2026-09-17 22:49:08 +09:00
424 lines
20 KiB
Matlab
424 lines
20 KiB
Matlab
% PIE - generalized version of the ptychographic iterative engine
|
|
%
|
|
%[self, cache, fourier_error] = PIE(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
|
|
%
|
|
%
|
|
% Publications most relevant to the Difference-Map implementation
|
|
% Odstrcil, M., Baksh, P., Boden, S. A., Card, R., Chad, J. E., Frey, J. G., & Brocklesby, W. S
|
|
% "Ptychographic coherent diffractive imaging with orthogonal probe relaxation."
|
|
% Optics express 24.8 (2016): 8360-8369.
|
|
|
|
|
|
function [self, cache, fourier_error] = PIE(self,par,cache,fourier_error,iter)
|
|
import engines.GPU.GPU_wrapper.*
|
|
import math.*
|
|
import utils.*
|
|
import engines.GPU.shared.*
|
|
import engines.GPU.PIE.*
|
|
|
|
%% MULTILAYER EXTENSION fix or remove in future
|
|
par.probe_modes = max(par.Nlayers, par.probe_modes);
|
|
par.object_modes = max(par.Nlayers, par.object_modes);
|
|
|
|
%disp(size(self.probe{1}))
|
|
if par.variable_probe
|
|
% expand the probe to the full size
|
|
probe_0 = self.probe{1};
|
|
self.probe{1} = reshape(self.probe{1},prod(self.Np_p),[]);
|
|
self.probe{1} = reshape(self.probe{1} * self.probe_evolution', self.Np_p(1), self.Np_p(2), []);
|
|
end
|
|
%disp(size(self.probe_evolution))
|
|
|
|
par.multilayer_object = par.Nlayers > 1;
|
|
par.multilayer_probe = false; % not supported anymore
|
|
probe_amp_corr = [0,0];
|
|
|
|
psi = cell(par.Nmodes, 1);
|
|
Psi = cell(par.Nmodes, 1);
|
|
probe_max = cell(par.probe_modes,1);
|
|
object_max = cell(par.object_modes,1);
|
|
aprobe2 = abs(mean(self.probe{1},3)).^2; % update only once per iteration
|
|
|
|
if (is_method(par, 'ePIE') && ...
|
|
... % empirical estimation when the hybrid PIE method should be used
|
|
par.grouping > self.Npos/sqrt( pi^2 * mean(Ggather(cache.MAX_ILLUM)) / max(aprobe2(:)))) % || ...% use hybrid ePIE in case of large grouping
|
|
% (~isempty(self.modes{end}.ASM_factor) && par.grouping > 1) || ...
|
|
% (is_method(par, 'ePIE')&& par.multilayer_object)
|
|
par.method = 'hPIE'; % hybrid ePIE
|
|
if iter == 1;verbose(1,'Switching to hybrid PIE method '); end
|
|
end
|
|
if is_method(par, {'hPIE'})
|
|
for ll = 1:(par.object_modes*par.Nscans)
|
|
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
|
|
for ll = 1:par.probe_modes
|
|
probe_upd_sum{ll}= (1+1i)*1e-8;
|
|
probe_illum_sum{ll} = 1e-8;
|
|
end
|
|
end
|
|
|
|
|
|
for ll = 1:par.Nlayers
|
|
obj_proj{ll} = Gzeros([self.Np_p, par.grouping], true);
|
|
end
|
|
if par.delta_p
|
|
grad = @get_grad_lsq; % dumped least squares gradient (preconditioner)
|
|
else
|
|
grad = @get_grad_flat; % common PIE gradient
|
|
end
|
|
|
|
%%%%%%%%%%%%%%%%%% ePIE algorithm %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
|
if par.grouping > 1
|
|
for ll = 1:par.probe_modes
|
|
if ~par.multilayer_object || ll == 1 % update only first layer probe
|
|
probe_max{ll} = Ggather(max2(abs(self.probe{ll}))).^2;
|
|
end
|
|
end
|
|
for ll = 1:par.object_modes
|
|
if ~par.multilayer_probe || ll == 1 % update only first layer object
|
|
object_max{ll} = Ggather(max2(abs(self.object{ll}(cache.object_ROI{:})))).^2;
|
|
end
|
|
end
|
|
end
|
|
|
|
if any(isnan(object_max{1})) || any(isnan(probe_max{1}))
|
|
error('Object or probe is nan, try smaller grouping')
|
|
end
|
|
|
|
|
|
|
|
% use already precalculated 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;
|
|
|
|
for ind_ii = randperm(length(indices))
|
|
g_ind = indices{ind_ii};
|
|
for ll = 1:max([par.probe_modes, par.object_modes,par.Nlayers])
|
|
% generate indices of the used probes
|
|
if par.variable_probe && ll == 1
|
|
p_ind{ll} = g_ind;
|
|
% elseif par.multilayer_object && ll > 1
|
|
% p_ind{ll} = 1:length(ii);
|
|
else % single probe only
|
|
if par.share_probe %|| ll > 1 % share incoherent modes
|
|
p_ind{ll} = 1;
|
|
else
|
|
if all(scan_ids{ind_ii} == scan_ids{ind_ii}(1))
|
|
p_ind{ll} = scan_ids{ind_ii}(1);
|
|
else
|
|
p_ind{ll} = scan_ids{ind_ii};
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
%% load data to GPU (if not loaded yet)
|
|
modF = get_modulus(self, cache, g_ind);
|
|
mask = get_mask(self, cache.mask_indices, g_ind);
|
|
noise = get_noise(self, par, g_ind);
|
|
|
|
|
|
% get objects projections
|
|
|
|
|
|
|
|
for layer = 1:par.Nlayers
|
|
ll = 1;
|
|
obj_proj{layer} = get_views(self.object, obj_proj{layer},layer,ll,g_ind, cache, scan_ids{ind_ii},[]);
|
|
end
|
|
% get illumination probe
|
|
for ll = 1:par.probe_modes
|
|
if ~par.multilayer_object || ismember(ll, [1,par.Nlayers:par.probe_modes])
|
|
probe{ll} = self.probe{min(ll,end)}(:,:,min(end,p_ind{ll}));
|
|
end
|
|
end
|
|
|
|
for ll = 1:max([par.probe_modes, par.object_modes, par.Nlayers]) % Nlayers
|
|
%% fourier propagation
|
|
|
|
if (ll == 1 && (par.multilayer_object || par.multilayer_probe) )
|
|
probe{ll} = self.probe{min(ll,end)}(:,:,min(end,p_ind{ll}));
|
|
end
|
|
if ll > 1 && par.multilayer_object || par.grouping == 1 % raw ePIE
|
|
probe_max{ll} = max(Ggather(max2(abs(probe{ll})))).^2;
|
|
end
|
|
%disabled by YJ. seems like a bug
|
|
%{
|
|
if ll > 1 && par.multilayer_probe || par.grouping == 1 % raw ePIE
|
|
object_max{ll} = max(Ggather(max2(abs(obj_proj{ll})))).^2;
|
|
end
|
|
%}
|
|
if (ll == 1 && par.apply_subpix_shift)
|
|
probe{ll} = apply_subpx_shift(probe{ll}, self.modes{min(end,ll)}.sub_px_shift(g_ind,:) );
|
|
end
|
|
|
|
probe{ll} = apply_subpx_shift_fft(probe{ll}, self.modes{1}.probe_fourier_shift(g_ind,:));
|
|
|
|
|
|
% get projection of the object and probe
|
|
psi{ll} = bsxfun(@times, obj_proj{min(ll,end)}, probe{min(ll,end)});
|
|
Psi{ll} = fwd_fourier_proj(psi{ll} , self.modes{min(end, ll)});
|
|
|
|
if ll < par.Nlayers && par.multilayer_object
|
|
probe{ll+1} = Psi{ll};
|
|
end
|
|
|
|
if ll < par.Nlayers && par.multilayer_probe
|
|
obj_proj{ll+1} = Psi{ll};
|
|
end
|
|
end
|
|
|
|
% get intensity (modulus) on detector including different corrections
|
|
aPsi = get_reciprocal_model(self, Psi, modF, mask,iter, g_ind, par, cache);
|
|
|
|
|
|
%%%%%%%%%%%%%%%%%%%%%%% LINEAR MODEL CORRECTIONS END %%%%%%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%
|
|
if iter > 0 && par.get_error && (mod(iter,min(20, 2^(floor(2+iter/50)))) == 0 || (iter < 20)) || iter == par.number_iterations % calculate only sometimes to make it faster
|
|
[fourier_error(iter,g_ind)] = get_fourier_error(modF, aPsi, noise,mask, par.likelihood);
|
|
end
|
|
|
|
|
|
|
|
if (par.multilayer_object || par.multilayer_probe)
|
|
[Psi(end),R] = modulus_constraint(modF,aPsi,Psi(end), mask, noise, par, 1); % apply only on the last layer !!!
|
|
else
|
|
[Psi,R] = modulus_constraint(modF,aPsi,Psi, mask, noise, par,1);
|
|
end
|
|
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
|
|
if strcmp(par.likelihood, 'poisson') % calculate only for the first mode
|
|
%% automatically find optimal step-size, note that for Gauss it is 1 !!
|
|
cache.beta_xi(g_ind) = gradient_descent_chi_solver(self,modF, aPsi2, R,mask, g_ind, mean(cache.beta_xi), cache);
|
|
end
|
|
|
|
if(par.multilayer_object || par.multilayer_probe)
|
|
ind_modes = par.Nlayers:-1:1;
|
|
else
|
|
ind_modes = 1:max(par.probe_modes, par.object_modes);
|
|
end
|
|
|
|
|
|
for ll = ind_modes
|
|
layer = ll;
|
|
chi = back_fourier_proj(Psi{min(end,ll)}, self.modes{min(end,ll)})-psi{min(end,ll)};
|
|
|
|
|
|
%% get optimal gradient lenghts
|
|
object_update=0; probe_update=0;m_probe_update= 0;
|
|
|
|
if iter >= par.object_change_start && (ll <= max(par.Nlayers, par.object_modes) || par.apply_multimodal_update)
|
|
object_update = Gfun(grad,chi, probe{min(ll,end)},...
|
|
probe_max{min(end,ll)}(1,1,min(end,p_ind{ll})),par.delta_p);
|
|
end
|
|
|
|
|
|
if iter >= par.probe_change_start && ll <= max(par.probe_modes, par.Nlayers)
|
|
%% find optimal probe update !!!
|
|
probe_update = Gfun(grad,chi,obj_proj{min(end,ll)},...
|
|
object_max{min(ll,end)}, par.delta_p);
|
|
m_probe_update = mean(probe_update,3);
|
|
end
|
|
|
|
if ((ll == 1 && ~(par.multilayer_object || par.multilayer_probe)) || ...
|
|
(ll == par.Nlayers && (par.multilayer_object || par.multilayer_probe))) && ...
|
|
(par.beta_LSQ || iter >= par.probe_position_search)
|
|
%% variable step extension, apply only the first mode except the 3PIE case
|
|
%% it will use the same alpha for the higher modes !!!!
|
|
|
|
[cache.beta_probe(g_ind),cache.beta_object(g_ind)] = gradient_projection_solver(self,chi,obj_proj{min(end,ll)},probe{ll},...
|
|
object_update, m_probe_update,p_ind{ll}, par, cache);
|
|
if any(g_ind ==1)
|
|
verbose(1,'Average alpha p:%3.3g o:%3.3g ', mean(cache.beta_probe),mean(cache.beta_object));
|
|
end
|
|
end
|
|
|
|
|
|
|
|
|
|
%%%%%%%%%%%%%%%%%%%%% PROBE UPDATE %%%%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%
|
|
%% update probe first
|
|
if ((iter >= par.probe_change_start) && ll <= max(par.probe_modes,par.Nlayers) ...
|
|
&& is_method(par, 'PIE')) || ...
|
|
(par.multilayer_object && ll > 1 && ll <= par.probe_modes)
|
|
% only in case of first layer probe, otherwise update interprobes
|
|
|
|
beta_probe = get_vals(cache.beta_probe,g_ind) .* get_vals(cache.beta_xi,g_ind);
|
|
|
|
if is_method(par, {'ePIE', 'hPIE'})
|
|
%%%%%%%%%%% update probe %%%%%%%%%%%%%%%%%%%%%%%%%%
|
|
probe{ll} = Gfun(@upd_probe_Gfun,probe{ll},probe_update, beta_probe);
|
|
|
|
if (ll == 1 && par.apply_subpix_shift)
|
|
probe{ll} = apply_subpx_shift(probe{ll} , -self.modes{min(end,ll)}.sub_px_shift(g_ind,:));
|
|
end
|
|
|
|
|
|
if iter >= par.probe_change_start
|
|
if (par.variable_probe && ll == 1)
|
|
self.probe{ll}(:,:,p_ind{ll}) = probe{ll}; % slowest line for large datasets !!!!!
|
|
elseif (~par.multilayer_object || ll == 1) && ll <= par.probe_modes
|
|
self.probe{ll} = mean(probe{ll},3); % merge information from all the shifted probes if needed
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
if iter > par.probe_fourier_shift_search && ll == 1
|
|
% search position corrections in the Fourier space, use
|
|
% only informatiom from the first mode, has to be after
|
|
% the probes updated
|
|
self.modes{1} = gradient_fourier_position_solver(chi, obj_proj{1},probe{1},self.modes{1}, g_ind);
|
|
end
|
|
|
|
if iter >= par.probe_position_search
|
|
% find optimal position shift that minimize chi{1} in current iteration
|
|
%[pos_update, cache] = gradient_position_solver(self, chi, obj_proj{1},probe{1,layer}, g_ind, iter, cache);
|
|
%modified by YJ to prevent bug
|
|
[pos_update, ~,~,cache] = gradient_position_solver(self, chi, obj_proj{1},probe{1,layer}, g_ind, iter, cache, par);
|
|
|
|
self.modes{1}.probe_positions(g_ind,:)=self.modes{1}.probe_positions(g_ind,:)+pos_update;
|
|
end
|
|
%%%%%%%%%%%%%%%%%%%%% OBJECT UPDATE %%%%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%
|
|
if iter >= min([par.object_change_start]) && ....
|
|
( ll <= max(par.Nlayers, par.object_modes) || par.apply_multimodal_update )
|
|
if ll ~= 1 && ~(par.multilayer_object || par.multilayer_probe) ; continue; end
|
|
|
|
if iter >= par.object_change_start % && ~(par.multilayer_probe && ll > 1) % the objects are just empty
|
|
beta_object = get_vals(cache.beta_object,g_ind) .* get_vals(cache.beta_xi,g_ind);
|
|
if par.share_object
|
|
obj_ids = 1; % update only the first object
|
|
else
|
|
obj_ids = unique(scan_ids{ind_ii}); % update only the objects processed in this block
|
|
end
|
|
|
|
if any(beta_object ~= 1)
|
|
object_update = bsxfun(@times, object_update, beta_object);
|
|
end
|
|
|
|
|
|
if is_method(par, 'ePIE') % use always in nearfield
|
|
% classical ePIE, faster constraint application, but it will fail with too high grouping
|
|
self.object = set_views(self.object, object_update,layer, obj_ids, g_ind, cache, scan_ids{ind_ii},[]);
|
|
elseif is_method(par, 'hPIE') %% hybrid PIE
|
|
if par.Nscans == 1 || par.share_object
|
|
ind_tmp = 1;
|
|
else
|
|
ind_tmp = 1+par.object_modes* ((1:par.Nscans)-1);
|
|
end
|
|
for kk = ind_tmp
|
|
obj_illum_sum{kk,layer}(:) = 0;
|
|
object_upd_sum{kk,layer}(:) = 0;
|
|
end
|
|
object_update = bsxfun(@times, object_update, aprobe2); % make is more like dumped LSQ solution
|
|
[object_upd_sum,obj_illum_sum] = set_views_rc(object_upd_sum,obj_illum_sum, object_update,aprobe2,layer,obj_ids, g_ind, cache, scan_ids{ind_ii},[]);
|
|
for kk = ind_tmp
|
|
self.object{kk,layer} = Gfun(@object_update_Gfun, self.object{kk,layer},object_upd_sum{kk,layer}, obj_illum_sum{kk,layer}, cache.MAX_ILLUM(min(kk,end)));
|
|
end
|
|
else
|
|
error('Unimplemented method %s ', par.method)
|
|
end
|
|
end
|
|
end
|
|
if ll > 1 && par.multilayer_object
|
|
% % apply rescaling to make scaling correction
|
|
Psi{ll-1} = probe{ll};
|
|
end
|
|
if ll > 1 && par.multilayer_probe
|
|
Psi{ll-1} = obj_proj{ll} + object_update;
|
|
end
|
|
end
|
|
|
|
if check_avail_memory < 0.2 || ~par.keep_on_gpu
|
|
% slow step that is not needed if there is enough memory ,
|
|
% it can slow down almost twice !!!
|
|
clear Psi R aPsi2 psi probe_update object_update chi
|
|
if par.keep_on_gpu
|
|
warning('Low GPU memory')
|
|
end
|
|
end
|
|
|
|
end
|
|
|
|
|
|
|
|
|
|
if par.multilayer_object
|
|
self.probe = self.probe(1);
|
|
end
|
|
|
|
if par.variable_probe && iter >= par.probe_change_start
|
|
%disp(size(self.probe{1}))
|
|
[self.probe{1}, self.probe_evolution] = apply_SVD_filter(self.probe{1}, par.variable_probe_modes+1, self.modes{1});
|
|
%disp(size(self.probe{1}))
|
|
|
|
elseif par.variable_probe && iter < par.probe_change_start
|
|
self.probe{1} = probe_0;
|
|
elseif ( ~isempty(self.probe_support)) && iter >= par.probe_change_start
|
|
self.probe{1} = apply_probe_contraints(self.probe{1}, self.modes{1});
|
|
end
|
|
|
|
if iter == 0
|
|
% apply initial correction for the probe intensity and return
|
|
% it seems to be safer to underestimate the probe amplitude for variable probe method
|
|
probe_amp_corr = 0.5*sqrt(probe_amp_corr(1) / probe_amp_corr(2)); %% calculate ratio between modF^2 and aPsi^2
|
|
|
|
%modified by YJ to fix bug with multilayer
|
|
if par.multilayer_object
|
|
self.probe{1} = self.probe{1}*probe_amp_corr;
|
|
else
|
|
|
|
for ii = 1:par.probe_modes
|
|
self.probe{ii} = self.probe{ii}*probe_amp_corr;
|
|
end
|
|
end
|
|
verbose(2,'Probe amplitude corrected by %.3g',probe_amp_corr)
|
|
return
|
|
end
|
|
|
|
end
|
|
|
|
function probe = upd_probe_Gfun(probe,probe_update, alpha_p)
|
|
probe = probe + alpha_p.*probe_update;
|
|
end
|
|
|
|
function grad = get_grad_flat(chi,proj,max2, delta )
|
|
% ePIE method
|
|
grad = (1/max2) * chi .* conj(proj) ;
|
|
end
|
|
|
|
function grad = get_grad_lsq(chi,proj,max2, delta )
|
|
% dumped LSQ method
|
|
aproj = abs(proj) ;
|
|
grad = chi .* conj(proj) .* aproj ./ (aproj.^2 + delta*max2)./sqrt(max2);
|
|
end
|
|
|
|
function object = object_update_Gfun(object,object_upd_sum, obj_illum_sum, max)
|
|
object = object + object_upd_sum ./ (obj_illum_sum+1e-9* max);
|
|
end
|
|
|
|
function array = get_vals(array, ind)
|
|
if isscalar(array)
|
|
return
|
|
else
|
|
array = reshape(array(ind),1,1,[]);
|
|
end
|
|
end
|