% This is a supplemental material to the article
% "Three-dimensional propagation in near-field tomographic X-ray phase retrieval"
% by A. Ruhlandt and T. Salditt.
%
% First, a 3d phantom is created, projected and propagated, allowing a
% comparison between 2d and 3d propagation.
% Second, a 2d and a 3d iterative reconstruction scheme can be carried out
% to illustrate the benefit of enforcing tomographical consistency during
% the phase retrieval.


%% settings for phantom and propagators

N=256;                  % N^2: number of pixels, N^3: number of voxels
lambda = 0.01;          % wavelength in size of pixels
dist = 2000;            % propagation distance in size of pixels
angles=round(N*pi/2);	% number of projections (satisfy angular sampling here)


%% create 3d and 2d propagators

k= 2*pi/lambda;
k_min2= (2*pi/N).^2;
[X,Y,Z] = meshgrid(1:N,1:N,1:N);

% propagator from object to detection plane
obj2det = ifftshift( exp(1i*dist*sqrt(k^2 - k_min2*(X-N/2-1).^2 - k_min2*(Y-N/2-1).^2 - k_min2*(Z-N/2-1).^2)) );


% propagator from detection to object plane
det2obj = 1./obj2det;
obj2det2d  = squeeze(obj2det(:,:,1));
det2obj2d= squeeze(det2obj(:,:,1));
figure(1); imagesc(real(obj2det2d)); title 'shifted propagator to avoid fftshifts'; colorbar;


%% create an object

delta_min = 0.001;          % minimal value of k*delta for one voxel
delta_max = 0.01;           % maximal value
beta_min = 0;               % minimal value of k*beta for one voxel
beta_max = delta_max * 0.07;% maximal value

N_spheres = 20;             % number of spheres within the volume
sphere_min_r = 10;          % minimum radius
sphere_max_r = 25;          % maximum radius

N_cubes = 15;               % number of cubes within the volume
cubes_min_a = 12;           % minimum side length
cubes_max_a = 35;           % maximum side length

border = 40;                % border from the outside where no object should appear

object = zeros(N,N,N);

% add spheres
for i=1:N_spheres
    r = floor(sphere_min_r + (sphere_max_r-sphere_min_r)*rand);
    
    [X,Y,Z] = meshgrid(-r:r,-r:r,-r:r);
    sph = X.^2+Y.^2+Z.^2;
    sph(sph<r^2)=1;
    sph(sph>1)=0;
    
    pos_x = 1+border+round(rand*(N-1-size(sph,2)-2*border));
    pos_y = 1+border+round(rand*(N-1-size(sph,1)-2*border));
    pos_z = 1+border+round(rand*(N-1-size(sph,3)-2*border));
    
    delta = delta_min + (delta_max-delta_min)*rand;
    beta = beta_min + (beta_max-beta_min)*rand;
    sph = -(1i*delta*sph + beta*sph);

    object(pos_y:pos_y+size(sph,1)-1, pos_x:pos_x+size(sph,2)-1, pos_z:pos_z+size(sph,3)-1) = object(pos_y:pos_y+size(sph,1)-1, pos_x:pos_x+size(sph,2)-1, pos_z:pos_z+size(sph,3)-1) + sph;
end

% add cubes
for i=1:N_cubes
    a = floor(cubes_min_a + (cubes_max_a-cubes_min_a)*rand);
    
    pos_x = 1+border+round(rand*(N-1-a-2*border));
    pos_y = 1+border+round(rand*(N-1-a-2*border));
    pos_z = 1+border+round(rand*(N-1-a-2*border));
    
    delta = delta_min + (delta_max-delta_min)*rand;
    beta = beta_min + (beta_max-beta_min)*rand;
    val = -(1i*delta + beta);
    
    object(pos_y:pos_y+a-1, pos_x:pos_x+a-1, pos_z:pos_z+a-1) = object(pos_y:pos_y+a-1, pos_x:pos_x+a-1, pos_z:pos_z+a-1) + val;
end

% show a first projection
imagesc(squeeze(imag(sum(object,3)))); colorbar; title 'projection of k*delta'; colorbar;


%% project phantom to equidistant angles between 0 and 180 deg

thetas=[0:180/angles:180-(90/angles)];      % projection angles

for i=1:N
    temp=radon(squeeze(object(i,:,:)),thetas);
    projs(:,:,i)=temp(57:57+255,:);         % rounded values for this demonstration
    i
end
figure(2); imagesc(real(squeeze(projs(:,1,:)))); colormap 'gray'; colorbar; title 'absorption';
figure(3); imagesc(imag(squeeze(projs(:,1,:)))); colormap 'gray'; colorbar; title 'phase shift';


%% propagate 2d projections to detection plane

for i=1:numel(thetas)
    psi = squeeze(exp(projs(:,i,:)));
    psi = ( ifft2( obj2det2d.*fft2(psi) ) );
    detector(:,i,:) = abs(psi);         % save modulus of the propagated projections

