% Solve the cases in the J. Opt. Soc. Am. A 2005 paper: Wave-front reconstruction from its gradients.
% S2p package: functions for calculation of phase phi from shears Sx, Sy
% Inputs: 
% 	Nx, Ny: dimensions
% 	Sx,Sy = shears, differences in phase, of size (1:Nx,1:Ny)
% 	    Sx(1:Ny,1:Nx-1) = phi(1:Ny,2:Nx)-phi(1:Ny,1:Nx-1)
% 	    Sx(1:Ny,Nx) = 0  or some other method of extension
% 	    
% 	    Sy(1:Ny-1,1:Nx) = phi(2:Ny,1:Nx)-phi(1:Ny-1,1:Nx)
% 	    Sx(Ny,1:Nx) = 0  or some other method of extension
% Output:  
% 	phi(1:Ny,1:Nx) = phase
% 
% functions:
%   phi=cos_s2p (Sx,Sy) - current nonperiodic method
%   phi=fft_s2p (Sx,Sy) - previous periodic FFT
%   
%   [var,dsx,dsy]=err_s2p(phi,Sx,Sy) - calculate the errors (how good a solution)
%   var = variance = sum{ |dsx|^2 + |dsy|^2 }
%   dsx = difference = phi(y,x+1)-phi(y,x)-Sx(y,x)
%   dsy = difference = phi(y+1,x)-phi(y,x)-Sy(y,x)
%
% Amos Talmi and Erez Ribak 2005

% Hit any key to continue
help main_s2p, %pause
close all; clear all
Nx = 64; Ny = 128; 		% define array size and other useful arrays
X1 = [1:Nx]; Y1 = [1:Ny]; oneX = ones (1, Nx); oneY = ones (1, Ny);
S = zeros (Ny, Nx);  		% initialize slope arrays to zero
preparrays 			% create analytic and random arrays and calc slopes without and with noise 

sxc = ['oneY'' * oneX                         ' ; ...               % given x slopes
        'oneY'' * (X1+10)                      ' ; ...
        '(Y1+10)''*oneX                        ' ; ...
        '(Y1-16.5)'' * ((X1/64).* (1- (X1/64)))' ; ...
        'Px                                   ' ; ...
        'Ax                                   ' ; ...
        'Bx                                   '];
syc = ['S ' ; ...                                                  % given y slopes
        'S ' ; ...
        'S ' ; ...
        'S ' ; ...
        'Py' ; ...
        'Ay' ; ...
        'By'];
pc = ['S' ; ...                                                     % given y slopes
        'S' ; ...
        'S' ; ...
        'S' ; ...
        'P' ; ...
        'A' ; ...
        'A'];
cases = ['Sx = 1, Sy = 0                            ' ; ...    % list given slopes
        'Sx = n+10 ,  Sy = 0                       ' ; ...
        'Sx(n+½, m) = m+10, Sy = 0                 ' ; ...
        'Sx(n+½, m) = (m-16.5)*n/64(1-n/64), Sy = 0' ; ...
        'Analytic input                            ' ; ...
        'Random input                              ' ; ...
        'Random input with noise on slopes         '];

format short; format compact 
for c = 1:7
    Sx = eval (sxc (c, :));  		% use x slopes above
    Sy = eval (syc (c, :)); 		% use y slopes above
    phi = eval (pc (c, :));  		% use function above (analytic, random cases)
    
    phi1 = cos_s2p (Sx, Sy);  		% calc function from slopes, our method
    phi2 = fft_s2p (Sx, Sy);  		% calc function from slopes, periodic 
    
    [var1, dsx, dsy] = err_s2p (phi1, Sx, Sy);   % errors in our method
    max_dev_cos = max (max (max (abs (dsx))), max (max (abs (dsy))));
    [var2, dsx, dsy] = err_s2p (phi2, Sx, Sy);   % errors in FFT method
    max_dev_fft = max (max (max (abs (dsx))), max (max (abs (dsy))));
    
    std_cos = sqrt (abs (var1)/Nx/Ny); 		% normalised standard deviations
    std_fft = sqrt (abs (var2)/Nx/Ny); 		% normalised standard deviations
%   disp (['case ' int2str(c) ': ' cases(c,:)]) % list cases above  
%   disp ('STDcos MaxDevCos STDFFT MaxDevFFT' )
%   disp ([std_cos, max_dev_cos, std_fft, max_dev_fft])
    fprintf('Case %d: %s. STDcos: %8.4g; MaxDevCos: %8.4g; STDFFT: %8.4g; MaxDevFFT: %8.4g\n', c, cases(c,:),[max_dev_cos, max_dev_fft, std_cos, std_fft])
    
    if c > 4 % analytic and random functions, without and with slopes noise
        figure (c); 
        subplot (131); imagesc(phi ); axis image;axis xy; colorbar h; title (cases(c,:))
