initial commit

ptycho/+engines/+GPU_MS/+initialize/check_inputs.m

@@ -0,0 +1,178 @@
+% INITIAL_CHECKS check if the inputs are valid or try to correct them 
+% 
+% [self,par] = initial_checks(self, par)
+% 
+% ** self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ** par       structure containing parameters for the engines 
+%
+% returns: 
+% ++ self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ++ par       structure containing parameters for the engines 
+
+function [self,par] = check_inputs(self, par)
+    import engines.GPU_MS.shared.*
+    import engines.GPU_MS.GPU_wrapper.*
+    import math.*
+    import utils.*
+
+[self.Np_o(1),self.Np_o(2),~] = size(self.object{1});
+[self.Np_p(1),self.Np_p(2),~] = size(self.probe{1});
+par.Nrec  = 1;
+par.Nscans = length(self.reconstruct_ind);
+
+if ischar(par.extension)
+    par.extension = {par.extension};
+end
+
+for ii = 1:numel(self.object)
+	assert(all(isfinite(self.object{ii}(:))), 'Provided object contains nan / inf')
+end
+
+for ii = 1:numel(self.probe)
+	assert(all(isfinite(self.probe{ii}(:))), 'Provided probes contain nan / inf')
+end
+
+Np_d =  size(self.diffraction);
+if any(self.Np_p ~= Np_d(1:2)) % && isempty(self.modF_ROI)
+    error('Size of probe and data is different')
+end
+
+tmp = self.diffraction(1:self.Np_p(1)*7:end);  % get some small sample 
+tmp = tmp * 2^(par.upsampling_data_factor*2); % remove upsampling effects 
+if par.compress_data && max(abs((tmp - round(tmp)))) > 0.2
+    verbose(1,'Data are not integers, cannot use compression')
+    par.compress_data = false;
+end
+
+%%%%%%%%%%%%% accelerated solver %%%%%%%%%%%%%%%%%%%
+if  par.accelerated_gradients_start < par.number_iterations && ~is_method(par, 'MLs')
+    verbose(3, 'accelerated_gradients_start < number_iterations	 is supported only for MLc engine')
+    par. accelerated_gradients_start = inf; 
+end
+
+% if  par.accelerated_gradients_start < par.number_iterations && par.momentum > 0 && is_method(par, 'ML')
+%     error('accelerated_gradients_start < inf cannot be used if momemtum > 0 ')
+% end
+
+%%%%%%%%%%%%% variable probe %%%%%%%%%%%%%%%%%%%%%%%%
+if ~par.variable_probe
+    par.variable_probe_modes = 0; 
+end
+
+if par.variable_probe && par.variable_probe_modes  > 0 && ~is_method(par, {'PIE', 'ML'})
+    warning('Variable probe implemented only for PIE and ML')
+    par.variable_probe = false;
+end
+
+if par.variable_probe &&  par.variable_probe_modes == 0 
+	error('Choose more than 0 variable_probe_modes for OPRP')
+    par.variable_probe_modes = 1; 
+end
+
+if par.variable_probe && ~par.share_probe && is_method(par, 'PIE')
+    par.share_probe = true;
+    % variable probe means automatically shared variable probe 
+end
+
+if ~is_method(par, {'PIE', 'ML'}) && strcmpi(par.likelihood, 'poisson')
+    warning('Poisson likelihood supported only for PIE methods')
+    par.likelihood = 'L1';
+end
+
+if ~ismember(lower( par.likelihood), {'l1','poisson'})
+    error('Unsupported error estimation')
+end
+   
+%%%%%%%%%%%%%% check if position correction is allowed 
+if ~ is_method(par, {'PIE', 'ML'}) && par.probe_position_search < par.number_iterations
+	verbose(2, 'Position correction supported only for PIE/ML methods ')
+    par.probe_position_search = inf; 
+end
+
+if any(~ismember(par.probe_geometry_model, {'scale', 'asymmetry', 'rotation', 'shear'}))
+	missing_option = setdiff(par.probe_geometry_model, {'scale', 'asymmetry', 'rotation', 'shear'} ); 
+	error('Unsupported geometry model option:  "%s"', missing_option{1}) 
+end
+
+if par.probe_position_search < par.number_iterations && par.detector_scale_search < par.number_iterations && any(ismember(par.probe_geometry_model,'scale')) 
+    error('Do not use probe_position_search with probe_geometry_model==''scale'' and detector_scale_search together')
+end
+
+%%%%%% checks for the multilayer method  %%%%%%%%%%%%%%%%
+%Note: self.z_distance is first initialized in load_from_p.m, where a
+%vacuum layer is appended: self.z_distance = [p.delta_z, inf] for far-field 
+par.Nlayers = length(self.z_distance);
+if par.Nlayers > 1 && isinf(self.z_distance(end)) 
+    % Added by ZC: exclude the last vacuum (inf) layer for multisluce
+    par.Nlayers = par.Nlayers - 1;
+end
+
+assert(sum(~isfinite(self.z_distance)) <= 1, 'Provided distanced of layers are not possible to be used')
+
+if par.Nlayers > 1 && ~is_method(par, {'PIE', 'ML'})
+    error('Multilayer extension is supported only for PIE/ML methods')
+end
+
+% Added by ZC. allow user to specify the layer used for position correction
+if ~isfield(par,'layer4pos') || isempty(par.layer4pos) 
+    par.layer4pos = ceil(par.Nlayers/2);
+end
+
+%%%%%%%%%%  fast scanning %%%%%%%%%%%%%%%%%%%%%%%%%
+if is_used(par, 'fly_scan') && ~is_method(par, {'PIE', 'ML'})
+    error('Fly scan is  supported only for PIE/ML methods')
+end
+
+if is_used(par, 'fly_scan')
+    if par.Nmodes == 1
+        warning('Flyscan has no effect with a single mode')
+        par.extension = setdiff(par.extension, 'fly_scan');
+        par.apply_subpix_shift= true;
+    end
+    par.Nrec = par.Nmodes;
+%     par.apply_multimodal_update = true;
+end
+
+%%%%%%%% nearfield %%%%%%%%%%%%%%%%%%%%%%%%%%%
+if par.estimate_NF_distance < par.number_iterations && isinf(self.z_distance(end))
+    error('estimate_NF_distance valid only for nearfield mode')
+end
+
+%%%%%%%%%%%%%%% OTHER %%%%%%%%%%%%%%%%%%%%
+if strcmpi(par.likelihood, 'poisson') && par.background_detection && ~isinf(par.background_detection)
+    error('Background detection does not work well with Poisson likelihood')
+end
+
+if prod(self.Np_p) *self.Npos >  intmax('int32') && par.keep_on_gpu && is_method(par, {'MLs', 'ePIE'})
+    warning('Dataset as more than 2147483647 elements (max of int32). Set par.keep_on_gpu to false')
+    par.keep_on_gpu = false;
+end
+
+if any(self.noise(:) == 0) && par.relax_noise
+    warning('Some values of expected noise are 0')
+    self.noise = max(0.5, self.noise); 
+end
+
+if par.Nrec > max([par.Nmodes, par.probe_modes , par.object_modes])
+    warning('Number of modes is too high')
+end
+
+if length(self.probe_positions) ~= self.Npos
+    self.probe_positions = [];
+end
+
+if par.mirror_objects && par.Nscans ~= 2 
+	error('Object mirroring is supported only for two scans')
+end
+
+%%%%%% position correction %%%%% 
+if ~is_method(par, {'PIE', 'ML'}) && par.probe_position_search < par.number_iterations
+    warning('Position corrections works only for PIE/ML methods')
+end
+
+if is_method(par, {'PIE', 'ML'}) && par.probe_position_search < par.number_iterations && ~(par.apply_subpix_shift || is_used(par,'fly_scan'))
+    verbose(2,'Subpixel shifting is strongly recommended for position refinement => enforcing par.apply_subpix_shift = true') 
+	par.apply_subpix_shift = true; 
+end
+
+end

ptycho/+engines/+GPU_MS/+initialize/get_defaults.m

@@ -0,0 +1,131 @@
+% INITIALIZE generate list of default parameters 
+% [param] = initialize 
+% 
+%
+% returns: 
+% ++ param       structure containing parameters for the engines 
+
+function [param] = get_defaults
+
+    %%%%%%%%%%%%%% GPU SETTINGS %%%%%%%%%%%%%%%%%%%%%%%%%%
+    param.use_gpu = true;        % use GPU if possible 
+    param.keep_on_gpu = true;    % keep the data all the time on GPU
+    param.compress_data = true;  % apply online compress on the GPU data 
+    param.gpu_id = []; % default GPU id, [] means choosen by matlab
+    param.check_gpu_load = true;
+    param.obj_size_limit_on_gpu = inf; % maximum object size (in MB) allowed on gpu. Automatically use cpu if exceed the limit.
+    
+    %% basic recontruction parameters 
+    %% PIE 
+    param.beta_object = 1;
+    param.beta_probe = 1;  % step size, faster convergence , more instable ?? 
+    %% DM
+    param.pfft_relaxation = 0.1; 
+    param.probe_inertia = 0.3; % add inertia to the probe reconstruction to avoid oscilations 
+    %% general 
+    param.share_probe = true;
+    param.share_object = false;
+    param.delta = 0;  % press values to zero out of the probe area !!  illim < max*delta is removed 
+    param.relax_noise = 0.0;  % relaxation for noise, lower => slower convergence, more robust 
+    param.positivity_constraint_object = 0; % enforce weak positivity in object
+    param.amplitude_threshold_object = inf; % enforce maximum amplitude to object. Values larger than the threshold is set to 1
+    param.Nmodes = 1;  %  number of multi apertures , always better to start wih one !! 
+    param.probe_modes = 1; % number of probes 
+    param.object_modes = 1;  %  number of multi apertures , always better to start wih one !! 
+    param.probe_change_start = 1;  % iteration when the probe reconstruction is started
+    param.object_change_start = 1;% iteration when the object reconstruction is started
+    param.number_iterations = 300 ; 
+    param.grouping = inf;
+    param.method = 'MLs';
+    param.likelihood = 'L1' ; % l1 or poisson,   - choose which likelihood should be used for solver, poisson is suported only for PIE 
+    param.verbose_level = 1;
+    param.plot_results_every = 50;
+
+
+    %tilted_sample_LB : I don't think this is needed
+
+    param.tilt_x = 0.0; %%LB
+    param.tilt_y = 0.0; %%LB
+
+
+    param.remove_residues = false; % autodetect and remove phase residua 
+    param.extension = ''; 
+
+    %% data handling 
+    param.upsampling_data_factor = 0;           % assume that the data were created by upsampling using function utils.unbinning 
+
+    param.damped_mask = 5e-3;  % if damped_mask = 0 -> do nothing, if 1>x>0  ->  push masked regions weakly towards measured magnitude value in each iteration
+    
+    param.background_detection = false; 
+    param.background_width = inf;
+
+    %% ADVANCED OPTIONS   
+    
+    param.object_regular =  [0, 0]; %  enforce smoothness !!!, use between [0-0.1 ]
+    param.remove_object_ambiguity = true;    % remove intensity ambiguity between the object and the probes 
+    param.variable_probe = false;           % Use SVD to account for variable illumination during a single (coupled) scan
+    param.apply_subpix_shift = false;       % apply FFT-based subpixel shift, important for good position refinement but it is slow
+
+    param.probe_geometry_model = {'scale', 'asymmetry', 'rotation', 'shear'};  % list of free parameters in the geometry model
+    param.probe_position_search = inf;
+    param.apply_relaxed_position_constraint = true; %added by YJ: allow position update without geom model constraint
+    param.update_pos_weight_every = inf; %added by YJ: allow position weight to be updated multiple times. Default = inf: only calculate once
+    param.max_pos_update_shift = 0.1; %added by YJ: allow user to specify the maximum position update allowed in each iteration. Default = 0.1 (pixel).
+    param.probe_position_search_momentum = 0; % added by YJ. enable momentum acceleration for position correction. Default = 0: no acceleration.
+
+    param.probe_fourier_shift_search = inf; 
+    param.estimate_NF_distance = inf;
+    param.detector_rotation_search = inf;   % rotation of the detector axis with respect to the sample axis, similar as rotation option in the position refinement geometry model but works also for 0/180deg rotation shared scans 
+    param.detector_scale_search = inf;      % pixel scale of the detector, can be used to refine propagation distance in ptycho 
+
+    param.apply_multimodal_update = false; % use thibault modes to get higher signal, it can cause isses, not real gain  if blur method is used 
+    param.probe_backpropagate = 0; 
+    param.beta_LSQ = 0.9;       % use predictive step length
+    param.delta_p = 0.1;     % LSQ damping constant 
+    param.variable_probe_modes = 1; % OPRP settings 
+    param.variable_probe_smooth = 0;% OPRP settings 
+    param.variable_intensity = false; % account fort variable intensity
+    param.relaxed_object_constrain = 0; % enforce known object (inputs.object_orig)
+    param.probe_position_error_max = 10e-9; % max expected error of the stages 
+    param.probe_fourier_shift_search = inf; 
+    param.momentum = 0;             % use mementume accelerated gradient decsent method 
+    
+    param.regularize_layers = 0;    % 0<R<1 -> apply regularization on the reconstructed layers 
+    param.preshift_ML_probe = true; % multilayer ptycho extension: if true, assume that the provided probe is reconstructed in center of the sample. 
+    param.layer4pos = [];             % Added by ZC. speficy which layer is used for position correction 
+    param.init_layer_select = [];       % Added by YJ. Select layers in the initial object for pre-processing If empty (default): use all layers.
+    param.init_layer_preprocess = '';   % Added by YJ. Specify how to pre-process initial layers
+                                       % '' or 'all' (default): use all layers (do nothing)
+                                       % 'avg': average all layers 
+                                       % 'avg1': keep one averaged layer 
+                                       % 'interp': interpolate layers using spline method. Need to specify desired depths in init_layer_interp
+    param.init_layer_interp = [];     % Specify desired depths for interpolation. The depths of initial are [1:Nlayer_init]. If empty (default), no interpolation                    
+    param.init_layer_append_mode = '';  % Added by YJ. Specify how to initialize extra layers
+                                       % '' or 'vac' (default): add vacuum layers
+                                       % 'edge': append 1st or last layers
+                                       % 'avg': append averaged layer
+    param.init_layer_scaling_factor = 1;  % Added by YJ. Scale all layers. Default: 1 (no scaling). Useful when delta_z is changed
+
+    param.initial_probe_rescaling = true;  % find the optimal scaling correction for the provided probe guess in the initial iteration 
+    param.accelerated_gradients_start = inf;  % use accelerated gradients to speed up the convergence
+    param.align_shared_objects = false;      % align multiple objects from various scans 
+
+    % extra analysis
+    param.get_fsc_score = false;         % measure evolution of the Fourier ring correlation during convergence 
+    param.mirror_objects = false;        % mirror objects, useful for 0/180deg scan sharing 
+    param.align_shared_objects = false;   % align the objects before sharing them onto single one 
+
+    % fly scans 
+    param.flyscan_offset = 0; 
+    param.flyscan_dutycycle = 1;
+    %rng('default');
+    %rng('shuffle');
+
+    % convergence check - stop reconstruction if fourier error is larger than the previous one by given (relative) threshold. 
+    param.fourier_error_threshold = inf; % default: no convergence check.
+
+    % I/O
+    param.save_init_probe = false;       % Added by YJ. If true, save initial probe function in the .mat output file. Default is false.
+    param.save_images = {'obj_ph','obj_ph_sum','obj_ph_stack','probe'}; % Added by YJ. Save intermediate results as tiff images. 
+    % Options: {'obj_ph','obj_ph_sum','obj_ph_stack','obj_mag','obj_ph_sum','obj_mag_stack','probe_mag','probe'}
+end

