mirror of
https://github.com/c-sooyoung/fold_slice.git
synced 2026-09-17 22:49:08 +09:00
354 lines
18 KiB
Matlab
354 lines
18 KiB
Matlab
%DM Difference-Map algorithm
|
||
%
|
||
% [self, cache, psi_dash, fourier_error ] = DM(self,par,cache,psi_dash,fourier_error, iter)
|
||
%
|
||
%
|
||
% ** self structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
|
||
% ** par structure containing parameters for the engines
|
||
% ** psi_dash complex projection from previous iteration, or emppty in first iteration
|
||
% ** 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
|
||
% ++ psi_dash complex projection from previous iteration, or emppty in first iteration
|
||
% ++ fourier_error array [Npos,1] containing evolution of reconstruction error
|
||
%
|
||
% Publications most relevant to the Difference-Map implementation
|
||
% + P. Thibault, M. Dierolf, A. Menzel, O. Bunk, C. David, F. Pfeiffer,
|
||
% "High-Resolution Scanning X-ray Diffraction Microscopy," Science 321, 379-382 (2008)
|
||
% + P. Thibault, M. Dierolf, O. Bunk, A. Menzel, F. Pfeiffer,
|
||
% "Probe retrieval in ptychographic coherent diffractive imaging,"
|
||
% Ultramicroscopy 109, 338–343 (2009)
|
||
|
||
% Academic License Agreement
|
||
%
|
||
% Source Code
|
||
%
|
||
% Introduction
|
||
% • This license agreement sets forth the terms and conditions under which the PAUL SCHERRER INSTITUT (PSI), CH-5232 Villigen-PSI, Switzerland (hereafter "LICENSOR")
|
||
% will grant you (hereafter "LICENSEE") a royalty-free, non-exclusive license for academic, non-commercial purposes only (hereafter "LICENSE") to use the cSAXS
|
||
% ptychography MATLAB package computer software program and associated documentation furnished hereunder (hereafter "PROGRAM").
|
||
%
|
||
% Terms and Conditions of the LICENSE
|
||
% 1. LICENSOR grants to LICENSEE a royalty-free, non-exclusive license to use the PROGRAM for academic, non-commercial purposes, upon the terms and conditions
|
||
% hereinafter set out and until termination of this license as set forth below.
|
||
% 2. LICENSEE acknowledges that the PROGRAM is a research tool still in the development stage. The PROGRAM is provided without any related services, improvements
|
||
% or warranties from LICENSOR and that the LICENSE is entered into in order to enable others to utilize the PROGRAM in their academic activities. It is the
|
||
% LICENSEE’s responsibility to ensure its proper use and the correctness of the results.”
|
||
% 3. THE PROGRAM IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR
|
||
% A PARTICULAR PURPOSE AND NONINFRINGEMENT OF ANY PATENTS, COPYRIGHTS, TRADEMARKS OR OTHER RIGHTS. IN NO EVENT SHALL THE LICENSOR, THE AUTHORS OR THE COPYRIGHT
|
||
% HOLDERS BE LIABLE FOR ANY CLAIM, DIRECT, INDIRECT OR CONSEQUENTIAL DAMAGES OR OTHER LIABILITY ARISING FROM, OUT OF OR IN CONNECTION WITH THE PROGRAM OR THE USE
|
||
% OF THE PROGRAM OR OTHER DEALINGS IN THE PROGRAM.
|
||
% 4. LICENSEE agrees that it will use the PROGRAM and any modifications, improvements, or derivatives of PROGRAM that LICENSEE may create (collectively,
|
||
% "IMPROVEMENTS") solely for academic, non-commercial purposes and that any copy of PROGRAM or derivatives thereof shall be distributed only under the same
|
||
% license as PROGRAM. The terms "academic, non-commercial", as used in this Agreement, mean academic or other scholarly research which (a) is not undertaken for
|
||
% profit, or (b) is not intended to produce works, services, or data for commercial use, or (c) is neither conducted, nor funded, by a person or an entity engaged
|
||
% in the commercial use, application or exploitation of works similar to the PROGRAM.
|
||
% 5. LICENSEE agrees that it shall make the following acknowledgement in any publication resulting from the use of the PROGRAM or any translation of the code into
|
||
% another computing language:
|
||
% "Data processing was carried out using the cSAXS ptychography MATLAB package developed by the Science IT and the coherent X-ray scattering (CXS) groups, Paul
|
||
% Scherrer Institut, Switzerland."
|
||
%
|
||
% Additionally, any publication using the package, or any translation of the code into another computing language should cite for difference map:
|
||
% P. Thibault, M. Dierolf, A. Menzel, O. Bunk, C. David, F. Pfeiffer, High-resolution scanning X-ray diffraction microscopy, Science 321, 379–382 (2008).
|
||
% (doi: 10.1126/science.1158573),
|
||
% for maximum likelihood:
|
||
% P. Thibault and M. Guizar-Sicairos, Maximum-likelihood refinement for coherent diffractive imaging, New J. Phys. 14, 063004 (2012).
|
||
% (doi: 10.1088/1367-2630/14/6/063004),
|
||
% for mixed coherent modes:
|
||
% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
|
||
% and/or for multislice:
|
||
% E. H. R. Tsai, I. Usov, A. Diaz, A. Menzel, and M. Guizar-Sicairos, X-ray ptychography with extended depth of field, Opt. Express 24, 29089–29108 (2016).
|
||
% (doi: 10.1364/OE.24.029089).
|
||
% 6. Except for the above-mentioned acknowledgment, LICENSEE shall not use the PROGRAM title or the names or logos of LICENSOR, nor any adaptation thereof, nor the
|
||
% names of any of its employees or laboratories, in any advertising, promotional or sales material without prior written consent obtained from LICENSOR in each case.
|
||
% 7. Ownership of all rights, including copyright in the PROGRAM and in any material associated therewith, shall at all times remain with LICENSOR, and LICENSEE
|
||
% agrees to preserve same. LICENSEE agrees not to use any portion of the PROGRAM or of any IMPROVEMENTS in any machine-readable form outside the PROGRAM, nor to
|
||
% make any copies except for its internal use, without prior written consent of LICENSOR. LICENSEE agrees to place the following copyright notice on any such copies:
|
||
% © All rights reserved. PAUL SCHERRER INSTITUT, Switzerland, Laboratory for Macromolecules and Bioimaging, 2017.
|
||
% 8. The LICENSE shall not be construed to confer any rights upon LICENSEE by implication or otherwise except as specifically set forth herein.
|
||
% 9. DISCLAIMER: LICENSEE shall be aware that Phase Focus Limited of Sheffield, UK has an international portfolio of patents and pending applications which relate
|
||
% to ptychography and that the PROGRAM may be capable of being used in circumstances which may fall within the claims of one or more of the Phase Focus patents,
|
||
% in particular of patent with international application number PCT/GB2005/001464. The LICENSOR explicitly declares not to indemnify the users of the software
|
||
% in case Phase Focus or any other third party will open a legal action against the LICENSEE due to the use of the program.
|
||
% 10. This Agreement shall be governed by the material laws of Switzerland and any dispute arising out of this Agreement or use of the PROGRAM shall be brought before
|
||
% the courts of Zürich, Switzerland.
|
||
|
||
|
||
function [self, cache, psi_dash, fourier_error ] = DM(self,par,cache,psi_dash,fourier_error, iter)
|
||
import utils.*
|
||
import math.*
|
||
import engines.GPU.shared.*
|
||
import engines.GPU.GPU_wrapper.*
|
||
|
||
psi = cell(par.Nmodes, 1);
|
||
Psi = cell(par.Nmodes, 1);
|
||
|
||
%%%%%%%%%%%%%%%%%%%%%%%%% Difference maps algorithm %%%%%%%%%%%%%%%%%%%
|
||
beta = 1;
|
||
gamma = 1;
|
||
relax_mask = 1; % smoothly change relaxation of the mask
|
||
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
||
probe_norm = norm2(self.probe{1});
|
||
object_modes = length(self.object);
|
||
probe_amp_corr = [0,0];
|
||
|
||
for ll = 1:par.object_modes
|
||
obj_proj{ll} = Gzeros([self.Np_p, 0], true);
|
||
end
|
||
for ll = 1:object_modes
|
||
obj_illum{ll,1} = Gzeros(self.Np_o);
|
||
obj_update{ll,1} = Gzeros(self.Np_o, true);
|
||
end
|
||
for ll = 1:par.probe_modes
|
||
probe_illum{ll} = Gzeros(size(self.probe{ll}));
|
||
probe_update{ll} = Gzeros(size(self.probe{ll}), true);
|
||
end
|
||
|
||
% use precalculated blocks to be solved in parallel
|
||
indices = cache.preloaded_indices_simple{1}.indices;
|
||
scan_ids = cache.preloaded_indices_simple{1}.scan_ids;
|
||
Nind = length(indices);
|
||
|
||
p_ind = cell(max(par.probe_modes, par.object_modes),1);
|
||
for jj = 1:Nind
|
||
g_ind = indices{jj};
|
||
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
if par.share_probe || ll > 1 % 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
|
||
p_ind{ll} = scan_ids{jj};
|
||
end
|
||
end
|
||
end
|
||
|
||
% load on GPU if needed
|
||
if ~par.keep_on_gpu
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
try; psi_dash{ll,jj} = Garray(psi_dash{ll,jj}); end
|
||
end
|
||
end
|
||
|
||
%% fourier propagation
|
||
for ll = 1:par.object_modes
|
||
obj_proj{ll} = get_views(self.object, obj_proj{ll},1,ll, g_ind, cache, scan_ids{jj},[]);
|
||
end
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
probe{ll} = self.probe{min(ll,end)}(:,:,p_ind{ll});
|
||
psi{ll} = bsxfun(@times, obj_proj{min(ll, end)}, probe{ll});
|
||
if isempty(psi_dash{ll,jj}); psi_dash{ll,jj} = psi{ll}; end % initial guess
|
||
% P_M (P_0(psi))
|
||
Psi{ll} = Gfun(@DM_update_psi, gamma,psi{ll}, psi_dash{ll,jj} );
|
||
Psi{ll} = fwd_fourier_proj(Psi{ll}, self.modes{min(ll,end)} );
|
||
end
|
||
|
||
%% load data to GPU (if not loaded yet)
|
||
modF = get_modulus(self, cache, g_ind,true,jj);
|
||
mask = get_mask(self, cache.mask_indices, g_ind);
|
||
|
||
% get intensity (modulus) on detector including different corrections
|
||
aPsi = get_reciprocal_model(self, Psi, modF, mask,iter, g_ind, par,cache);
|
||
|
||
if iter > 0 && (par.get_error && (mod(iter,min(20, 2^(floor(2+iter/50)))) == 0 || iter < 20) || iter == par.number_iterations )
|
||
[fourier_error(iter,g_ind)] = get_fourier_error(modF, aPsi, [],mask);
|
||
end
|
||
|
||
if ~isempty(mask)
|
||
mask = single(1)-mask;
|
||
mask = relax_mask + (min(relax_mask, par.pfft_relaxation) - relax_mask ) * mask;
|
||
else
|
||
mask = par.pfft_relaxation;
|
||
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));
|
||
for ll = 1:size(psi_dash,2); psi_dash{ll,jj} = []; end
|
||
continue
|
||
end
|
||
|
||
%% fourier (modulus) contraint
|
||
Psi = modulus_constraint(modF,aPsi,Psi, mask, [], par, 0 );
|
||
|
||
|
||
aPsi = []; modF = []; mask = [];
|
||
|
||
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
Psi{ll} = back_fourier_proj(Psi{ll}, self.modes{min(end,ll)});
|
||
psi_dash{ll,jj} = Gfun(@DM_update, psi_dash{ll,jj} ,beta, Psi{ll},psi{ll});
|
||
end
|
||
Psi = [];
|
||
|
||
% get back from GPU if required
|
||
if ~par.keep_on_gpu
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
psi_dash{ll,jj} = Ggather(psi_dash{ll,jj});
|
||
end
|
||
end
|
||
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
|
||
psi_dash = cell(size(psi_dash));
|
||
verbose(2,'Probe amplitude corrected by %.3g',probe_amp_corr)
|
||
return
|
||
end
|
||
|
||
%% iterative solver of the overlap constraint, important for initial convergence
|
||
for kk = 1:10
|
||
for ll = 1:object_modes
|
||
obj_illum{ll}(:) = 0;
|
||
obj_update{ll}(:) = 0i;
|
||
end
|
||
for ll = 1:par.probe_modes
|
||
probe_illum{ll}(:) = 0;
|
||
probe_update{ll}(:) = 0i;
|
||
end
|
||
|
||
probe_0 = self.probe;
|
||
%% obj_update , obj_illum,probe_update,probe_illum is not reset to make it more stable
|
||
for jj = 1:length(indices)
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
if par.share_probe || ll > 1 % 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
|
||
p_ind{ll} = scan_ids{jj};
|
||
end
|
||
end
|
||
end
|
||
for ll = 1:par.probe_modes
|
||
probe{ll} = self.probe{ll}(:,:,p_ind{ll});
|
||
cprobe{ll} = conj(probe{ll});
|
||
aprobe{ll} = real(probe{ll}.*cprobe{ll});
|
||
end
|
||
|
||
% move to GPU (if not there yet)
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
psi_dash{ll,jj} = Garray(psi_dash{ll,jj});
|
||
end
|
||
g_ind = indices{jj};
|
||
|
||
for ll = 1:par.object_modes
|
||
obj_proj{ll} = get_views(self.object, obj_proj{ll},1,ll, g_ind, cache, scan_ids{jj},[]);
|
||
end
|
||
if iter >= par.probe_change_start
|
||
for ll = 1:par.probe_modes
|
||
%% update probe
|
||
[probe_update{ll},probe_illum{ll}] = QQ_probe(psi_dash{ll,jj}, obj_proj{min(end,ll)}, probe_update{ll},probe_illum{ll}, p_ind{ll});
|
||
end
|
||
end
|
||
|
||
if iter >= par.object_change_start
|
||
%% update object
|
||
[obj_update,obj_illum] = QQ_object(psi_dash{ll,jj}, obj_update,obj_illum, aprobe{min(ll,end)}, cprobe{min(ll,end)}, g_ind,min(ll, object_modes),scan_ids{jj},cache);
|
||
end
|
||
|
||
% get back from GPU if required
|
||
if ~par.keep_on_gpu
|
||
for ll = 1:max(par.probe_modes, par.object_modes)
|
||
psi_dash{ll,jj} = Ggather(psi_dash{ll,jj});
|
||
end
|
||
end
|
||
|
||
end
|
||
for ll = 1:max([par.probe_modes, object_modes])
|
||
if iter >= par.probe_change_start && ll <= par.probe_modes
|
||
% add some inertia to prevent oscilations
|
||
self.probe{ll} = update_probe(self, self.probe{ll} , probe_update{ll} , probe_illum{ll} , par, ll);
|
||
end
|
||
if iter >= par.object_change_start && ll <= object_modes
|
||
self.object{ll} = Gfun(@update_object, self.object{ll}, obj_update{ll}, obj_illum{ll}, cache.MAX_ILLUM(ll)*1e-4, par.probe_inertia);
|
||
end
|
||
end
|
||
|
||
if iter > par.probe_change_start
|
||
% if both object and probe are recontructed, solve the
|
||
% realspace constraint iterativelly
|
||
min_iter = 1 + par.keep_on_gpu; % if par.keep_on_gpu==true, then the code is so slow that it is not worthy skip the norm calculation
|
||
min_change = 0.01;
|
||
if kk >= min_iter || verbose >= 2
|
||
dprobe = max(norm2(self.probe{1} - probe_0{1}) ./ probe_norm);
|
||
verbose(2,'Update probe difference: %3.2f%%', dprobe*100)
|
||
end
|
||
if kk >= min_iter && dprobe < min_change % change is below 1% and at least 2 iterations were done
|
||
break
|
||
end
|
||
end
|
||
end
|
||
|
||
if verbose > 2
|
||
for kk = 1:object_modes
|
||
Nresid = sum2(utils.findresidues(self.object{kk}(cache.object_ROI{:})) > 0.1);
|
||
if Nresid > 0
|
||
verbose(1,'Number of residua in object %i: %i', kk, Nresid)
|
||
end
|
||
end
|
||
end
|
||
|
||
end
|
||
function [probe_update,probe_illum] = QQ_probe(psi, obj_proj, probe_update,probe_illum, p_ind)
|
||
import engines.GPU.GPU_wrapper.*
|
||
%% update probe
|
||
[upd, illum] = Gfun(@QQ_probe_Gfun, psi,obj_proj);
|
||
if size(probe_update,3)==1 % shared probe
|
||
probe_update = probe_update+sum(upd,3);
|
||
probe_illum = probe_illum+sum(illum,3);
|
||
else
|
||
probe_update(:,:,p_ind) = probe_update(:,:,p_ind)+sum(upd,3);
|
||
probe_illum(:,:,p_ind) = probe_illum(:,:,p_ind)+sum(illum,3);
|
||
end
|
||
end
|
||
function norm = get_probe_norm_aux(P0, P1)
|
||
norm = abs(P0-P1).^2;
|
||
end
|
||
|
||
function [upd, illum] = QQ_probe_Gfun(psi,proj)
|
||
upd = psi .* conj(proj) ;
|
||
illum = abs(proj).^2;
|
||
end
|
||
function [obj_update,obj_illum] = QQ_object(psi,obj_update,obj_illum, aprobe, cprobe, g_ind,ll,scan_ids,cache)
|
||
import engines.GPU.shared.*
|
||
%% update object
|
||
psi = psi .* cprobe;
|
||
[obj_update,obj_illum] = set_views_rc(obj_update,obj_illum,psi,aprobe,1,ll, g_ind, cache, scan_ids,[]);
|
||
end
|
||
function psi = DM_update_psi(gamma,psi, psi_dash )
|
||
% real space update function for difference maps
|
||
psi = (1+gamma)*psi - gamma*psi_dash;
|
||
end
|
||
function psi_dash = DM_update(psi_dash,beta, psi_tmp,psi )
|
||
% update funciton for difference maps
|
||
psi_dash = psi_dash + beta * ( psi_tmp - psi ) ;
|
||
end
|
||
function object = update_object(object, object_upd, object_illum, delta, inertia)
|
||
% apply also some inertia in the object update
|
||
object = object*inertia + (1-inertia)*object_upd./ (object_illum+delta);
|
||
end
|
||
function probe = update_probe(self, probe, probe_update, probe_illum, par, probe_id)
|
||
import engines.GPU.shared.*
|
||
probe_new = probe_update ./ (probe_illum+1e-6);
|
||
% apply probe support on the first probe mode
|
||
if probe_id == 1
|
||
probe_new = apply_probe_contraints(probe_new, self.modes{probe_id});
|
||
end
|
||
% add some inertia to prevent oscilations
|
||
probe = par.probe_inertia*probe + (1-par.probe_inertia)*probe_new;
|
||
end
|