if c>6,   subplot (131); imagesc(phi1); axis image;axis xy; colorbar h; title (cases(c,:)), end
        phi1 = phi1 - phi; phi1 = (phi1 - min (phi1(:))) .* R2; 	% difference from input
        phi2 = phi2 - phi; phi2 = (phi2 - min (phi2(:))) .* R2; 	% difference from input
        subplot (132); imagesc(phi1); axis image;axis xy; colorbar h; title ('cos error')
        subplot (133); imagesc(phi2); axis image;axis xy; colorbar h; title ('fft error')
    end
    
end
format

%% preparearrays

% create known analytic array and slopes
% create random array by smoothing a random array with a power spectrum -1 
% calculate slopes from random array, and then add 50%  noise to the slopes

Qx = 3.7; Qy = 2.4; 							% periods of analytic array
phi = cos (Qy* (0.5:Ny-0.5)*pi/Ny)'*cos (Qx* (0.5:Nx-0.5)*pi/Nx) ;	% create array
phi= 1e2 * (phi + 0.2 * oneY' * X1 / Nx - 0.4 * Y1' * oneX  / Ny);	% adds small linear drift
Px = S; Py = S;   P = phi;    						% slopes
Px (1:Ny, 1:Nx - 1) = phi (1:Ny, 2:Nx) - phi (1:Ny, 1:Nx - 1);
Py (1:Ny - 1, 1:Nx) = phi (2:Ny, 1:Nx) - phi (1:Ny - 1, 1:Nx);