ptycho/+engines/+GPU_MS/+initialize/get_parallel_blocks.m

@@ -0,0 +1,161 @@
+% GET_PARALLEL_BLOCKS Find the optimal groups to be solved in parallel on GPU/CPU
+% 
+%[cache, par] = get_parallel_blocks(self, par, cache)
+% 
+% ** 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
+%
+% returns: 
+% ++ self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ++ par       structure containing parameters for the engines 
+
+
+
+% 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 mixed coherent modes:
+% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
+% for LSQ-ML method 
+% M. Odstrcil, A. Menzel, M.G. Sicairos,  Iterative least-squares solver for generalized maximum-likelihood ptychography, Optics Express, 2018
+% for OPRP method 
+%  M. Odstrcil, P. Baksh, S. A. Boden, R. Card, J. E. Chad, J. G. Frey, W. S. Brocklesby,  "Ptychographic coherent diffractive imaging with orthogonal probe relaxation." Optics express 24.8 (2016): 8360-8369
+% 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). 
+% 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 [cache, par] = get_parallel_blocks(self, par, cache)
+    import engines.GPU_MS.GPU_wrapper.*
+    import utils.*
+    import engines.GPU_MS.shared.*
+
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %%%% FIND MAXIMAL GROUP SIZE IF GPU IS USED %
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+
+    grouping_0 = par.grouping; 
+
+    par.grouping = round(min(self.Npos, par.grouping));
+
+    global gpu 
+
+    if par.use_gpu
+        %% !! very empirical estimation of the GPU memory requirements !!!
+        % use the estimated memory requirements to prevent low memory issues 
+
+        [required_mem_constant] = estimate_req_memory(self, par, 0); 
+        [required_mem] = estimate_req_memory(self, par, 1); 
+        % precheck size of the block and try to optimize size of the groups
+        % allowed for given GPU 
+        max_group_size = floor( (gpu.AvailableMemory - required_mem_constant) ./ (required_mem - required_mem_constant));
+        max_group_size = min(self.Npos, max_group_size); 
+        verbose(1,'Maximal possible grouping %i', max_group_size);
+
+        % if group size was set to infinity, is the maximal group size possible 
+        if isinf(grouping_0)
+            par.grouping= max_group_size; 
+        else
+            % otherwise use max_group_size as a top limit 
+            par.grouping = min(par.grouping, max_group_size); 
+        end
+        
+        % adjust grouping to minimize overhead -> make the group sizes more
+        % equal 
+        if is_method(par, {'ML', 'PIE'})
+            % allows to calculate several scans together 
+            par.grouping = ceil(self.Npos/ceil(self.Npos/par.grouping));
+        else
+            % consider each scan separatelly
+            Npos_scan = cellfun(@length, self.reconstruct_ind); 
+            par.grouping = max(ceil(Npos_scan./ceil(Npos_scan./par.grouping)));
+        end
+
+        if par.grouping ~= grouping_0 
+            verbose(1,'Optimal grouping was changed from %i to %i ', grouping_0, par.grouping);
+        end
+        if par.grouping < 1
+            error('Too low memory, use smaller dataset or try ePIE')
+        end
+
+        verbose(1,'Selected grouping %i', par.grouping);
+
+    else
+        if is_method(par, {'DM', 'ML'})
+            par.grouping = self.Npos;
+        end
+    end
+
+    % precalculate distance matrix for pseudo ePIE / hPIE / MLs to get
+    % least overlapping indices 
+    if is_method(par, {'ML', 'PIE'})
+        if self.Npos/par.Nscans < 1e3 %added by YJ to save memory
+            for ll = 1:par.Nscans
+                dist_mat =  single(distmat(self.probe_positions_0(self.reconstruct_ind{ll},:))); 
+                dist_mat(dist_mat==0 |  dist_mat > max(self.Np_p)/2) = inf;
+                cache.distances_matrix{ll} = dist_mat;
+            end
+        end
+    end
+    if is_method(par, 'MLc')
+        % get higly overlapping subsets of indices for PIE / ML
+         [cache.preloaded_indices_compact{1}.indices,cache.preloaded_indices_compact{1}.scan_ids] = ...
+             get_close_indices(self, cache, par );     
+    elseif is_method(par, {'MLs', 'PIE'})
+        % preload order of indices , generate several of them to add randomness
+        for i = 1:min(par.number_iterations,10)  
+            [cache.preloaded_indices_sparse{i}.indices,cache.preloaded_indices_sparse{i}.scan_ids] = ...
+                get_nonoverlapping_indices(self, cache, par );
+        end
+    end
+    % get just some predefined sets of indices - RAAR, DM , !! order
+    % does not matter 
+    [cache.preloaded_indices_simple{1}.indices,cache.preloaded_indices_simple{1}.scan_ids] = ...
+         get_scanning_indices(self, cache, par );
+     
+
+end

ptycho/+engines/+GPU_MS/+initialize/init_solver.m