% uncomment the following lines for noisy data
%   photonsperpixel=10000;
%   det = squeeze(detector(:,i,:).^2);    % add noise to the intensities
%   det = det .* photonsperpixel;
%   det = 1E12*imnoise(floor(det)*1E-12,'poisson')./photonsperpixel;
%   detector(:,i,:) = sqrt(det);          % save the modulus
end
figure(4); imagesc(squeeze(detector(:,1,:))); colormap 'gray'; colorbar; title 'abs 2d propagation';


%% for comparison: propagate 3d volume and project afterwards

psi3d = ifftn(obj2det.*fftn(object));
messung = abs(psi3d+1);
figure(5); imagesc(fliplr(rot90(squeeze(sum(messung-1,2))+1,3))); colormap 'gray'; title 'abs projected 3d propagation'; colorbar;


%% RECONSTRUCTION 2d: step 1: phase retrieval

soft_coupling = 0.1;        % soft coupling parameter

for frame=1:size(detector,2)
    
    current_meass = squeeze(detector(:,frame,:));
    current_guess = ones(N,N);
    for ii=1:1000   % 1000 iterations
        % first: modulus constraint at detection plane
        ref_guess = ifft2( obj2det2d.*fft2(current_guess) );
        ref_guess = current_meass .* ref_guess./abs(ref_guess);
        ref_guess = ifft2( det2obj2d.*fft2(ref_guess) );

        % second: delta >= 0 and beta >= 0 at object position
        ref_guess = log(ref_guess);
        ref_guess = -abs(real(ref_guess)) -1i*abs(imag(ref_guess));
        
        % third: soft_coupling*delta >= beta
        gt=(imag(ref_guess)>real(ref_guess));
        ref_guess(gt) = real(ref_guess(gt))+1i*real(ref_guess(gt));
        
        gt=(real(ref_guess) < imag(ref_guess)*soft_coupling);
        ref_guess(gt) = soft_coupling*imag(ref_guess(gt))+1i*imag(ref_guess(gt));
        
        ref_guess = exp(ref_guess);
        current_guess = ref_guess;
    end
    it_reco2d(:,frame,:) = current_guess;
    
    % show the current guess
    figure(6);
    headline=sprintf('2d phase reconstruction (frame %i) after %i iterations',frame,ii);
    imagesc(angle(ref_guess)); title(headline); colormap 'gray'; colorbar;
    drawnow;
end


%% RECONSTRUCTION 2d: step 2: tomography

% filtered backprojection
for i=1:N
reco2d(:,:,i) = iradon(real(log(it_reco2d(:,:,i))),thetas,'linear','Ram-Lak',1,N);
reco2d(:,:,i) = reco2d(:,:,i) + 1i * iradon(imag(log(it_reco2d(:,:,i))),thetas,'linear','Ram-Lak',1,N);
i
end

% reproject the volume
for i=1:N
	reco2d_reproj(:,:,i)=radon(reco2d(:,:,i),thetas);
    i
end
figure(7);
imagesc(squeeze(imag(reco2d_reproj(:,1,:)))); title('2d iterative reco - reprojection'); colormap 'gray'; colorbar;



%% RECONSTRUCTION 3d: step 1: tomography

%reconstruct volume from holograms in a way it works as a modulus constraint
[X,Y] = meshgrid(1:N,1:N);
for i=1:N
     modulus3d(:,:,i)=iradon(detector(:,:,i)-1,thetas,'linear','Ram-Lak',1,N) + 1;
     i
end

%% RECONSTRUCTION 3d: step 2: phase retrieval
reco3d = ones(N,N,N);
% uncomment the following lines to use a GPU for faster reconstruction
% reco3d=gpuArray(reco3d);
% modulus3d=gpuArray(modulus3d);
% obj2det=gpuArray(obj2det);
% det2obj=gpuArray(det2obj);

for ii=1:1000   % 1000 iterations
    % first: modulus constraint at detection plane
    guess = ifftn( obj2det.*fftn(reco3d +1) );
    guess = modulus3d .* guess./abs(guess);
    guess = ifftn( det2obj.*fftn(guess -1) );

	% second: delta >= 0 and beta >= 0 at object position
	guess = -abs(real(guess)) -1i*abs(imag(guess));
    
    % third: beta <= soft_coupling*delta
    gt=(imag(guess)>real(guess));
	guess(gt) = real(guess(gt))+1i*real(guess(gt));
    
    gt=(real(guess) < imag(guess)*soft_coupling);
    guess(gt) = soft_coupling*imag(guess(gt))+1i*imag(guess(gt));
    
    reco3d=guess;
    ii
end


% reproject the volume
for i=1:N
	reco3d_reproj(:,:,i)=radon(reco3d(:,:,i),thetas);
    i
end

reco3d=gather(reco3d);                  % get the data back from the GPU
reco3d_reproj = gather(reco3d_reproj);
figure(8);
imagesc(squeeze(imag(reco3d_reproj(:,1,:)))); title('3d iterative reco - reprojection'); colormap 'gray'; colorbar;