% random function, smoothed
l = 2 * max (Nx, Ny); lh = l/2+1; x = -l/2:l/2-1;  xh = 1:l; 		%l/4:3*l/4; % sizes
[XX, YY]  = meshgrid (x, x); R2 = XX.^2+YY.^2; R2 (lh, lh) = 1; 	% create r^2 array
XX  = real (ifft2 (fft2 ((rand(l) - 0.5)) .* fftshift (R2.^-1)));  	% smoothed random function 
XX = 1e4 * (XX  + .03*ones(l,1)*x/l+ .02*x'*ones(1,l)/l); % add small slope
R2 = R2<950;      					% create circle
R2 = 1;                 				% comment this line for circle case
YY = XX .* R2;     					% circle or full array
if R2 (1) ~= 1, R2 = R2 (Y1-Ny/2+l/2, X1-Nx/2+l/2); end % centre arrays
phi = YY (Y1-Ny/2+l/2, X1-Nx/2+l/2) .* R2;             	% chop out relevant central part

Ax = S; Ay = S;  A = phi; 				% slopes
Ax (1:Ny, 1:Nx-1) = phi (1:Ny, 2:Nx) - phi (1:Ny, 1:Nx-1);
Ay (1:Ny-1, 1:Nx) = phi (2:Ny, 1:Nx) - phi (1:Ny-1, 1:Nx);

% add noise to slopes
avg_abs_Sx= sqrt (trace(Ax*Ax')/ Nx );  avg_abs_Sy = sqrt (trace(Ay*Ay')/ Nx );
Noise = 50;    						 %noise strength, percent
Bx = Ax + (avg_abs_Sx/8*sqrt(Noise)/100)*(poissrnd(Noise, Ny, Nx) - Noise);      % add noise
By = Ay + (avg_abs_Sy/8*sqrt(Noise)/100)*(poissrnd(Noise, Ny, Nx) - Noise);      % Strength/100 noise


function phi = cos_s2p (Sx, Sy)
%function phi = cos_s2p (Sx, Sy, Nx, Ny)
% phi = cos_s2p (Sx, Sy, Nx, Ny) solves exactly the shears - to - phase problem with Cos series
% Sx (y, x) = X_shear = phi (y, x + 1) - phi (y, x)
% Sy (y, x) = Y_shear = phi (y + 1, x) - phi (y, x)
% phi, Sx, Sy = (1:Ny, 1:Nx) Sx (y, Nx) = Sy (Ny, x) = 0
% You may use init_s2p () for dummy Sx, Sy

% To calculate the Sine Transform, we use FFT of 2N points
[Ny, Nx] = size (Sx);
oneX = ones (1, 2*Nx); oneY = ones (1, 2*Ny);

% calc the fourier transform: Sx is sin (x) cos (y)!!
% - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 

tQy = [0:pi/Ny/2:pi]; tQy = exp (-i*tQy); ttqy = tQy.' ; 
tQx = [0:pi/Nx/2:pi]; tQx = exp (-i*tQx);
QX = [2:Nx]; QY = [2:Ny];
cQX = 2 + 2*Nx - QX; cQY = 2 + 2*Ny - QY;
toneY = oneY';
K = 2i/Nx/Ny;

work = zeros (2*Ny, 2*Nx);
work (1:Ny, QX) = Sx (1:Ny, QX - 1); 
work (1:Ny, cQX) = -work (1:Ny, QX);		% work (2 + b, y) = work (2Nx - b, y) = Sx (1 + b, y)
gw = fft2 (work);				% gw (qy) = sum{ exp[ (qy - 1) (y - 1)2pi/2Ny]*F (y)

%Fourier Cos (y) sin (x) transform of Sx
SQx (QY, QX) = (0.5*K)* (gw (QY, QX).* (ttqy (QY)*oneX (QX)) - gw (cQY, QX).* (ttqy (cQY)*oneX (QX)) );
SQx (1, QX) = K*gw (1, QX);
SQx (1:Ny, 1) = 0;

work = zeros (2*Ny, 2*Nx);
work (QY, 1:Nx) = Sy (QY - 1, 1:Nx);
work (cQY, 1:Nx) = - work (QY, 1:Nx);			%anti - symmetry in y
gw = fft2 (work);

SQy (QY, QX) = (0.5*K)* (gw (QY, QX).* (toneY (QY)*tQx (QX)) - gw (QY, cQX).* (toneY (QY)*tQx (cQX)) );
SQy (QY, 1 ) = K*gw (QY, 1);
SQy (1, 1:Nx) = 0;

%calc the Qg factors: Qgx = 2*sin (Q*pi/2/Nx) = 2*imag (tQx (Q - 1))
Qgx = -2*imag (tQx); Qgy = -2*imag (tQy);
Qgx (1) = 1e-50;

%calc FQ = fi_Q
for x = 1:Nx, 
    for y = 1:Ny, 
        FQ (y, x) = - (Qgx(x) * SQx(y, x) + Qgy(y) * SQy(y, x) )/ (Qgx(x)^2 + Qgy(y)^2);
    end ;
end ;
FQ (1, 1) = 0;

% inverse FFT
work = zeros (2*Ny, 2*Nx);
X1 = [1:Nx]; Y1 = [1:Ny];
work (Y1, X1) = FQ (Y1, X1).* (ttqy (Y1)*tQx (X1));
work (cQY, X1) = FQ (QY, X1).* (conj (ttqy (QY))*tQx (X1));
work (Y1, cQX) = FQ (Y1, QX).* (ttqy (Y1)*conj (tQx (QX)));
work (cQY, cQX) = FQ (QY, QX).* (conj (ttqy (QY))*conj (tQx (QX)));
%Note: work (Ny + 1, :) = work (:, Nx + 1) = 0 !!
gw =  real (fft2 (work));
phi (Y1, X1) = 0.25*gw (Y1, X1);


function [var, Dsx, Dsy] = err_s2p (phi, Sx, Sy)
%function [var,Dsx,Dsy]= err_s2p(phi,Sx,Sy,Nx,Ny)
% [variance, Dsx, Dsy]=err_s2p(phi,Sx,Sy,Nx,Ny) returns errors in solution
%  Dsx= phi(y,x+1)-phi(y,x)-Sx(y,x)
%  Dsy= phi(y+1,x)-phi(y,x)-Sy(y,x)
%   variance=sum(Dsx^2+Dsy^2)

[Ny, Nx] = size (Sx);
Dsx=zeros(Ny,Nx); Dsy=Dsx;
Dsx(1:Ny,1:Nx-1)=phi(1:Ny,2:Nx)-phi(1:Ny,1:Nx-1)-Sx(1:Ny,1:Nx-1);
Dsy(1:Ny-1,1:Nx)=phi(2:Ny,1:Nx)-phi(1:Ny-1,1:Nx)-Sy(1:Ny-1,1:Nx);
var=trace(Dsx*Dsx') +trace(Dsy*Dsy');


function phi = fft_s2p (Sx, Sy)
% phi = fft_s2p (Sx, Sy) solves the shears - to - phase problem using periodic exp (ik*2pi/N) series
% Sx (y, x) = X_shear = phi (y, x + 1) - phi (y, x)
% Sy (y, x) = Y_shear = phi (y + 1, x) - phi (y, x)
% phi, Sx, Sy = (1:Ny, 1:Nx) 
% You may use init_s2p () for dummy Sx, Sy

[Ny, Nx] = size (Sx);
oneX = ones (1, Nx); oneY = ones (1, Ny);
% calc the fourier transform: Sx is exp (-ikx*x*2pi/Nx)*exp (-iky*y*2pi/Ny)
SQx = fft2 (Sx); SQy = fft2 (Sy);
%calc the Qg factors.: Qgx = 2*sin (Q*pi/Nx) 
QX = [0:Nx - 1]; QY = [0:Ny - 1];
Qgx = exp (i*2*pi/Nx*QX) - 1; 	% 2*sin (X*pi/Nx) if attributed to the middle;
Qgy = exp (i*2*pi/Ny*QY) - 1; 	% 2*sin (Y*pi/Ny);
Qgx (1) = eps;			                    % avoid small number division

%calc FQ = fi_Q
for qx = 1:Nx, 
    for qy = 1:Ny, 
        FQ (qy, qx) = (conj (Qgx (qx))*SQx (qy, qx) + conj (Qgy (qy))*SQy (qy, qx) )/ ( Qgx (qx)*conj (Qgx (qx)) + Qgy (qy)*conj (Qgy (qy)) );
    end ;
end ;
FQ (1, 1) = 0;
phi = real (ifft2 (FQ));  % inverse fft