@@ -0,0 +1,696 @@
+% INITIALIZE_SOLVER initialize GPU ptycho reconstruction, generate cache values, fftshift data, etc 
+% 
+% [self, cache] = initialize_solver(self,par) 
+% 
+% ** self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ** par       structure containing parameters for the engines 
+%
+% returns: 
+% ++ self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ++ cache     structure with precalculated values to avoid unnecessary overhead
+
+function [self, cache] = init_solver(self,par)
+        
+    import engines.GPU_MS.shared.*
+    import math.*
+    import utils.*
+    import plotting.*
+    import engines.GPU_MS.GPU_wrapper.*
+    
+    verbose(struct('prefix','GPU/CPU_MS-engine-init'))
+
+    par.Nscans = length(self.reconstruct_ind); %number of scans
+    cache.skip_ind = setdiff(1:self.Npos,[self.reconstruct_ind{:}]); %  wrong datasets  to skip 
+        
+    if ~any(self.probe_support(:))
+        self.probe_support = []; 
+    end
+    %% avoid probe to be larger than a certain oversampling !!!! 
+    if isempty(self.probe_support) 
+        par.probe_backpropagate = 0;
+    end
+
+    if  ~isempty(self.background) && any(self.background(:) > 0)
+        Background = self.background;
+    elseif  par.background_detection
+        Background = 0; 
+    else
+        Background = [];  %  array of background light 
+    end
+    
+    Noise = [];
+    %% prepare data / noise / mask 
+    if par.relax_noise &&  ~isempty(self.noise) &&  strcmp(par.likelihood, 'L1')
+        Noise = self.noise;
+        Noise = (sqrt(posit(self.diffraction + Noise)) - sqrt(posit(self.diffraction - Noise)))/2;
+        Noise(self.diffraction == 0) = 1;
+        disp('Using measured noise')
+        Noise = max(0.5, Noise); 
+    end
+    
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %%%%%%% PREPARE MASK AND DATA %%%%%%%%%%%%%%%
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+  
+    %% prepare mask , note that bool in matlab has size of uint8 !!
+    cache.mask_indices = [];
+    if any(self.mask(:))
+        Mask = [];
+        % single mask 
+        if all(all(mean(self.mask,3) == self.mask(:,:,1)))
+            Mask = self.mask(:,:,1);
+        else
+            % mask for each scan
+            for ll = 1:par.Nscans
+                ind = self.reconstruct_ind{ll};
+                %if there is only one repeated mask over whole scan 
+                if all(all(all(bsxfun(@eq, self.mask(:,:,ind), self.mask(:,:,ind(1))))))
+                    Mask(:,:,ll) = self.mask(:,:,ind(1)); 
+                end
+                cache.mask_indices(ind) = ll;
+            end
+        end
+        if isempty(Mask)
+            % mask for each position 
+            Mask = self.mask;  % otherwise just store original
+            cache.mask_indices(ind) = 1:self.Npos;
+        end
+        % important to save memory 
+        if all(Mask(:) == 1 | Mask(:) == 0)
+            Mask = logical(Mask );
+        else
+            Mask = uint8(Mask*255);  % if there are nonlogical values in mask, store them as uint8 to save memory 
+        end
+    else
+        Mask = [];
+    end
+    
+    %% prepare diffraction data 
+    Diffraction = self.diffraction;  % diffraction is intensity, not amplitude, comment by ZC
+    if par.upsampling_data_factor
+        % downsample the data down to original size to save memory 
+        Diffraction = utils.binning_2D(Diffraction, 2^par.upsampling_data_factor) * (2^(2*par.upsampling_data_factor)); 
+        if ~isempty(Mask)
+            Mask = utils.binning_2D(Mask, 2^par.upsampling_data_factor) == 1; 
+        end
+    end
+    
+    Diffraction = single(max(0,Diffraction));
+
+    
+    if ~isempty(Mask)
+        if size(Mask,3) == par.Nscans && par.Nscans  > 1
+            for ll = 1:par.Nscans
+                ind = self.reconstruct_ind{ll};
+                Diffraction(:,:,ind) = Diffraction(:,:,ind) .* ~Mask(:,:,ll); 
+            end
+        else
+            Diffraction = Diffraction .* ~Mask;
+        end
+    end
+      
+    
+    if  ~isinf(self.z_distance(end)) %  && mod(Ninf,2)~=0
+        % assume inputs already fftshifted, but in case of nearfield
+        % fftshift it back for the ASM propagator 
+        Noise = fftshift_2D(Noise);
+        Diffraction = fftshift_2D(Diffraction);
+        Mask = fftshift_2D(Mask);
+    end
+
+    
+    %%%% compress data if requested %%%%%%
+    if par.compress_data
+        DATA_MAX = quantile(max2(abs(Diffraction)), 1-1e-2); 
+        C_factor_0 = 2;  % compression factor  >=2 seems to be safe, >=4 is pratically lossless 
+        if par.compress_data == 1 || DATA_MAX < 2^(2*8) / C_factor_0^2
+            Diffraction = sqrt(single(Diffraction));
+            if  DATA_MAX < 2^(2*8) / C_factor_0^2
+                % simple sqrt compression to 8 bits 
+                verbose(1, 'Online data compression to 8-bits')
+                Diffraction = uint8(C_factor_0*Diffraction);
+                cache.C_factor = C_factor_0;
+            elseif DATA_MAX < 2^(2*16) / 16^2
+                % failsafe option: sqrt compression to 16 bits 
+                verbose(1, 'Online data compression to 16-bits')
+                cache.C_factor = 16;  % use compression factor 16, to be super safe just because we have space
+                Diffraction = uint16(cache.C_factor*Diffraction);
+            else
+                error('Online compression will fail')
+            end
+        elseif par.compress_data == 2
+            % SVD subtraction compression to 8 bits (failsafe is compression to 16bits)
+            % additionally remove some SVD modes 
+            Diffraction = sqrt(single(Diffraction));
+            Nmodes = par.Nscans;
+            [U,S,V] = fsvd(reshape(Diffraction,prod(self.Np_p),[]), Nmodes);
+            ind_relevant = diag(S).^2/sum(diag(S).^2) > 1e-2; % more than 1% of power
+            cache.US_diffraction = (U(:,ind_relevant)*S(ind_relevant,ind_relevant));
+            cache.V_diffraction = V(:,ind_relevant);
+            svd_Diffraction = round(reshape(cache.US_diffraction*cache.V_diffraction',[self.Np_p, self.Npos]));
+
+            %% compress 
+            cDiffraction = single(Diffraction) - svd_Diffraction;
+            % reestimate optimal compression factor to keep values < 128 
+            C_factor = min(C_factor_0, 128/quantile(max2(abs(cDiffraction)), 1-1e-2));
+            if C_factor > 3
+                verbose(1, 'Online data compression to 8-bits + SVD')
+                cache.C_factor = C_factor;
+                cache.US_diffraction =  cache.US_diffraction;
+                cache.V_diffraction = cache.V_diffraction*C_factor;
+                Diffraction = int8(cDiffraction*C_factor);
+            elseif DATA_MAX < 2^(2*16) / 16^2
+                % sqrt compression to 16 bits 
+                %warning(sprintf('Too high online compression of data, it may cause problems\n Compression factor is %2.2f but should be >= 2\n Switching from 8 to 16bits',C_factor))
+                verbose(1, 'Online data compression to 16-bits')
+                C_factor = 16;
+                cache.C_factor = C_factor;
+                Diffraction = uint16(C_factor*Diffraction); 
+            else
+                error('Online compression will fail')
+            end
+            
+            clear svd_Diffraction cDiffraction 
+
+        else
+            error('Unimplented level of compression')
+        end
+    else
+       % precalculate sqrt from the data, store as singles 
+        Diffraction = sqrt(single(max(0,Diffraction))); % diffraction is amplitude, comment by ZC
+    end
+
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %% write back the data arrays 
+    self.diffraction = Diffraction; % diffraction is amplitude, comment by ZC
+    self.mask = Mask; 
+    self.noise = Noise;    
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    
+    
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %%%%%%% PREPARE GEOMETRY, PROPAGATION, MODES%
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    
+    % precalculate ASM factor for propagation distance recovery 
+    [ASM_difference] = near_field_evolution_gradient(self.Np_p, self.lambda,  self.pixel_size .*self.Np_p );
+    cache.ASM_difference = fftshift(ASM_difference);
+    
+    % custom propagator to account for tilted plane sample - this chuck is
+    % basically not used since the lines bellow are commented out - LB
+    if any(par.p.sample_rotation_angles(1:2)) && check_option(par.p, 'apply_tilted_plane_correction', 'propagation') 
+        % get propagators to the tilted plane 
+        [tilted_plane_propagate_fwd, tilted_plane_propagate_back] = ...
+            get_tilted_plane_propagators(Garray(self.probe{1}), ...
+                                        [par.p.sample_rotation_angles(1:2),0],...
+                                        self.lambda, self.pixel_size); 
+    else
+        tilted_plane_propagate_fwd = [];  tilted_plane_propagate_back = [];
+    end
+    
+  
+    if ~iscell(self.affine_matrix)
+        self.affine_matrix = {self.affine_matrix}; 
+    end
+    
+    %Note: par.Nmodes is # of probe modes (p.probe_modes). Assigned in
+    %load_from_p.m
+%     modes = cell(max(par.Nmodes, par.Nlayers),1);
+    modes = cell(par.Nlayers,1); % corrected by Zhen Chen
+    
+    %Comment by YJ: par.Nmodes seems to be # of modes for A-fly scan
+%     for i = 1:max(par.Nmodes, par.Nlayers) 
+    %par.Nlayers equals to the # of slices in the object (excluding the
+    %vacuum layer)
+    for i = 1:par.Nlayers % modified by ZC
+        verbose(2,'Creating new modes files ')
+        modes{i}.lambda =  self.lambda;
+        
+        % decompose affine matrix into scale, asymmetry, rotation, shear 
+        %affine = scale*[1+asym/2,0; 0,1-asym/2]*[cosd(rot), sind(rot); -sind(rot), cosd(rot)] * [1,0;tand(shear),1];
+        
+        affine_matrix = self.affine_matrix{min(i,end)};
+        [scale, asymmetry, rotation, shear]  = decompose_affine_matrix(affine_matrix);
+        
+        % store initial geometry parameters 
+        modes{i}.scales =  repmat(scale, 1,par.Nscans); 
+        modes{i}.asymmetry =  repmat(asymmetry, 1,par.Nscans); 
+        modes{i}.shear =   repmat(shear, 1,par.Nscans); 
+        modes{i}.rotation =  repmat(rotation, 1,par.Nscans); 
+        modes{i}.affine_matrix =  repmat(affine_matrix, 1,1,par.Nscans); 
+        modes{i}.shift_scans =  zeros(2, par.Nscans); 
+        modes{i}.probe_scale_upd = 0; 
+        modes{i}.probe_rotation = ones(1,par.Nscans) * par.sample_rotation_angles(3); % one rotation per scan 
+        if par.mirror_objects
+            modes{i}.probe_rotation = modes{i}.probe_rotation .* [1,-1];  % flip the coordinates for mirrored object (ie 0 vs 180deg rotation)
+        end
+        modes{i}.probe_rotation_all = zeros(self.Npos,1);
+        for jj = 1:par.Nscans
+            modes{i}.probe_rotation_all(self.reconstruct_ind{jj}) = modes{i}.probe_rotation(jj); % one rotation per scan 
+        end
+        
+        distance = self.z_distance(min(end,i));
+
+        if ~isinf(distance)
+            verbose(2, 'Layer %i distance %g um  ', i, distance*1e6 )
+        end
+        modes{i}.distances = distance;
+        if is_used(par, 'fly_scan') && (~isfield(modes{i}, 'probe_positions') || isempty(modes{i}.probe_positions) )
+            %% get positions for fly scans 
+             self = prepare_flyscan_positions(self, par); 
+             modes{i}.probe_positions = self.modes{i}.probe_positions; %added by YJ. seems like a bug
+             modes{i}.probe_positions_0 = self.probe_positions_0; %added by YJ. seems like a bug
+        else
+           %% get positions for normal tomo 
+            try  %  try to reuse the positions of there are saved 
+                modes{i}.probe_positions = self.modes{i}.probe_positions;
+                verbose(2,'Using saved positions')
+            catch
+                if (modes{i}.scales(end) == modes{1}.scales(end)) && ~isempty(self.probe_positions)
+                    modes{i}.probe_positions = self.probe_positions;
+                    verbose(0,'Using saved positions')
+                else
+                    verbose(2,'Using original positions')
+                    modes{i}.probe_positions = (affine_matrix*self.probe_positions_0')';
+                end
+            end
+            try
+                modes{i}.probe_positions_0 = self.modes{i}.probe_positions_0;
+            catch
+                modes{i}.probe_positions_0 = self.probe_positions_0;
+            end
+
+        end
+        modes{i}.probe_positions_update = { zeros(size(modes{i}.probe_positions)) };
+        modes{i}.probe_positions_all = {modes{i}.probe_positions};
+        modes{i}.probe_positions_weight = zeros(self.Npos, 1);   
+        if isfield(self, 'probe_fourier_shift') && ~isempty(self.probe_fourier_shift)  && i == 1
+            modes{i}.probe_fourier_shift = self.probe_fourier_shift; 
+        else
+            modes{i}.probe_fourier_shift =  zeros(self.Npos,2);
+        end
+
+        if ~isempty(self.probe_support) && i <= par.Nrec
+            modes{i}.probe_support = self.probe_support; 
+            if i == 1
+                verbose(2,'Using real-space probe support')
+            end
+        else 
+            modes{i}.probe_support = [];
+        end
+        
+        if ~isempty(self.probe_support_fft) && i <= par.Nrec && ~check_option(par,'probe_support_tem')
+            modes{i}.probe_support_fft = fftshift(self.probe_support_fft); 
+            if i == 1
+                verbose(2,'Using far-field probe support')
+            end                      
+        elseif check_option(par,'probe_support_tem') % not shift for TEM aperture, by Zhen Chen
+            modes{i}.probe_support_fft = self.probe_support_fft; 
+        else 
+            modes{i}.probe_support_fft = [];
+        end
+        
+        F = mean( self.pixel_size)^2 .* mean(self.Np_p) /   (modes{i}.lambda * modes{i}.distances); 
+        if F ~= 0
+            verbose(3,'Nearfield propagation: Fresnel number/Npix %3.3g', F)
+        end
+        scale = modes{i}.scales(end);
+        modes{i}.ASM_factor = [] ;
+        modes{i}.cASM_factor = [] ;
+
+        if ~isinf(modes{i}.distances(end)) % Forward Fresnel propagator in k-space, ASM, commented by ZC
+            %% near field factor            
+            %ASM =  exp( modes{i}.distances(end)* cache.ASM_difference);
+            % modified by YJ: use H instead of dH (which is an approximation)
+            %[~,ASM,~,~] = near_field_evolution(ones(self.Np_p), modes{i}.distances(end), self.lambda,  self.pixel_size .*self.Np_p, true );
+            tiltx = par.tilt_x*1e-3;
+            tilty = par.tilt_y*1e-3;
+            [~,ASM,~,~] = near_field_evolution(ones(self.Np_p), modes{i}.distances(end), self.lambda,  self.pixel_size .*self.Np_p, true,tiltx,tilty );
+            ASM = fftshift(ASM);
+            modes{i}.ASM_factor = ASM;
+            modes{i}.cASM_factor = conj(ASM);
+        end
+        
+        %% far field factor
+        modes{i}.FAR_factor = [];
+        modes{i}.cFAR_factor = conj(modes{i}.FAR_factor);
+  
+        if isinf( par.probe_backpropagate)
+             modes{i}.support_fwd_propagation_factor = inf; 
+             modes{i}.support_back_propagation_factor = -inf;
+        elseif par.probe_backpropagate ~= 0
+            [~, modes{i}.support_fwd_propagation_factor] = utils.prop_free_nf( self.probe{1}(:,:,1), par.probe_backpropagate,...
+                modes{i}.lambda,  self.pixel_size ./ scale );
+  
+            modes{i}.support_fwd_propagation_factor = fftshift( modes{i}.support_fwd_propagation_factor );
+            modes{i}.support_back_propagation_factor = conj(modes{i}.support_fwd_propagation_factor);
+        else
+            modes{i}.support_fwd_propagation_factor = [];
+            modes{i}.support_back_propagation_factor = [];
+        end
+        
+        %I am a bit confused on why this is commented out?? - the tilted
+        %plane propogators are caleed in the ptycho solver
+        % modes{i}.tilted_plane_propagate_fwd = tilted_plane_propagate_fwd;
+        %modes{i}.tilted_plane_propagate_back = tilted_plane_propagate_back;
+        modes{i}.tilted_plane_propagate_fwd = [];
+        modes{i}.tilted_plane_propagate_back = [];
+    end
+    
+      
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %%%%%%% PREPARE PROBES, INCOHERENT MODES %%%
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+   
+    probe_0 = mean(self.probe{1},3);
+    probe = cell(1,par.probe_modes);  % Added by ZC and fixed by YJ
+    for i = 1:par.probe_modes
+        try
+            probe{i} = self.probe{i};
+                %% test if the probe size ok for the variable probe settings etc 
+            assert( size(probe{i},4) == 1+par.variable_probe_modes || ...
+                ~(par.variable_probe) || i > 1)
+            assert((size(probe{i},3) ==1 || par.variable_probe) || ...
+                   (size(probe{i},3) == par.Nscans && ~par.share_probe)  )  % no variable prob extension and multiple probes used 
+            assert(size(probe{i},3) == par.Nscans || par.share_probe || par.variable_probe, 'Wrong probe size for not shared probe option')
+        catch
+            if i <= par.Nrec || is_used(par, 'fly_scan')
+                verbose(2, 'Creating probe')
+
+                if ~par.share_probe && size(probe{i},3) == 1
+                     % dont share probe between scans 
+                     probe{i} = repmat(probe_0,[1,1,par.Nscans]);
+                end
+                if  (par.variable_probe && par.variable_probe_modes > 0) && i == 1
+                    verbose(2,'Creating variable probe ')
+                    probe{i}(:,:,:,2:1+par.variable_probe_modes) = ...
+                        randn([self.Np_p, size(probe{i},3), par.variable_probe_modes])+randn([self.Np_p,size(probe{i},3), par.variable_probe_modes])*1i;
+                    continue 
+                end
+            end
+            if length(probe) < i   % none of above 
+               % simply create slightly shifted modes in fourier domain, it is useful for
+               % inital guess of incoherent modes after orthogonalization 
+                step = median(diff(self.probe_positions_0)); 
+                probe{i} = 0.01*fftshift(imshift_fft(fftshift(probe_0), randn, randn, false));
+            end
+            
+            % fill the unreconstructed positions if the OPRP method is used
+            if par.variable_probe && is_method(par, 'PIE') && i ==1 
+                ind_wrong = setdiff(1:self.Npos, [self.reconstruct_ind{:}]);
+                probe{i}(:,:,ind_wrong) = repmat(mean(probe{i},3),1,1,length(ind_wrong));
+            end
+         end
+    end
+    if par.probe_modes > par.Nrec
+        %  orthogonalization of incoherent probe modes 
+        if is_used(par, 'fly_scan')
+            probe_tmp = probe;
+            % orthogonalize the modes with all the other shifted modes 
+            for i = 1:par.Nrec
+                dx = median(modes{i}.probe_positions - modes{1}.probe_positions);
+                probe_tmp{i} = imshift_fft(probe_tmp{i}, dx); 
+            end
+            probe_tmp = ortho_modes(probe_tmp);  % perform othogonalization 
+            probe(1+par.Nrec:par.probe_modes) = probe_tmp(1+par.Nrec:par.probe_modes);
+        else
+            ind = [1,1+par.Nrec:par.probe_modes];  % skip polyvave/multilayer probe_tmp 
+            probe(ind) = ortho_modes_eig(probe(ind));  %% slightly better  
+        end
+    end
+        
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %%%%%%% PREPARE OBJECT, MULTILAYER OBJECT %%%
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+     
+    %% updated illumination
+    aprobe2 = abs(self.probe{1}(:,:,1)).^2; 
+    for ll = 1:par.Nscans
+        if par.share_object
+            ind = [self.reconstruct_ind{:}];
+        else
+            ind = self.reconstruct_ind{ll};
+        end
+        [cache.oROI_s{1}] = find_reconstruction_ROI( modes{1}.probe_positions,self.Np_o, self.Np_p);
+        % avoid oscilations by adding momentum term 
+        illum_sum_0{ll} = Ggather(set_views(Gzeros(self.Np_o), Garray(aprobe2), 1,1, ind, cache));
+    end
+    
+    %% multilayer extension 
+    % modified by YJ for more dynamic initialization
+
+    % Step 1: choose specific layers from initial object file
+    N_layer_input_obj = size(self.object,2);
+    par.init_layer_select(par.init_layer_select<0) = [];
+    par.init_layer_select(par.init_layer_select>N_layer_input_obj) = [];
+    if ~isempty(par.init_layer_select)
+        object_temp = cell(size(self.object,1),length(par.init_layer_select));
+        for ll = 1:par.Nscans
+            for jj=1:length(par.init_layer_select)
+                object_temp{ll,jj} = self.object{ll,par.init_layer_select(jj)};
+            end
+        end
+        self.object = object_temp;
+    end
+    
+    % Step 2: pre-process layers
+    switch par.init_layer_preprocess
+        case 'avg' % only use the averaged layer
+            verbose(0,'Average initial layers')
+            for ll = 1:par.Nscans
+                obj_avg = prod(cat(3,self.object{ll,:}),3); 
+                obj_avg = abs(obj_avg).*exp(1i*phase_unwrap(angle(obj_avg))/size(self.object,2));
+                for jj=1:size(self.object,2)
+                    self.object{ll,jj} = obj_avg;
+                end
+            end
+        case 'avg1' % only use the averaged layer
+            verbose(0,'Average initial layers and only keep one')
+            object_temp = cell(size(self.object,1),1);
+            for ll = 1:par.Nscans
+                obj_avg = prod(cat(3,self.object{ll,:}),3); 
+                obj_avg = abs(obj_avg).*exp(1i*phase_unwrap(angle(obj_avg))/size(self.object,2));
+                object_temp{ll,1} = obj_avg;
+            end
+            self.object = object_temp;
+        case 'interp' % interpolate layers
+            if ~isempty(par.init_layer_interp)
+                verbose(0,'Interpolate %d initial layers to %d layers', size(self.object,2), length(par.init_layer_interp))
+                for ll = 1:par.Nscans
+                    obj_temp = cat(3,self.object{ll,:});
+                    [N_obj_y,N_obj_x,N_obj_z] = size(obj_temp); 
+                    [X,Y,Z] = meshgrid(linspace(1,N_obj_x,N_obj_x),linspace(1,N_obj_y,N_obj_y),linspace(1,N_obj_z,N_obj_z));
+                    [Xq,Yq,Zq] = meshgrid(linspace(1,N_obj_x,N_obj_x),linspace(1,N_obj_y,N_obj_y),par.init_layer_interp);
+                    obj_temp = interp3(X,Y,Z,obj_temp,Xq,Yq,Zq,'spline');
+                    for jj=1:size(obj_temp,3)
+                        self.object{ll,jj} = obj_temp(:,:,jj);
+                    end
+                end
+            end
+        case {'','all'} % default: keep all layers
+            % nothing to do
+        otherwise
+            error('Invalid init_layer_preprocess!')
+    end
+    
+    % Step 3: add or remove layers based on par.Nlayers
+    if size(self.object,2) > par.Nlayers
+        warning('Initial object has more layers than Nlayers')
+        for ll = 1:par.Nscans
+            self.object{ll,1} = prod(cat(3,self.object{ll,:}),3); 
+        end
+        self.object(:,2:end) = [];
+    end
+    
+    if size(self.object,2) < par.Nlayers
+        N_add = par.Nlayers - size(self.object,2);
+        verbose(0,'Add %d more layers from %d layer(s)', N_add, size(self.object,2))
+        for ll = 1:size(self.object,1) %loop over scans
+            obj{ll} = self.object(ll,:);
+            switch par.init_layer_append_mode
+                case 'avg' %Not sure when this is useful, but I'll keep it for now
+                    verbose(0,'Append averaged layer')
+                    obj_pre = prod(cat(3,self.object{ll,:}),3); 
+                    obj_pre = abs(obj_pre).*exp(1i*phase_unwrap(angle(obj_pre))/size(self.object,2));
+                    obj_post = obj_pre;
+                case 'edge'
+                    verbose(0,'Append 1st/last layer')
+                    obj_pre = self.object{ll,1};
+                    obj_post = self.object{ll,end};
+                case {'','vac'}
+                    verbose(0,'Append vacuum layer')
+                    %obj_pre = ones(self.Np_o, 'single') + 1e-9i*randn(self.Np_o, 'single');
+                    obj_pre = ones(self.Np_o, 'single');
+                    obj_post = obj_pre;
+                otherwise
+                    error('Invalid init_layer_append_mode!')
+            end
+            for ii = 1:N_add
+                if mod(ii, 2) == 1
+                    obj{ll}{end+1} = obj_post; % add slice at the end 
+                else
+                    obj{ll}(2:end+1) = obj{ll};
+                    obj{ll}{1} = obj_pre; % add slice at the beginning 
+                end
+            end
+        end
+        self.object = cat(1, obj{:}); %combine all scans
+    end
+    
+    % if object has more layers but only one is needed
+    if size(self.object,2) > 1 && par.Nlayers == 1
+       for ll = 1:par.Nscans
+           object{ll,1} = prod(cat(3,self.object{ll,:}),3); 
+       end
+       self.object = object; 
+    end
+
+    % At this point: size(self.object,3) should equal to par.Nlayers
+    %{ 
+    %% MO's code, should be useless now. I'll keep it for now in case of
+    bugs in steps 1-3. 
+    for j = 1:par.Nlayers  
+        for i = 1:max(1, par.Nscans * ~par.share_object) % loop over scans
+            try
+                object{i,j} = self.object{min(end,i),j};
+                object{i,j}(1);
+            catch
+                verbose(0, 'add transparent slice') % add extra layers 
+                %% add fully transparent slice at the end 
+                object{i,j} = ones(self.Np_o, 'single');
+                if size(self.object,2) == 1
+                    % swap order of the new layers to keep the original
+                    % reconstruction in center
+                    object(i,:) = object(i,end:-1:1); 
+                end
+            end
+        end
+    end
+    %}
+    % Step 4: assign self.object to object and rescale layers if needed
+    if par.init_layer_scaling_factor~=1
+        verbose(0,'Rescale each layer by %f', par.init_layer_scaling_factor)
+    end
+    for j = 1:par.Nlayers
+        for i = 1:max(1, par.Nscans * ~par.share_object) % loop over scans
+            if par.init_layer_scaling_factor~=1
+                object_temp = self.object{min(end,i),j};
+                object_temp_ph = phase_unwrap(angle(object_temp))*par.init_layer_scaling_factor;
+                object{i,j} = abs(object_temp).*exp(1i.*object_temp_ph);
+            else
+                object{i,j} = self.object{min(end,i),j};
+            end
+        end
+    end
+    
+    for i = 1:numel(object)
+        object{i} = single(object{i});
+        object{i} =  complex(object{i});
+    end
+
+    for i = 1:numel(probe)
+        probe{i} = single(probe{i});
+        probe{i}  = complex(probe{i});
+    end
+    
+    %% STORE RESULTS TO SELF CLASS
+    self.object = object; 
+    self.probe = probe; 
+    self.modes = modes; 
+    self.diffraction = Diffraction; 
+    self.noise = Noise; 
+    self.mask = Mask; 
+    self.background = reshape(Background,1,1,[]);        
+    
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %%%%%%%%%%% PRECALCULATE USEFUL VALUES %%%%%%%
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+              
+    if ~isfield(self, 'probe_evolution' ) 
+        % initial coefficients for OPRP approximation 
+        self.probe_evolution(:,1) = ones(self.Npos,1);  % first mode is constant 
+    end
+    new_probe_modes_ind = 1+(size(self.probe_evolution,2):par.variable_probe_modes); 
+    self.probe_evolution(:,new_probe_modes_ind) = 1e-6*randn(self.Npos,length(new_probe_modes_ind));
+           
+    if par.variable_probe 
+        pnorm = norm2(self.probe{1});
+        self.probe{1}(:,:,:,2:end) = self.probe{1}(:,:,:,2:end) ./ pnorm(1,1,:,2:end);
+        self.probe_evolution(:,2:end) = self.probe_evolution(:,2:end) .* squeeze(mean(pnorm(1,1,:,2:end),3))';
+    end
+    
+
+    if par.background_detection || ~isempty(self.background)
+        %% auto-estimate background correction distribution 
+        if isempty(self.mask)
+            mask = 0;
+        else
+            mask = self.mask; 
+        end
+        if par.background_detection
+            background_weight = sum(self.diffraction.^2,3) ./ max(1,sum(~mask,3)); 
+            background_weight = imgaussfilt(background_weight,1);
+            background_weight = 1./sqrt(max(1e-3, background_weight)) .* ~any(mask,3);
+            background_weight = max(0,background_weight - 0.3*mean(background_weight(:)));
+            cache.background_weight = ( fftshift_2D(background_weight / sum2(background_weight)));
+        end
+        
+        if isinf(par.background_width)
+            cache.background_profile = 1;
+        else 
+            mdiffr = Garray(fftshift(mean(get_modulus(self,cache,1:self.Npos,false).^2,3)));
+
+            W = par.background_width;
+            X = (-self.Np_p(1):self.Np_p(1)-1);
+            Y = (-self.Np_p(2):self.Np_p(2)-1);
+            [X,Y] = meshgrid(X,Y);
+
+            background_profile = exp(-sqrt( (X/W(1)).^2 +(Y/W(1)).^2)); 
+            background_profile = conv2(mdiffr,background_profile, 'same');
+            background_profile = background_profile / max2(background_profile);
+            
+            background_profile = utils.crop_pad(background_profile,self.Np_p); 
+            cache.background_profile = gather(fftshift(background_profile)); 
+ 
+        end
+         
+        if ~isempty(self.diffraction_deform_matrix)
+            apply_deform = @(x,D)single(reshape(full(D * double(reshape(x,[],size(x,3)))), size(x)));
+            % apply deformation effects caused by tilted sample , be sure to enforce the mask before
+            % interpolation, the hotpixels can spread around after the correction
+            if isscalar(cache.background_profile)
+                cache.background_profile = ones(self.Np_p, 'single'); 
+            end
+            cache.background_profile = apply_deform(cache.background_profile, self.diffraction_deform_matrix');
+        end
+    else
+         cache.background_profile_weight = 1; 
+    end
+
+    for  ll = 1:par.Nscans
+        illum_sum_0{ll} = Ggather(illum_sum_0{ll}); 
+        cache.MAX_ILLUM(ll) = max(illum_sum_0{ll}(:));
+        cache.illum_sum_0{ll} = illum_sum_0{ll}; 
+    end
+
+    %% precalculate illumination ROIs
+    cache = precalculate_ROI(self,cache, Ggather(sqrt(aprobe2)));
+
+    %% prepare mask needed for subpixel shifts of object views 
+    cache.apodwin = single(0.1+0.9*tukeywin(self.Np_p(1),0.05) .* tukeywin(self.Np_p(2), 0.05)');
+
+    if par.initial_probe_rescaling
+        %% initial rescaling of probe intensity , just a very rough guess     
+        % modified by ZC, propagate probe to far field
+        mean_aPsi = mean2(abs(fft2_safe(self.probe{1}(:,:,1))).^2); 
+        % old method: self.modes{end} is the vaccum layer
+        %mean_aPsi = mean2(abs(fwd_fourier_proj(self.probe{1}(:,:,1), self.modes{end})).^2); 
+        
+        mean_diffraction_intensity = mean(mean2(self.diffraction(:,:,randi(self.Npos, [10,1])).^2));  % take roughly average intensity  % bug fixed by Zhen Chen, previous no ^2
+
+        for ii = 1:par.probe_modes
+            self.probe{ii} = self.probe{ii} * sqrt( mean_diffraction_intensity / mean_aPsi); 
+        end   
+    end
+
+end

ptycho/+engines/+GPU_MS/+initialize/load_from_p.m

@@ -0,0 +1,457 @@
+% LOAD_FROM_P load parameters from the p-structure to param and self structures for GPU
+% engine 
+% 
+% [self, param] = load_from_p(self, param, p)
+% 
+% 
+% ** self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ** param       structure containing parameters for the engines 
+% ** p          ptychoshelves p structure 
+%
+% returns: 
+% ++ self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ++ param       structure containing parameters for the engines 
+
+function [self, param, p] = load_from_p(param, p)
+    import math.*
+    import utils.*
+    import engines.GPU_MS.shared.*
+    
+    [Np_p(1),Np_p(2),Npos] = size( p.fmag);
+    self.reconstruct_ind = p.scanidxs;
+
+    %% load default variables with different name from the main ptycho code 
+    param.Nmodes = p.probe_modes; 
+    param.plot_results_every = p.plot.interval; 
+    
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %% additional features added by YJ
+    %for output intermediated results, 
+    param.fout = p.fout;    
+    if isfield(p, 'save_results_every')
+        param.save_results_every = p.save_results_every;
+        if param.save_results_every<=p.number_iterations
+            verbose(0, 'Intermediate results will be saved in the directory below every %i iterations.', param.save_results_every)
+            verbose(0,param.fout)
+        end
+    end
+    if isfield(p, 'extraPrintInfo')
+        param.extraPrintInfo = p.extraPrintInfo;
+    end
+    if isfield(p, 'affine_matrix')
+        param.affine_matrix_init = p.affine_matrix;
+    end
+    if isfield(p, 'beam_source')
+        param.beam_source = p.beam_source;
+    end
+    if isfield(p, 'TV_lambda')
+        param.TV_lambda = p.TV_lambda;         
+    end
+    
+    if isfield(p,'avg_photon_threshold') && p.avg_photon_threshold>=0
+        avg_photon_threshold = p.avg_photon_threshold;
+    else %default
+        if isfield(param,'beam_source') && strcmp(param.beam_source,'electron')
+            avg_photon_threshold = 0.0001;
+        else
+            avg_photon_threshold = 0.01;
+        end
+    end
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %%
+    % if defined in p, use from p otherwise use defaults in param 
+    try;param.object_regular = p.reg_mu; end
+    try;param.probe_inertia = p.probe_regularization; end
+
+    param.tilt_x = p.tilt_x; %%LB
+    param.tilt_y = p.tilt_y; %%LB
+
+
+    if check_option(p,'opt_errmetric','poisson')
+        param.likelihood = 'poisson'; 
+    else
+        param.likelihood = 'L1'; 
+    end
+    if get_option(p,'background_width') &&  get_option(p,'binning')
+        param.background_width = param.background_width / 2^p.binning;
+    end
+
+    %% load variables from the main ptycho code and merge it with the defaults 
+    for field = fieldnames(param)'
+        field = field{1}; 
+        if isfield(p, field)
+            param.(field)  = p.(field);
+        end
+    end
+          
+    % set verbosity for GPU engine 
+    param.verbose_level = max(-2,p.verbose_level-2); % adjust verbosity for GPU code , verbose_level 0 is enough for commmon use
+    verbose(param.verbose_level)
+
+    % load additional reconstructed parameters , otherwise use default 
+    for item = {{'background',[]}, {'intensity_corr',[]}, {'probe_fourier_shift',[]}, {'rotation',0},{'shear',0},{'relative_pixel_scale',1}}
+        item = item{1}; 
+        if isfield(p, item{1}) && ~isempty(p.(item{1}))
+            self.(item{1}) = p.(item{1});
+        else
+            self.(item{1}) = item{2}; 
+        end
+    end
+    if any(ismember(fieldnames(p), {'shear', 'rotation', 'relative_pixel_scale'}))  && isfield(p, 'positions_0')
+        warning('Reseting probe positions to original values')
+        p.positions = p.positions_0; 
+        p = rmfield(p, 'positions_0'); 
+    end
+    
+    if isempty(p.affine_matrix)
+       p.affine_matrix = diag([1,1]); 
+    end
+    
+    self.diffraction_deform_matrix = [];
+
+    if ~check_option(p,'asize_presolve') 
+        param.Np_p_presolve = []; 
+    else
+        param.Np_p_presolve = min(p.asize_presolve, p.asize); 
+    end
+  
+    % Other parameters 
+    param.fourier_ptycho = check_option(p,'fourier_ptycho'); 
+    param.upsampling_data_factor = p.detector.upsampling; 
+    
+    % Offaxis ptychography correction 
+    if check_option(p, 'sample_rotation_angles')
+        param. sample_rotation_angles = p.sample_rotation_angles;  % 3x1 vector rotation around [X,Y,beam] axes in degrees , apply a correction accounting for tilted plane oR the sample and ewald sphere curvature (high NA correction)
+    else
+        param. sample_rotation_angles = [0,0,0]; % conventional ptychography
+    end
+    
+    self.pixel_size = p.dx_spec;
+    self.Np_p = Np_p; 
+    Nscans = length(self.reconstruct_ind); 
+
+    %%%%%%%%%%%%%%%%%%%%%%
+    %% load probes%%%%%%%%
+    %%%%%%%%%%%%%%%%%%%%%%
+    
+    if param.share_probe
+        p.share_probe_ID(:) = 1; % enforce single object if sharing is requested
+    end
+    
+    % normalization for consistency with the other CPU engines 
+    probes = single(p.probes ./ (prod(sqrt(Np_p))*2*p.renorm)); 
+    
+    % force to a column cell convenient for multilayer, Added by ZC
+    %self.probe = cell (min(p.probe_modes, size(probes,4)),1); 
+    for i = 1:min(p.probe_modes, size(probes,4))
+        % variable probe 
+        if isfield(p, 'probe_PCA') && ~isempty(p.probe_PCA) && i == 1   && size(p.probe_PCA.eigen_vec,1) == p.asize(1) && param.variable_probe && is_method(param, 'PIE')
+            verbose(1,'Loading PCA probe (%i)', i)
+            self.probe{i} = reshape(p.probe_PCA.eigen_vec,prod(p.asize),[]) * p.probe_PCA.evolution' /(prod(sqrt(Np_p))*2*p.renorm);% normalization for consistency with the CPU code;
+        elseif  isfield(p, 'probe_variable') && ~isempty(p.probe_variable) && i == 1   && size(p.probe_variable.eigen_vec,1) == p.asize(1) && param.variable_probe && is_method(param, 'ML')
+             % ML methods, OPRP approx
+            verbose(0,'Loading variable probe (%i)', i)
+            % constant part 
+            self.probe{i}(:,:,:,1) = probes(:,:,:,1); 
+            % variable part  
+            self.probe{i}(:,:,:,2) = p.probe_variable.eigen_vec /(prod(sqrt(Np_p))*2*p.renorm);
+            % evolution of the variable part 
+            self.probe_evolution =  p.probe_variable.evolution ;
+            assert(length(self.probe_evolution)==Npos, 'Wrong size of variable probe coefficients')
+        else
+            % constant probe 
+            verbose(1,'Loading constant probe (%i)', i)
+            if param.share_probe
+                self.probe{i} = mean(probes(:,:,:,i),3); 
+            else
+                if size(p.probes,3) == Nscans 
+                    % one probe for each scan 
+                    self.probe{i} = probes(:,:,:,i);
+                else
+                    % rather take only first to avoid issues 
+                    self.probe{i} = probes(:,:,1,i);
+                end
+            end
+        end
+    end
+    
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %% prepare support contraints%%%%
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+        
+    % provide estimate of the probe support
+    if check_option(p, 'probe_mask') && check_option(p,'use_probe_support') && any(p.probe_mask)
+        self.probe_support = ~p.probe_mask; 
+    elseif check_option(p,'probe_support_radius') && p.probe_support_radius < sqrt(2)
+         % very useful for DM code 
+         [X,Y] = meshgrid((-p.asize(2)/2+1:p.asize(2)/2)/p.asize(2), (-p.asize(1)/2+1:p.asize(1)/2)/p.asize(1));
+         self.probe_support = sqrt(X.^2+Y.^2) < p.probe_support_radius/2; 
+    else
+        self.probe_support = [];
+    end
+   
+    % estimate of the probe support in detector plane 
+    if check_option(p,'probe_support_fft') &&  ~check_option(p, 'prop_regime', 'nearfield') && ~check_option(p,'probe_support_tem') % exclude TEM aperture mask by ZC
+        if ~check_option(p.model, 'probe_focal_length') && ~check_option(p.model, 'probe_outer_zone_width')
+            error('Missing  model.probe_focal_length and model.probe_outer_zone_width of Fresnel zone plate' )
+        end        
+        if ~check_option(p.model, 'probe_outer_zone_width')
+            p.model.probe_outer_zone_width = p.lambda * p.model.probe_focal_length / p.model.probe_diameter; 
+        end
+        FZP_cone_diameter = p.lambda* p.z/(p.model.probe_outer_zone_width * p.ds); 
+        % add some extra space 
+        FZP_cone_diameter = FZP_cone_diameter * 1.2; 
+        [X,Y] = meshgrid(-p.asize(2)/2+1:p.asize(2)/2, -p.asize(1)/2+1:p.asize(1)/2);
+        self.probe_support_fft = utils.imgaussfilt2_fft(sqrt(X.^2+Y.^2) < FZP_cone_diameter/2, FZP_cone_diameter/50); 
+        af_probe = sqrt(abs(fftshift(fft2(self.probe{1}(:,:,1))))); 
+        [cx, cy] = center(max(0,af_probe-0.1*max(af_probe(:)))); 
+        self.probe_support_fft = imshift_fast(self.probe_support_fft, -cx, -cy,[], 'nearest'); 
+        self.probe_support_fft = max(0, min(1, self.probe_support_fft)); 
+        verbose(1, 'Using farfield probe support constraint')
+    elseif check_option(p,'probe_support_tem')  % TEM aperture mask by ZC
+        mask_dp=abs(fft2(p.probe_initial));
+        mask_dp(mask_dp>1.0)=1;
+        mask_dp(mask_dp<0.1)=0;
+        self.probe_support_fft = logical(mask_dp);
+    else
+        self.probe_support_fft = []; 
+    end
+    
+    %%%%%%%%%%%%%%%%%%%%%%
+    %% load object%%%%%%%% 
+    %%%%%%%%%%%%%%%%%%%%%%
+    
+    if param.share_object
+        p.share_object_ID(:) = 1; % enforce single object if sharing is requested
+    end
+    
+    % correct positions for sample tilt 
+    positions = p.positions ; 
+    if isfield(p, 'positions_0')
+        positions_0 = p.positions_0 ;
+    end
+    
+    verbose(1, 'Load and pad object and probe')
+    % update current size of the object
+    for i = 1:length(p.object)
+        for j = 1:size(p.object{i},4) % load multiple layers of the object 
+            self.object{i,j} = p.object{i}(:,:,1,j);
+        end
+        p.object_size(i,:) = size(self.object{i,1});
+    end
+    
+    % calculate optimal size !! find minimal object to fit all scans 
+    Np_o = max(p.object_size,[],1); 
+    
+    % expand object size if the probe p
+    if p.number_iterations > check_option(p, 'probe_position_search') && is_method(param, {'ML', 'PIE'})
+        extra = 0.2; % add plenty of extra space for geometry refinement 
+    else
+        extra = 0.05; % do just a little of extra space 
+    end
+
+    % shift the positions to account for the expanded object size AND
+    % center them !!! (GPU code assumes positions to be centered, better for scale / probe positions are unknown)
+
+    %%%%%%%%%%%%%%%%%%%%%%
+    %% load positions %%%%
+    %%%%%%%%%%%%%%%%%%%%%%
+    
+    for i = unique(p.share_object_ID)
+        ind = [self.reconstruct_ind{p.share_object_ID == i}];
+        position_offset = 1+floor((max(positions(ind,:))-min(positions(ind,:)))/2 + min(positions(ind,:)) );
+        if isfield(p, 'positions_0')
+            self.probe_positions_0(ind,:) = positions_0(ind,:) - position_offset;
+            self.probe_positions(ind,:)   = positions(ind,:)   - position_offset;
+        else
+            self.probe_positions_0(ind,:) = positions(ind,:)   - position_offset;
+            self.probe_positions = [];
+        end
+    end
+    
+    % get object extent 
+    self.Np_o = max(Np_o,  ceil((1+extra) * ( self.Np_p +  (max(self.probe_positions_0) - min(self.probe_positions_0)) ))); 
+    % store object size without padding, useful for plotting
+    p.object_size = max(p.object_size, ceil(( self.Np_p +  (max(self.probe_positions_0) - min(self.probe_positions_0)) ))); 
+    self.probe_positions_0 = self.probe_positions_0(:,[2,1]); %swap x&y axis
+    
+    if ~isempty(self.probe_positions)
+        self.probe_positions = self.probe_positions(:,[2,1]);
+    end
+    self.Npos = Npos;
+    
+    % only a relative correction with respect to the affine matrix already
+    % applied in p-struct 
+    for ii = 1:Nscans
+        self.affine_matrix{ii} = diag([1,1]); 
+    end
+
+    
+    %%%%%%%%%%%%%%%%%%%%%%
+    %% adjust object %%%%%
+    %%%%%%%%%%%%%%%%%%%%%%
+    for i = 1:size(self.object,1)
+        for layer = 1:size(self.object,2)
+            % if needed expand the object to allow position refinement 
+            % and shift for consistency with the CPU code
+            self.object{i,layer} = imshift_fast(self.object{i,layer},1,1,self.Np_o, 'nearest', mean(self.object{i,layer}(:)));
+        end
+    end
+    
+        
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    %% load data, mask noise %%%%
+    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+    
+    verbose(1, 'Preparing data and masks')
+    assert(all(isfinite(p.fmag(:))), 'Provide p.fmag contains NaN/Inf')
+    self.noise  = [];
+    self.diffraction = (single(p.fmag .* p.fmask) / single(p.renorm) ).^2;
+    self.mask = (~p.fmask);
+    
+    if check_option(p, 'damped_mask')
+        % if relaxed mask is used, try to fill the smallest gaps (hot pixels) by neighbors 
+        mask_ind = find(self.mask); 
+        self.diffraction(mask_ind) = self.diffraction(min(mask_ind+1, numel(self.diffraction))); 
+    end
+
+    low_photon_count_dp = sum(sum(self.diffraction)) / prod(p.asize) < avg_photon_threshold;
+    if any(low_photon_count_dp)
+        error('%0.2f%% diffraction patterns has average photon count < %f', sum(low_photon_count_dp)/size(self.diffraction,3)*100, avg_photon_threshold)
+    end
+    
+    %% automatic data centering / flipping / tilted plane correction 
+    if check_option(p, 'auto_center_data') || check_option(p, 'custom_data_flip') || check_option(p, 'sample_rotation_angles')
+        self.diffraction = fftshift_2D(self.diffraction);
+        self.mask = fftshift_2D(self.mask);
+        warning on 
+        warning off backtrace
+        
+        if check_option(p, 'auto_center_data')
+            warning('Enforcing automatic data centering')
+            for ii = 1:Nscans
+                [x0,y0]=math.center(abs(fftshift(fft2(fftshift(self.probe{1}(:,:,min(end,ii)))))));
+                W = mean2(self.diffraction(:,:,self.reconstruct_ind{ii}));
+                W = ((W - min(W)) / (max(W)-min(W))).^4;  % give more weight to the more transpared regions (air)
+                avg_pattern = mean(W.*sqrt(max(0,single(self.diffraction(:,:,self.reconstruct_ind{ii})))),3);
+                [x,y]=math.center(avg_pattern);
+                x = round(x-x0); y = round(y-y0); 
+                self.diffraction(:,:,self.reconstruct_ind{ii}) = imshift_fast(self.diffraction(:,:,self.reconstruct_ind{ii}),x,y);
+                self.mask(:,:,self.reconstruct_ind{ii}) = imshift_fast(self.mask(:,:,self.reconstruct_ind{ii}), x,y);
+                fprintf('Data in scan %i shifted by %i %i pixels\n', ii, x,y); 
+            end
+            if ~isempty(self.probe_support_fft )
+                self.probe_support_fft  = imshift_fast(self.probe_support_fft , x,y);
+            end
+        end
+        
+        % apply custom flip of the diffraction data 
+        if check_option(p, 'custom_data_flip')  && any(p.custom_data_flip)
+            warning('Applying custom data flip: %i %i %i ', p.custom_data_flip(1), p.custom_data_flip(2), p.custom_data_flip(3))
+            if p.custom_data_flip(1)
+                self.diffraction =  flipud(self.diffraction);
+                self.mask        =  flipud(self.mask); 
+            end
+            if p.custom_data_flip(2)
+                self.diffraction =  fliplr(self.diffraction); 
+                self.mask        =  fliplr(self.mask); 
+            end
+            if p.custom_data_flip(3)
+                self.diffraction =  permute(self.diffraction, [2,1,3]); 
+                self.mask        =  permute(self.mask, [2,1,3]); 
+            end
+        end
+
+        %     
+         if isfield(p, 'sample_rotation_angles') &&  any(p.sample_rotation_angles) && check_option(p, 'apply_tilted_plane_correction', 'diffraction') 
+            %% OFFAXIS PTYCHOGRAPHY CORRECTION 
+            if utils.verbose > -1
+                warning('Applying tilted plane correction: %3.5g %3.3g %3.3g\n Note that current implementation assumes low NA illumination, if this is not true, the central diffraction cone can be malformed', p.sample_rotation_angles(1), p.sample_rotation_angles(2), p.sample_rotation_angles(3))
+            end
+            % create matrix of deformation to apply effects similar to 
+            deform_mat =  get_tilted_plane_correction_matrix(max(self.Np_p), ...
+                                                p.z ,p.detectors{1}.pixel_size, ...
+                                                p.sample_rotation_angles(1),...
+                                                p.sample_rotation_angles(2),...
+                                                p.sample_rotation_angles(3));
+            if self.Np_p(1) ~= self.Np_p(2)
+                % quick fix for asymmetric probe dimensions
+                blank = true(self.Np_p); 
+                blank_ind = find(utils.crop_pad(blank,[ max(self.Np_p), max(self.Np_p)])); 
+                deform_mat = deform_mat(blank_ind,blank_ind); 
+            end
+                
+            apply_deform = @(x,D)single(reshape(full(D * double(reshape(x,[],size(x,3)))), size(x)));
+            plotting.smart_figure(3423)
+            ax(1)=subplot(1,2,1);
+            imagesc(log(1+max(self.diffraction, [],3))); axis off image xy ; colormap(plotting.franzmap)
+            title('Diffraction BEFORE tilted plane correction')
+            % apply deformation effects caused by tilted sample , be sure to enforce the mask before
+            % interpolation, the hotpixels can spread around after the correction
+            self.diffraction = apply_deform(self.diffraction .* ~self.mask, deform_mat');
+            self.mask =        apply_deform(self.mask, deform_mat') > 0;
+            % plotting.imagesc3D(log(1+max(self.diffraction,[],3))); grid on  
+            ax(2)=subplot(1,2,2);
+            imagesc(log(1+max(self.diffraction, [],3))); axis off image xy ; colormap(plotting.franzmap)
+            title('Diffraction AFTER tilted plane correction')
+            plotting.suptitle('Effect of tilted ptychography correction')
+            linkaxes(ax, 'xy')
+            drawnow
+            self.diffraction_deform_matrix = deform_mat;
+        end
+
+        warning on 
+
+        self.diffraction = ifftshift_2D(self.diffraction);
+        self.mask = ifftshift_2D(self.mask);
+    end
+    
+    self.filename = [p.detector.data_prefix, p.run_name]; 
+    
+    % other basic parameters 
+    self.path = '';
+    if check_option(p, 'prop_regime', 'nearfield')
+        self.z_distance = p.z; 
+    else
+        self.z_distance = inf;  % farfield  
+    end
+    % multilayer extension 
+    if isfield(p, 'delta_z')
+        assert(all(isfinite(p.delta_z)), 'Some of the provided layer distanced delta_z is not finite' )
+        self.z_distance = [p.delta_z(:)', self.z_distance]; 
+    end
+    
+    self.lambda = p.lambda;
+    if check_option(p,'diff_pattern_blur')
+        self.diff_pattern_blur=p.diff_pattern_blur;
+    else
+        self.diff_pattern_blur = 0;  % incoherent smoothing 
+    end
+    self.modes = []; 
+           
+    % keep p structure for plotting purposes 
+    param.p = p;
+
+    % if requested, shift the average probe to center and shift the object
+    % to correspond to the probe shift
+    if check_option(p, 'auto_center_probe')
+        [x,y] = center(mean(abs(self.probe{1}(:,:,:,1)))); 
+        for ii = 1:numel(self.probe)
+            self.probe{ii} = imshift_fft(self.probe{ii}, -x,-y);
+        end
+        for ii = 1:numel(self.object)
+            self.object{ii} = imshift_fft(self.object{ii}, -x, -y);
+        end
+    end
+    
+    % check if all is ok (remove in future !!!)
+    positions = self.probe_positions_0(:,[2,1]); 
+    positions = bsxfun(@plus, positions, ceil(self.Np_o/2-self.Np_p/2));    
+    positions = round(positions); 
+    range = ([min(positions(:,1)), max(positions(:,1))+ Np_p(1), min(positions(:,2)), max(positions(:,2))+ Np_p(2)]);
+    if range(1) < 0 || range(2) > self.Np_o(1) || range(3) < 0 || range(4) > self.Np_o(2)
+        warning('Object size is too small, not enough space for probes !! ') 
+    end
+
+    
+end

ptycho/+engines/+GPU_MS/+initialize/prepare_flyscan_positions.m

@@ -0,0 +1,113 @@
+% PREPARE_FLYSCAN_POSITIONS from finit number of measured position interpolate possitions for each
+% measured frame when fly scan is used
+% 
+% self = prepare_flyscan_positions(self, par)
+% 
+% 
+% ** self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ** par       structure containing parameters for the engines 
+%
+% returns: 
+% ++ self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+
+% 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 mixed coherent modes:
+% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
+% for LSQ-ML method 
+% M. Odstrcil, A. Menzel, M.G. Sicairos,  Iterative least-squares solver for generalized maximum-likelihood ptychography, Optics Express, 2018
+% for OPRP method 
+%  M. Odstrcil, P. Baksh, S. A. Boden, R. Card, J. E. Chad, J. G. Frey, W. S. Brocklesby,  "Ptychographic coherent diffractive imaging with orthogonal probe relaxation." Optics express 24.8 (2016): 8360-8369
+% 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). 
+% 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 = prepare_flyscan_positions(self, par)
+import engines.GPU_MS.GPU_wrapper.*
+import math.*
+import utils.*
+import plotting.*
+
+
+    jumps = diff(self.probe_positions_0);
+    step = median(jumps,1);
+    jumps = sum(abs(jumps),2);
+    % empirical condition 
+    jumps = find(jumps > 10*median(jumps));
+    if length(jumps) < par.Nscans
+        % assume that smooth path is used 
+       %% ADVANCED FLY SCAN - SPIRAL
+        for ii = 1:par.Nscans
+            assert(~any(isfinite(par.probe_position_search)), 'Position refinement and fly scans not suported')
+
+            ind =  self.reconstruct_ind{ii};
+            pos = self.probe_positions_0(ind,:);
+            [ang, rad] = cart2pol(pos(:,1)-pos(1,1), pos(:,2)-pos(1,2)); 
+            ang = unwrap(ang); 
+            % get interpolate d positions of the sub probes 
+            ang_all = ang + (par.flyscan_offset -0.5+linspace(0,par.flyscan_dutycycle*(par.Nmodes-1)/par.Nmodes, par.Nmodes)   ).*[diff(ang);0]; 
+            rad_all = interp1(ang, rad, ang_all, 'pchip'); 
+            [X,Y] = pol2cart(ang_all, rad_all);
+
+            for ll = 1:par.Nmodes            
+                    self.modes{ll}.probe_positions(ind,:) = pos(1,1:2) + [X(:,ll), Y(:,ll)];
+                if iter == 1;  self.probe{ll} = self.probe{1}; end
+            end
+        end
+    else
+        %% ADVANCED FLY SCAN - LINE SCAN 
+        % interpolate the other modes into new positions 
+        pos = self.modes{1}.probe_positions;
+        for ll = 1:par.Nmodes            
+            ratio = par.flyscan_dutycycle*(ll-1)/par.Nmodes;
+            self.modes{ll}.probe_positions = pos(min((1:self.Npos)+1,self.Npos),:)*ratio + (1-ratio)*pos;
+            if ~isempty(jumps)
+                % expected step continuation
+                self.modes{ll}.probe_positions(jumps,:) = bsxfun(@plus, self.modes{ll}.probe_positions(jumps-1,:),step);
+            end         
+        end
+    end
+end

ptycho/+engines/+GPU_MS/+initialize/private/distmat.m

@@ -0,0 +1,157 @@
+% DISTMAT Compute a Distance Matrix for One or Two Sets of Points
+%
+% 
+% 
+% Copyright (c) 2015, Joseph Kirk
+% All rights reserved.
+% 
+% Redistribution and use in source and binary forms, with or without
+% modification, are permitted provided that the following conditions are
+% met:
+% 
+% * Redistributions of source code must retain the above copyright
+% notice, this list of conditions and the following disclaimer.
+% * Redistributions in binary form must reproduce the above copyright
+% notice, this list of conditions and the following disclaimer in
+% the documentation and/or other materials provided with the distribution
+% 
+% THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+% AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+% IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+% ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+% LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+% CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+% SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+% INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+% CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+% ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+% POSSIBILITY OF SUCH DAMAGE.
+% 
+
+% Filename: distmat.m
+%
+% Description: Computes a matrix of pair-wise distances between points in
+%              A and B, using one of {euclidean,cityblock,chessboard} methods
+%
+% Author:
+%       Joseph Kirk
+%       jdkirk630@gmail.com
+%
+% Date: 02/27/15
+%
+% Release: 2.0
+%
+% Inputs:
+%     A      - (required) MxD matrix where M is the number of points in D dimensions
+%     B      - (optional) NxD matrix where N is the number of points in D dimensions
+%                 if not provided, B is set to A by default
+%     METHOD - (optional) string specifying one of the following distance methods:
+%                 'euclidean'                       Euclidean distance (default)
+%                 'taxicab','manhattan','cityblock' Manhattan distance
+%                 'chebyshev','chessboard','chess'  Chebyshev distance
+%                 'grid','diag'                     Diagonal grid distance
+%
+% Outputs:
+%     DMAT   - MxN matrix of pair-wise distances between points in A and B
+%
+% Usage:
+%     dmat = distmat(a)
+%       -or-
+%     dmat = distmat(a,b)
+%       -or-
+%     dmat = distmat(a,method)
+%       -or-
+%     dmat = distmat(a,b,method)
+%
+% Example:
+%     % Pairwise Euclidean distances within a single set of 2D points
+%     xy = 10*rand(25,2);  % 25 points in 2D
+%     dmat = distmat(xy);
+%     figure; plot(xy(:,1),xy(:,2),'.');
+%     for i=1:25, text(xy(i,1),xy(i,2),[' ' num2str(i)]); end
+%     figure; imagesc(dmat); colorbar
+%
+% Example:
+%     % Pairwise Manhattan distances within a single set of 2D points
+%     xy = 10*rand(25,2);  % 25 points in 2D
+%     dmat = distmat(xy,'cityblock');
+%     figure; plot(xy(:,1),xy(:,2),'.');
+%     for i=1:25, text(xy(i,1),xy(i,2),[' ' num2str(i)]); end
+%     figure; imagesc(dmat); colorbar
+%
+% Example:
+%     % Pairwise Chebyshev distances within a single set of 2D points
+%     xy = 10*rand(25,2);  % 25 points in 2D
+%     dmat = distmat(xy,'chebyshev');
+%     figure; plot(xy(:,1),xy(:,2),'.');
+%     for i=1:25, text(xy(i,1),xy(i,2),[' ' num2str(i)]); end
+%     figure; imagesc(dmat); colorbar
+%
+% Example:
+%     % Inter-point Euclidean distances for 2D points
+%     xy = 10*rand(15,2);  % 15 points in 2D
+%     uv = 10*rand(25,2);  % 25 points in 2D
+%     dmat = distmat(xy,uv);
+%     figure; plot(xy(:,1),xy(:,2),'.');
+%     for i=1:15, text(xy(i,1),xy(i,2),[' ' num2str(i)]); end
+%     figure; plot(uv(:,1),uv(:,2),'.');
+%     for i=1:25, text(uv(i,1),uv(i,2),[' ' num2str(i)]); end
+%     figure; imagesc(dmat); colorbar
+%
+% See also:
+%
+function dmat = distmat(a,varargin)
+    
+    
+    % Set defaults
+    method = 'euclidean';
+    b = a;
+    
+    % Error check primary input
+    if ~isnumeric(a)
+        error('Expecting a matrix of floating point values for A input.');
+    end
+    
+    % Process optional inputs
+    for var = varargin
+        arg = var{1};
+        if ischar(arg)
+            method = arg;
+        elseif ~isempty(arg)
+            b = arg;
+        end
+    end
+    
+    % Check input dimensionality
+    [na,aDims] = size(a);
+    [nb,bDims] = size(b);
+    if (aDims ~= bDims)
+        error('Input matrices must have the same dimensionality.');
+    end
+    
+    % Create index matrices
+    [j,i] = meshgrid(1:nb,1:na);
+    
+    % Compute array of inter-point differences
+    delta = a(i,:) - b(j,:);
+    
+    % Compute distance by specified method
+    dmat = zeros(na,nb);
+    switch lower(method)
+        case {'euclidean','euclid'}
+            % Euclidean distance
+            dmat(:) = sqrt(sum(delta.^2,2));
+        case {'cityblock','city','block','manhattan','taxicab','taxi'}
+            % Cityblock distance
+            dmat(:) = sum(abs(delta),2);
+        case {'chebyshev','cheby','chessboard','chess'}
+            % Chebyshev distance
+            dmat(:) = max(abs(delta),[],2);
+        case {'grid','diag'}
+            dmat(:) = max(abs(delta),[],2) + (sqrt(2) - 1)*min(abs(delta),[],2);
+        otherwise
+            error('Unrecognized distance method %s',method);
+    end
+    
+end
+

ptycho/+engines/+GPU_MS/+initialize/private/get_close_indices.m

@@ -0,0 +1,300 @@
+% GET_CLOSE_INDICES simple based method to select indices for DM 
+% !! GPU needs the sets to be with similar , ideally the same sizes !!! 
+%
+% [indices_out, scan_ids_out] = get_close_indices(self, cache, par )          
+% 
+% 
+% ** 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
+%
+% returns: 
+% ++ indices_out   cell of arrays, contain indices of positions processed in parallel 
+% ++ scan_ids_out  cell of arrays, contain scan numbers for each position 
+
+% 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 mixed coherent modes:
+% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
+% for LSQ-ML method 
+% M. Odstrcil, A. Menzel, M.G. Sicairos,  Iterative least-squares solver for generalized maximum-likelihood ptychography, Optics Express, 2018
+% for OPRP method 
+%  M. Odstrcil, P. Baksh, S. A. Boden, R. Card, J. E. Chad, J. G. Frey, W. S. Brocklesby,  "Ptychographic coherent diffractive imaging with orthogonal probe relaxation." Optics express 24.8 (2016): 8360-8369
+% 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). 
+% 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 [indices_out, scan_ids_out] = get_close_indices(self, cache, par )      
+
+import math.*
+import utils.*
+
+grouping = par.grouping;
+
+% in case of a shared scan join together all positions to find the optimal groups 
+group_across_scans = true; 
+
+
+if par.share_object && group_across_scans
+    Nsets = 1; 
+else
+    Nsets = par.Nscans; 
+end
+cluster_refinement_time = 0;
+cluster_time = 0; 
+
+% in simplest case process all positions together 
+if Nsets == 1 && grouping >= self.Npos
+    indices_out = {[self.reconstruct_ind{:}]}; 
+    scan_ids_out{1} = []; 
+    for ii = 1:length(self.reconstruct_ind)
+        scan_ids_out{1} = [scan_ids_out{1}; ii*ones(length(self.reconstruct_ind{ii}),1)];
+    end
+    return
+end
+
+%rng default
+
+for kk = 1:Nsets
+    % take them sequentially but with random offset 
+    if par.share_object  && group_across_scans
+        % join all indices into one large set if the object is shared 
+        indices_0 = [self.reconstruct_ind{:}];
+        for ii = 1:length(self.reconstruct_ind)
+            scans_0(self.reconstruct_ind{ii}) = ii;
+        end
+    else
+        indices_0 = self.reconstruct_ind{kk};
+        scans_0 = kk * ones(size(indices_0)); % scan number 
+    end
+    N = length(indices_0);
+    Ngroups=ceil(N/grouping);
+    positions = self.probe_positions_0(indices_0,:);
+    Npos = length(positions);
+
+    % get initial set  distribution 
+    [groups, C, sum_D, D] = get_best_kmeans(positions, Ngroups); 
+
+    iter = 0;
+    t0 = tic;
+    
+    while true 
+        iter= iter +1;
+        [nbins,bins] = hist(groups, unique(groups));
+        % if less than 2 types of groups are present, finish 
+        Ngroups_sizes = length(unique(nbins));
+        % try to find distribution with most similar sets sizes, if not
+        % easy, end with suboptimal distribution  after 50 iterations 
+        if ( Ngroups_sizes <= max(2, ceil(iter/1e3)) && (Ngroups*grouping ~= N || iter > 1e3 )) ...
+                ||  Ngroups_sizes == 1 % choose suboptimal solution if better is not found soon
+            break
+        end
+
+        % find group with lowest number of members , add new points into
+        % this group 
+        min_group = bins(argmin(nbins));
+        large_groups = bins(nbins>grouping);
+        if isempty(large_groups) || any(ismember(min_group, large_groups)) ; break; end 
+        % choose closest position from the largest group to be moved to the
+        % smallest group 
+        ind_large = (D(:,min_group) == min(D(ismember(groups, large_groups), min_group)));
+
+        groups(ind_large) = min_group;
+       
+    end
+
+    % remove empty groups 
+    ugroups = unique(groups); 
+    Ngroups = length(ugroups); 
+    groups = sum((1:Ngroups) .*(groups == ugroups'),2);
+    
+    
+    
+    cluster_time = cluster_time + toc(t0); 
+
+    
+    for ii = 1:Ngroups
+        C(ii,:) = median(positions(groups == ii,:));
+    end
+    for ii = 1:Ngroups
+        D(:,ii) = (sum((positions - C(ii,:)).^2,2));
+    end
+
+        
+    t0 = tic;
+    %% find more compact refinement 
+    % find the most distanced points 
+    [~,sind] = sort(D,2);
+    % positions to be improved  -> find the best matching group 
+    optimal_group = sind(:,1); 
+
+    nonoptimal_ratio_0 = 1; 
+
+    for iter = 1:10
+        ind_switch  = (groups ~=  optimal_group);
+        nonoptimal_ratio = sum(ind_switch) / numel(ind_switch);
+        if nonoptimal_ratio > 0
+            verbose(0, 'Indexes to be switched:  %3.2g%% positions', nonoptimal_ratio * 100)
+        end
+        
+        if nonoptimal_ratio >= nonoptimal_ratio_0
+            break
+        end
+        nonoptimal_ratio_0 = nonoptimal_ratio; 
+
+        max_dist_0 = inf;
+        for i = 1:sum(ind_switch)
+            % calculate distance for each point from its group center 
+            center_dist = (D(sub2ind(size(D), (1:Npos)', groups)));
+            max_dist_0 = max(center_dist(ind_switch));
+            % start from the worst case
+            ind_worse = find(max(center_dist(ind_switch)) == center_dist, 1, 'first');
+            % initial group 
+            group_old = groups(ind_worse);
+            % better fitting group 
+            group_new = optimal_group(ind_worse);
+            % position to be switched in the new group 
+            ind_new = find(D(:,group_old) == min(D(groups == group_new, group_old)), 1, 'first');
+            % switch the group members  
+            groups(ind_worse) = group_new;
+            groups(ind_new) = group_old;
+            ind_switch([ind_worse, ind_new]) = 0;
+            if all(ind_switch == 0)
+                break
+            end
+        end
+      
+%          ind_switch = (groups ~=  sind(:,1));
+%          for ii = Ngroups
+%              clf
+%              hold all;
+%              ind = groups == ii;
+%              ax = plot(self.probe_positions_0(ind & ind_switch, 1), self.probe_positions_0(ind & ind_switch, 2), 'o');
+%              ax2 = plot(self.probe_positions_0(ind & ~ind_switch, 1), self.probe_positions_0(ind & ~ind_switch, 2), 'x');
+%              try; ax2.Color = ax.Color; end
+%              plot(C(ii,1),C(ii,2),'x','Linewidth', 2)
+% %              drawnow 
+% %             pause(1)
+%          end
+%          title(num2str(iter))
+%          axis tight equal
+%          pause(1)
+%        
+    end
+    
+    cluster_refinement_time = cluster_refinement_time + toc(t0); 
+
+
+    %% optimally sort the indices to help GPU
+    [nbins,bins] = hist(groups, unique(groups));
+    [~,ind] = sort(nbins,2,'descend');
+    for ii = 1:length(bins)
+        indices{kk}{ii} = indices_0((groups == bins(ind(ii))));
+        scan_ids{kk}{ii} = scans_0((groups == bins(ind(ii))));
+    end 
+    verbose(2,'=== Number of cluster sizes %i', length(unique(nbins)))
+end
+
+verbose(0,'=== Position clusters found in %i iterations in %3.2gs', iter, cluster_time)
+verbose(0,'=== Position clusters refined in %i iterations in %3.2gs', iter, cluster_refinement_time)
+
+
+%rng shuffle 
+
+if verbose()  > 1 && Ngroups_sizes > 1
+    warning('Unequal group sizes, it may cause slower calculation')
+end
+
+
+indices_out = horzcat(indices{:});
+scan_ids_out =  horzcat(scan_ids{:});
+
+
+if Ngroups  == 1 && Nsets == 1
+    %% merge groups from multiple scans into larger chunks if grouping is too large 
+    indices_out = {horzcat(indices_out{:})};
+    scan_ids_out =  {horzcat(scan_ids_out{:})};
+end
+
+
+
+%% sort them to minimize allocation of new projection matrices
+Nitems = cellfun(@length, indices_out); 
+
+if all(max(Nitems) - min(Nitems) <= 1) && all(Nitems > 100)
+    % just neglect one scanning position to keep the bunches with the same
+    % size -> faster run on GPU 
+    for i = 1:length(indices_out)
+        indices_out{i} = indices_out{i}(1:min(Nitems));
+        scan_ids_out{i} = scan_ids_out{i}(1:min(Nitems));
+    end
+else
+    [~,ind] = sort(Nitems(:),1,'descend' );
+    indices_out = indices_out(ind);
+    scan_ids_out = scan_ids_out(ind);
+end
+
+
+    
+end
+
+
+function [groups, C, sum_D, D] = get_best_kmeans(positions, Ngroups)
+    % make several guesses to get better Kmean distribution 
+    warning('off','stats:kmeans:FailedToConverge')
+    for i = 1:10
+        [groups{i}, C{i}, sum_D{i}, D{i}] = kmeans(positions, Ngroups);
+        nbins = hist(groups{i}, unique(groups{i}));
+        score(i) = std(nbins);
+    end
+    best = math.argmin(score); 
+    groups = groups{best};
+    C = C{best};
+    sum_D = sum_D{best};
+    D = D{best};
+end
+

ptycho/+engines/+GPU_MS/+initialize/private/get_nonoverlapping_indices.m

@@ -0,0 +1,203 @@
+%  GET_NONOVERLAPPING_INDICES a heuristic based method to select pseudorandom indices of non overlapping regions   
+% Note: It can be slow for large number of scanning positions 
+%
+% [indices_out, scan_ids_out] = get_nonoverlapping_indices(self, cache, par )      
+%
+% ** 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
+%
+% returns: 
+% ++ indices_out   cell of arrays, contain indices of positions processed in parallel 
+% ++ scan_ids_out  cell of arrays, contain scan numbers for each position 
+
+
+% 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 mixed coherent modes:
+% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
+% for LSQ-ML method 
+% M. Odstrcil, A. Menzel, M.G. Sicairos,  Iterative least-squares solver for generalized maximum-likelihood ptychography, Optics Express, 2018
+% for OPRP method 
+%  M. Odstrcil, P. Baksh, S. A. Boden, R. Card, J. E. Chad, J. G. Frey, W. S. Brocklesby,  "Ptychographic coherent diffractive imaging with orthogonal probe relaxation." Optics express 24.8 (2016): 8360-8369
+% 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). 
+% 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 [indices_out, scan_ids_out] = get_nonoverlapping_indices(self, cache, par )      
+
+% find groups accross the scans in order to further minimize overlap
+group_across_scans = true; %need to be true for sharing object amongs scans
+
+if group_across_scans
+% divide the grouping equally over all the scans 
+    grouping = ceil(par.grouping/par.Nscans);
+else
+    grouping = par.grouping;
+end
+max_groups = 0;
+
+for kk = 1:par.Nscans
+        
+    % setdiff sort the indices by size 
+    indices_0 = self.reconstruct_ind{kk};  % remove unwanted from the decision process 
+    ind_start(kk) = min(indices_0)-1;
+    indices_0 = indices_0 - ind_start(kk);  % 
+    Npos_tmp=length(indices_0);
+    % randomly permutate the indices    
+    indices_0 = indices_0(randperm(Npos_tmp)); 
+    max_groups = max(max_groups, ceil(Npos_tmp/grouping));
+    % fill it with some initial random guess 
+    for ii = 1:ceil(Npos_tmp/grouping) 
+         indices{kk}{ii} = indices_0(1+(ii-1)*grouping : min(Npos_tmp,ii*grouping));
+    end
+   
+    % no need for this method ot it calculation would be too long -> use
+    % just the random initial guess 
+    if (self.Npos/par.Nscans > 1e3 ) || (grouping == 1) || ~isfield(cache, 'distances_matrix')
+        %%%for ii = 1:length(indices{1}) %why length(indices{1})? Bug?
+        for ii = 1:length(indices{kk}) %modified by YJ to prevent error when different scans have differernt number of positions
+            scan_ids{kk}{ii} = ones(1,length(indices{kk}{ii}))*kk;  % note their scan origin
+        end
+        continue 
+    end  % hope that for large number of positions the random statistics will be enough
+    
+    try
+
+     update_score = 0;
+     for i = 1:ceil(Npos_tmp/grouping)-1
+         id = indices{kk}{i};
+         dist_mat_small = cache.distances_matrix{kk}(id,id);
+         for ii = 0:2*length(indices{kk}{i+1})  % go twice through all positions 
+             j = 1+mod(ii, length(indices{kk}{i+1}));
+             min_dist = 1./sum(1./dist_mat_small.^2); % find the shortest distance between the probes
+             if all(isinf(min_dist)) % all(isnan(min_dist))
+                 break
+             end
+             [~,min_dist_ind] = min(min_dist);
+            % make a swap with the j position in i+1 index array
+             tmp = indices{kk}{i+1}(j);
+             indices{kk}{i+1}(j) = indices{kk}{i}(min_dist_ind);
+             indices{kk}{i}(min_dist_ind) = tmp;
+             
+             
+            % update distance matrix 
+            dist_mat_small_update = cache.distances_matrix{kk}(tmp,indices{kk}{i});
+            dist_mat_small(min_dist_ind,:) = dist_mat_small_update';
+            dist_mat_small(:,min_dist_ind) = dist_mat_small_update;
+         end
+         update_score = update_score +j;
+     end
+    catch
+        keyboard
+    end
+        
+
+     % fill the last group by the skip indieces but do not expand it 
+     skip_ind = cache.skip_ind(randperm(length(cache.skip_ind)));
+     indices{kk}{end} = [indices{kk}{end}, skip_ind(1:min(end, grouping-length(indices{kk}{end})))]; % join skip_ind back to the last (smallest) set    
+          
+     for ii = 1:length(indices{kk})
+        scan_ids{kk}{ii} = ones(1,length(indices{kk}{ii}))*kk;  % note their scan origin
+     end
+end
+
+
+if group_across_scans
+    indices_out = cell(max_groups,1);
+    scan_ids_out = cell(max_groups,1);
+    %% merge groups from difference scans into larger chunks if required 
+    for ii = 1:max_groups
+        indices_out{ii} = [];
+        scan_ids_out{ii} = [];
+        % from each scan add one group
+        for kk = 1:par.Nscans
+            if ii <= length(indices{kk})
+                indices_out{ii} = [indices_out{ii}, indices{kk}{ii}+ind_start(kk)];
+                scan_ids_out{ii} = [scan_ids_out{ii}, scan_ids{kk}{ii}];
+            end
+        end
+        if length(scan_ids_out) > 1 && length(scan_ids_out{end}) < grouping / 10
+            % if the a group is too small, merge it with the previous to
+            % reduce the overhead 
+            indices_out{end-1} = [indices_out{end-1}, indices_out{end}]; 
+            scan_ids_out{end-1} = [scan_ids_out{end-1}, scan_ids_out{end}];
+            scan_ids_out(end) = []; indices_out(end) = []; 
+        end
+    end
+else
+    indices_out = {};
+    for ii = 1:par.Nscans
+        for kk = 1:length(indices{ii})
+            indices_out = [indices_out, indices{ii}{kk}+ind_start(ii)];
+        end
+    end
+    scan_ids_out = [scan_ids{:}]';
+
+end
+
+%% sort them to minimize allocation of new projection matrices
+Nitems = cellfun(@length, indices_out); 
+
+if all(max(Nitems) - min(Nitems) <= 1) && all(Nitems > 100)
+    % just neglect one scanning position to keep the bunches with the same
+    % size -> faster run on GPU 
+    for i = 1:length(indices_out)
+        indices_out{i} = indices_out{i}(1:min(Nitems));
+        scan_ids_out{i} = scan_ids_out{i}(1:min(Nitems));
+    end
+else
+    [~,ind] = sort(Nitems(:),1,'descend' );
+    indices_out = indices_out(ind);
+    scan_ids_out = scan_ids_out(ind);
+end
+
+
+    
+end
+
+
+

ptycho/+engines/+GPU_MS/+initialize/private/get_scanning_indices.m

@@ -0,0 +1,147 @@
+% GET_SCANNING_INDICES simple based method to select indices for DM 
+% 
+% [indices_out, scan_ids_out] = get_scanning_indices(self, cache, par ) 
+%
+% ** 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
+%
+% returns: 
+% ++ indices_out   cell of arrays, contain indices of positions processed in parallel 
+% ++ scan_ids_out  cell of arrays, contain scan numbers for each position 
+
+
+% 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 mixed coherent modes:
+% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
+% for LSQ-ML method 
+% M. Odstrcil, A. Menzel, M.G. Sicairos,  Iterative least-squares solver for generalized maximum-likelihood ptychography, Optics Express, 2018
+% for OPRP method 
+%  M. Odstrcil, P. Baksh, S. A. Boden, R. Card, J. E. Chad, J. G. Frey, W. S. Brocklesby,  "Ptychographic coherent diffractive imaging with orthogonal probe relaxation." Optics express 24.8 (2016): 8360-8369
+% 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). 
+% 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 [indices_out, scan_ids_out] = get_scanning_indices(self, cache, par )      
+
+import engines.GPU_MS.GPU_wrapper.*
+import engines.GPU_MS.shared.*
+
+grouping = par.grouping;
+max_groups = 0;
+
+if self.Npos == grouping && par.Nscans == 1
+    indices_out = self.reconstruct_ind;
+    scan_ids_out = {ones(self.Npos,1)};
+    return
+end
+
+
+for kk = 1:par.Nscans
+     N = length(self.reconstruct_ind{kk});
+
+     % !! indices ordering has to be always the same for DM !!! 
+      
+     indices_0 = self.reconstruct_ind{kk};
+
+     max_groups = max(max_groups, ceil(N/grouping));
+
+     for ii = 1:ceil(N/grouping)
+         indices{kk}{ii} = indices_0(1+(ii-1)*grouping : min(end,ii*grouping));
+     end     
+     % fill the last group by the skip indices but do not expand it 
+     skip_ind = cache.skip_ind(randperm(length(cache.skip_ind)));
+     indices{kk}{end} = [indices{kk}{end}, skip_ind(1:min(end, grouping-length(indices{kk}{end})))]; % join skip_ind back to the last (smallest) set    
+     for ii = 1:length(indices{kk})
+        scan_ids{kk}{ii} = kk * ones(1,length(indices{kk}{ii}));  % note their scan origin
+     end 
+end
+
+% how many scans should be merged to reach the desired grouping 
+Njoin = ceil(par.grouping / (self.Npos/par.Nscans));
+
+
+if Njoin > 1 && par.Nscans > 1 && is_method(par, {'PIE', 'ML'})
+    % join several scan to improve performance 
+    indices_out = cell(ceil(par.Nscans/Njoin),1);
+    scan_ids_out = cell(ceil(par.Nscans/Njoin),1);
+    %% merge groups from difference scans into larger chunks 
+    for kk = 1:ceil(par.Nscans/Njoin)
+        indices_out{kk} = [];
+        scan_ids_out{kk} = [];
+        for ii = 1:Njoin
+            if kk+(ii-1)*ceil(par.Nscans/Njoin) <= par.Nscans
+            indices_out{kk} = [indices_out{kk}, indices{kk+(ii-1)*ceil(par.Nscans/Njoin)}{1}];
+            scan_ids_out{kk} = [scan_ids_out{kk}, scan_ids{kk+(ii-1)*ceil(par.Nscans/Njoin)}{1}];
+            end
+        end
+    end
+else
+    indices_out = horzcat(indices{:});
+    scan_ids_out =  horzcat(scan_ids{:});
+end
+
+
+%% sort them to minimize allocation of new projection matrices
+Nitems = cellfun(@length, indices_out); 
+
+if all(max(Nitems) - min(Nitems) <= 1) && all(Nitems > 100)
+    % just neglect one scanning position to keep the bunches with the same
+    % size -> faster run on GPU 
+    for i = 1:length(indices_out)
+        indices_out{i} = indices_out{i}(1:min(Nitems));
+        scan_ids_out{i} = scan_ids_out{i}(1:min(Nitems));
+    end
+else
+    [~,ind] = sort(Nitems(:),1,'descend' );
+    indices_out = indices_out(ind);
+    scan_ids_out = scan_ids_out(ind);
+end
+
+    
+end

ptycho/+engines/+GPU_MS/+initialize/rescale_inputs.m

@@ -0,0 +1,167 @@
+% RESCALE_INPUTS multigrid scheme method
+%
+% self = rescale_inputs(self, Np_p_new, rescale_data)
+%
+% 
+% ** self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ** Np_p_new  new size of the rescaled dataset 
+% ** rescale_data  true -> rescale also data + mask + noise arrays 
+%
+% returns: 
+% ++ self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+
+    
+% 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 mixed coherent modes:
+% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
+% for LSQ-ML method 
+% M. Odstrcil, A. Menzel, M.G. Sicairos,  Iterative least-squares solver for generalized maximum-likelihood ptychography, Optics Express, 2018
+% for OPRP method 
+%  M. Odstrcil, P. Baksh, S. A. Boden, R. Card, J. E. Chad, J. G. Frey, W. S. Brocklesby,  "Ptychographic coherent diffractive imaging with orthogonal probe relaxation." Optics express 24.8 (2016): 8360-8369
+% 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). 
+% 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 = rescale_inputs(self, Np_p_new, rescale_data)
+
+    import utils.*
+    import math.*
+
+    if isempty(Np_p_new)
+        return 
+    end
+    
+    scale = Np_p_new ./ self.Np_p; 
+    
+    if all(scale == 1)
+        return
+    end
+    
+    self.modes = [];
+
+    if ~isempty(self.diffraction_deform_matrix)
+        verbose(1,'Rescalling with diffraction_deform_matrix is not implemented')
+        self.diffraction_deform_matrix = [];
+    end
+    
+    self.Np_p = ceil(self.Np_p.*scale);
+    self.Np_o = ceil(self.Np_o.*scale);
+
+    if ~isempty(self.probe_positions) && any(self.probe_positions(:) ~= self.probe_positions_0(:))
+        self.probe_positions = scale([2,1]) .* self.probe_positions;
+    else
+        self.probe_positions = []; 
+    end
+    self.probe_positions_0 = scale([2,1]) .* self.probe_positions_0;
+    self.pixel_size = self.pixel_size ./ scale; 
+
+    for i = 1:length(self.probe)
+        % scale also intensity 
+        self.probe{i} = interpolateFT_addnoise(self.probe{i}, self.Np_p, 1 )./prod(scale);
+    end
+
+    for i = 1:numel(self.object)
+        self.object{i} = interpolateFT_addnoise(self.object{i}, self.Np_o , 1);
+    end
+    
+    if ~isempty(self.probe_support)  
+        self.probe_support = interpolateFT( self.probe_support, self.Np_p) ;    
+    end
+    if ~isempty(self.probe_support_fft)  
+        self.probe_support_fft = crop_pad( self.probe_support_fft, self.Np_p) ;    
+    end
+
+
+    if rescale_data
+        self.diffraction = fftshift_2D(self.diffraction);
+        self.diffraction = crop_pad(self.diffraction, self.Np_p);
+        self.diffraction = ifftshift_2D(self.diffraction);
+        
+        % force mask empty by Zhen Chen
+        self.mask=[];
+        if ~isempty(self.mask)
+            fill_value = 0.9;  % fill value in case of ptychographic "super resolution"
+            self.mask = fftshift_2D(self.mask);
+            self.mask = crop_pad(self.mask, self.Np_p, fill_value);
+            self.mask = ifftshift_2D(self.mask);
+        end
+        if ~isempty(self.noise)
+            self.noise = fftshift_2D(self.noise);
+            self.noise = crop_pad(self.noise, self.Np_p);
+            self.noise = ifftshift_2D(self.noise);
+
+        end
+    end
+    
+        
+
+end
+
+function [ imout ] = interpolateFT_addnoise(im,outsize, noise_amplitude)
+    % fill the empty regions in the FFT interpolated data by some weak
+    % random noise 
+    import math.fftshift_2D
+    import math.ifftshift_2D
+    import utils.crop_pad
+
+
+    Nout = outsize;
+    Nin = size(im);
+
+    imFT = fftshift_2D(fft2(im));
+
+    imout = crop_pad(imFT, outsize);
+
+    % add noise to avoid correlation between upsampled imaged from
+    % interpolation artefacts 
+    imout = imout + noise_amplitude * randn(outsize).*mean(min(abs(imFT))); 
+
+    imout = ifft2(ifftshift_2D(imout))*(Nout(1)*Nout(2)/(Nin(1)*Nin(2)));
+
+
+end
+

ptycho/+engines/+GPU_MS/+initialize/save_to_p.m

@@ -0,0 +1,177 @@
+% SAVE_TO_P save parameters and recosntrutions from  param and self structures to the p-structure
+%
+% p_out = save_to_p(self, param, p, fourier_error)
+%
+% 
+% ** self      structure containing inputs: e.g. current reconstruction results, data, mask, positions, pixel size, ..
+% ** param       structure containing parameters for the engines 
+% ** p          ptychoshelves p structure 
+% ** fourier_error  array [Npos,1] containing evolution of reconstruction error 
+%
+% returns: 
+% ** p_out       updated ptychoshelves p structure 
+
+
+
+% 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 mixed coherent modes:
+% P. Thibault and A. Menzel, Reconstructing state mixtures from diffraction measurements, Nature 494, 68–71 (2013). (doi: 10.1038/nature11806),
+% for LSQ-ML method 
+% M. Odstrcil, A. Menzel, M.G. Sicairos,  Iterative least-squares solver for generalized maximum-likelihood ptychography, Optics Express, 2018
+% for OPRP method 
+%  M. Odstrcil, P. Baksh, S. A. Boden, R. Card, J. E. Chad, J. G. Frey, W. S. Brocklesby,  "Ptychographic coherent diffractive imaging with orthogonal probe relaxation." Optics express 24.8 (2016): 8360-8369
+% 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). 
+% 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 p_out = save_to_p(self, param, p, fourier_error)
+   
+    import utils.*
+    Np_p = self.Np_p; 
+
+    p_out = p; 
+    
+
+    %% return calculated values back to the main code (p-structure)
+    p_out.probe_modes = length(self.probe);
+    p_out.numprobs = size(self.probe{1},3);
+    for ll = 1:p_out.probe_modes
+        if size(self.probe{ll},3) == self.Npos
+            % classical OPRP method 
+            for ii = 1:p.numscans
+                probes(:,:,ii,ll) = mean(self.probe{ll}(:,:,self.reconstruct_ind{ii}),3);
+            end
+             % save the variable modes 
+             [X,V] = extract_PCA(self.probe{ll}, param.variable_probe_modes);
+             p_out.probe_PCA.eigen_vec = X * (prod(sqrt(Np_p))*2*p.renorm);% normalization for consistency with the CPU code
+             p_out.probe_PCA.evolution = V;
+        else
+            % store constant part 
+            probes(:,:,1:size(self.probe{ll},3),ll) = self.probe{ll}(:,:,:,1);
+            if ndims(self.probe{ll}) == 4
+                p_ind = zeros(self.Npos,1);
+                for kk = 1:length(self.reconstruct_ind)
+                   p_ind = p_ind + kk*ismember(1:self.Npos,  self.reconstruct_ind{kk})';
+                end
+                p_out.probe_variable.eigen_vec = self.probe{ll}(:,:,:,2) * (prod(sqrt(Np_p))*2*p.renorm);% normalization for consistency with the CPU code
+                p_out.probe_variable.evolution = self.probe_evolution;
+            end
+        end
+    end
+            
+    
+    % normalization for consistency with the CPU code
+    probes = probes * (prod(sqrt(Np_p))*2*p.renorm); 
+    p_out.probes = probes; 
+   
+    
+    position_offset = 1+floor((p.object_size-self.Np_p)/2);
+
+    % for consistency with the CPU code revert the object to the original size 
+    p_out.numobjs = size(self.object,1);
+    p_out.object = cell(1,p_out.numobjs);
+    for i = 1:p_out.numobjs
+        obj_size = p.object_size(min(end,p.share_object_ID(i)),:); 
+        p_out.object{i} = single([]); 
+        for layer = 1:param.Nlayers
+            p_out.object{i}(:,:,1,layer) = imshift_fast(self.object{i,layer}, -1,-1, obj_size, 'nearest', mean(self.object{i,layer}(:)));
+        end
+    end
+
+    for i = 1:p.numscans
+        id = p.share_object_ID(i); 
+        obj_size = p.object_size(min(end,id),:); 
+        p_out.illum_sum{id} = imshift_fast(self.illum_sum{id}, -1,-1, obj_size, 'nearest');
+    end
+    
+    if param.probe_position_search < param.number_iterations
+        % store the refined positions 
+        p_out.positions = self.probe_positions(:,[2,1]);
+        % return to the original coordinates 
+        p_out.positions_0 = self.probe_positions_0(:,[2,1]);
+        for i = 1:length(self.reconstruct_ind)
+            ind = self.reconstruct_ind{i};
+            p_out.positions(ind,:) = p_out.positions(ind,:) + position_offset(p.share_object_ID(i),:);
+            p_out.positions_0(ind,:) = p_out.positions_0(ind,:) + position_offset(p.share_object_ID(i),:);
+        end
+    else
+        for i = 1:length(self.reconstruct_ind)
+            ind = self.reconstruct_ind{i};
+            if ~isempty(self.probe_positions)
+                p_out.positions(ind,:) = self.probe_positions(ind,[2,1]) + position_offset(p.share_object_ID(i),:);
+            else
+                p_out.positions(ind,:) = self.probe_positions_0(ind,[2,1]) + position_offset(p.share_object_ID(i),:);
+            end
+        end
+    end
+    
+
+    if param.probe_position_search < param.number_iterations || param.detector_rotation_search < param.number_iterations  || param.detector_scale_search < param.number_iterations
+        p_out = engines.GPU_MS.analysis.report_refined_geometry(self, param, p_out); 
+    end
+ 
+
+    % save additional reconstructed parameters  
+    for item = {'background', 'intensity_corr', 'probe_fourier_shift' }
+        try
+            p_out.(item{1}) = self.(item{1});
+        end
+    end
+    
+    % save error metrics 
+    ind_ok = any(~isnan(fourier_error),2); % plot only the reported values 
+    p_out.error_metric.value = nanmean(fourier_error(ind_ok,:),2);
+    p_out.error_metric.iteration = find(ind_ok);
+    if strcmp(param.likelihood,'poisson' )
+        p_out.error_metric.err_metric = 'poisson';
+    else
+        p_out.error_metric.err_metric = 'L1';
+    end
+    p_out.error_metric.method = ['GPU-',param.method, ' metric:' , param.likelihood ];
+
+    
+end