% B.I.Stepanov Institute of Physics
% National Academy of Sciences of Belarus

% Aerosol backscatter and lidar ratio retrieval algorithm utilizing Raman lidar 
% signals.

function raman

global k d;
global etaA etaM;
global hStep Nbeg Nfin N0 N;
global betaMref Omega betaM Lmeas;
global tauAMatrix smoothOp2 smoothOp1;
global iter;
global ThetaArr GammaArr RArr Psi1Arr Psi2Arr LmeasRArr LcalcArr;
global gammaMiddle;
global stage;

%-------------------------------------------------------------------------
disp('Reading input data...');

% Lidar signal weighting coefficient (by channel).
k = hdf5read('data.h5', '/input/k');
% Profile smoothness weighthing coefficient (by aerosol characteristic).
d = hdf5read('data.h5', '/input/d');

% Aerosol attenuation ratio between Raman and source lidar channels.
etaA = hdf5read('data.h5', '/input/etaA');
% Molecular backscatter ratio between Raman and source lidar channels.
etaM = hdf5read('data.h5', '/input/etaM');

% Height grid step, in meters.
hStep = hdf5read('data.h5', '/input/hStep');
% First index of lidar signal data selected for retrieval (by channel).
% Add 1 to obtain Matlab's 1-based indices.
Nbeg = hdf5read('data.h5', '/input/Nbeg') + 1;
% Reference point index (by channel).
Nref = hdf5read('data.h5', '/input/Nref') + 1;
% Last index of lidar signal data selected for retrieval (by channel).
Nfin = hdf5read('data.h5', '/input/Nfin') + 1;

% Molecular backscatter coefficient at the reference point (by channel).
betaMref = hdf5read('data.h5', '/input/betaMref');

% Measured lidar signal, divided by its average value around the reference 
% point (by height, channel).
S = hdf5read('data.h5', '/input/S');
% Weighting coefficients for lidar signal (by height, channel).
Omega = hdf5read('data.h5', '/input/Omega');
% Molecular backscatter coefficient (by height, channel).
betaM  = hdf5read('data.h5', '/input/betaM');
% Molecular optical thickness, relative to the reference point, with 
% positive values below it (by height, channel).
tauM = hdf5read('data.h5', '/input/tauM');

% Zenith angle of the lidar track, in degrees (by height, channel).
Z = hdf5read('data.h5', '/input/Z');

% Initial approximation to aerosol backscatter ratio (by height).
Theta0 = hdf5read('data.h5', '/input/Theta0');
% Initial approximation to aerosol lidar ratio (by height).
Gamma0 = hdf5read('data.h5', '/input/Gamma0');
% Initial approximation to lidar correction coefficients (by channel).
R0 = hdf5read('data.h5', '/input/R0');
% Maximal and minimal allowed values for the lidar correction coefficients.
RLower = hdf5read('data.h5', '/input/RLower');
RUpper = hdf5read('data.h5', '/input/RUpper');

TolFun = hdf5read('data.h5', '/input/TolFun');
TolX = hdf5read('data.h5', '/input/TolX');

% Boundaries of aerosol profiles to be retrieved.
N0 = min(Nbeg);
N = max(Nfin);

%-------------------------------------------------------------------------
% Check consistency of input array dimensions.
assert(all(size(k) == [2, 1]));
assert(all(size(d) == [3, 1]));

assert(all(size(etaA) == [1, 1]) && etaA > 0);
assert(all(size(etaM) == [1, 1]) && etaM > 0);

assert(all(size(hStep) == [1, 1]) && hStep > 0);
assert(all(size(Nbeg) == [2, 1]));
assert(all(size(Nref) == [2, 1]));
assert(all(size(Nfin) == [2, 1]));
assert(all(size(betaMref) == [2, 1]));

assert(all(size(S) == [N, 2]));
assert(all(size(Omega) == [N, 2]));
assert(all(size(betaM) == [N, 2]));
assert(all(size(tauM) == [N, 2]));

assert(all(size(Z) == [N, 2]));

assert(all(size(Theta0) == [N, 1]));
assert(all(size(Gamma0) == [N, 1]));
assert(all(size(R0) == [2, 1]));
assert(all(size(RLower) == [2, 1]));
assert(all(size(RUpper) == [2, 1]));

assert(all(size(TolFun) == [1, 1]) && TolFun > 0);
assert(all(size(TolX) == [1, 1]) && TolX > 0);

%-------------------------------------------------------------------------
% Measured lidar signal, divided by its average value around the reference 
% point, with molecular component removed, without the backscatter ratio
% correction factor (by height, channel).
Lmeas = zeros(N, 2);

Lmeas(Nbeg(1):Nfin(1), 1) = S(Nbeg(1):Nfin(1), 1) .* ...
    exp(-2 * tauM(Nbeg(1):Nfin(1), 1));

Lmeas(Nbeg(2):Nfin(2), 2) = S(Nbeg(2):Nfin(2), 2) .* ...
    exp(-(1 + 1 / etaM) * tauM(Nbeg(2):Nfin(2), 2));
    
for j = 1:2
    assert(all(Omega(Nbeg(j):Nfin(j)) > 0))
    
    % Extend arrays related to lidar signal equations below the lower
    % signal boundaries. Thus signal equations in the near-to-the-ground
    % region will be the same, and retrieved concentrations in that region
    % will tend to be near-constant.
    Lmeas(1:Nbeg(j)-1, j) = Lmeas(Nbeg(j), j);
    Omega(1:Nbeg(j)-1, j) = Omega(Nbeg(j), j);
    betaM(1:Nbeg(j)-1, j) = betaM(Nbeg(j), j);
    Z(1:Nbeg(j)-1, j) = Z(Nbeg(j), j);
end

% Matrices that transform extinction coefficient vectors to optical
% thickness vectors, relative to the reference points (by channel).
tauAMatrix = cell(2, 1);

for j = 1:2
    % Steps along the lidar track, with zenith angles taken into account.
    effStep = hStep ./ cos(degtorad(Z(1:Nfin(j), j)));
    
    tauAMatrixTmp = tauIntegrationMatrix(...
        Nbeg(j), Nref(j), Nfin(j), effStep);
    
    tauAMatrix{j} = tauAMatrixTmp(Nbeg(j):Nfin(j), Nbeg(j):Nfin(j));
end

% Matrices that transform concentration profiles to vector evaluating their 
% somoothness.
smoothOp2 = secondDifferenceOperator(N-N0+1, hStep);
smoothOp1 = firstDifferenceOperator(N-N0+1, hStep);

%-------------------------------------------------------------------------
% Current iteration number + 1 (1 means no iterations have completed yet).
iter = 0;

% The following arrays will be used to store snapshots of the optimization 
% process at every iteration as well as the retrieval results. They will 
% be 3-dimensional; new layers will be added to these arrays in 'optfun'.

% Preallocate space for 2 iterations to make sure that the arrays will remain 
% 3-dimensional even in the non-realistic case where initial approximation 
% itself triggers termination of the retrieval process.

% Aerosol backscatter ratio profile.
ThetaArr = zeros(N, 1, 2);
% Aerosol lidar ratio profile.
GammaArr = zeros(N, 1, 2);
% Lidar correction coefficients.
RArr = zeros(1, 2, 2);

% Residuals for lidar equations.
Psi1Arr = zeros(1, 2, 2);
% Residuals for smoothness criteria.
Psi2Arr = zeros(1, 3, 2);

% Measured lidar signal, divided by its average value around the reference 
% point, with molecular component removed (by height, channel).
LmeasRArr = zeros(N, 2, 2); 
% Calculated lidar signal, divided by its value at the reference point,
% with molecular component removed (by height, channel).
LcalcArr = zeros(N, 2, 2);

%-------------------------------------------------------------------------
disp('Starting optimization...');

% First stage of the optimization process: lidar ratio is a scalar.
stage = 0;

x0 = [Theta0(N0:N); sum(Gamma0(N0:N) / (N-N0+1)); R0];

l0 = zeros(size(x0));
u0 = ones(size(x0)) * inf;

% Allow negative values in aerosol backscatter ratio, which may be caused
% by artifitial noise in perturbance analysis.
l0(1:N-N0+1) = -inf;

% Limit variability of the lidar correction coefficients.
l0((N-N0+1)+1 + 1 : (N-N0+1)+1 + 2) = RLower;
u0((N-N0+1)+1 + 1 : (N-N0+1)+1 + 2) = RUpper;

gammaMiddle = 0.0;

optopt = optimset('Display', 'Iter', 'DerivativeCheck', 'off', ...
    'TolFun', TolFun, 'TolX', TolX, 'Jacobian', 'on');

% Invoke optimizer.
lsqnonlin(@optfun, x0, l0, u0, optopt);

gammaMiddle = sum(GammaArr(N0:N, 1, iter)) / (N-N0+1);

stage = 1;

x0 = [ThetaArr(N0:N, 1, iter); ones(N-N0+1, 1) * gammaMiddle; ...
    RArr(1, 1, iter); RArr(1, 2, iter)];

l0 = zeros(size(x0));
u0 = ones(size(x0)) * inf;

% Allow negative values in aerosol backscatter ratio, which may be caused
% by artifitial noise in perturbance analysis.
l0(1:N-N0+1) = -inf;

% Limit variability of the lidar correction coefficients.
l0((N-N0+1)*2 + 1 : (N-N0+1)*2 + 2) = RLower;
u0((N-N0+1)*2 + 1 : (N-N0+1)*2 + 2) = RUpper;

% Invoke optimizer for the second time.
lsqnonlin(@optfun, x0, l0, u0, optopt);

%-------------------------------------------------------------------------
disp('Saving output data...')

hdf5write('data.h5', '/output/iterCount', iter, 'WriteMode', 'append');

hdf5write('data.h5', '/output/Theta', ThetaArr, 'WriteMode', 'append');
hdf5write('data.h5', '/output/Gamma', GammaArr, 'WriteMode', 'append');
hdf5write('data.h5', '/output/R', RArr, 'WriteMode', 'append');

hdf5write('data.h5', '/output/Psi1', Psi1Arr, 'WriteMode', 'append');
hdf5write('data.h5', '/output/Psi2', Psi2Arr, 'WriteMode', 'append');

hdf5write('data.h5', '/output/LmeasR', LmeasRArr, 'WriteMode', 'append');
hdf5write('data.h5', '/output/Lcalc', LcalcArr, 'WriteMode', 'append');

disp('Done.')
end

%-------------------------------------------------------------------------
% Optical thickness integration operator.
%
% Returns a matrix that transforms a vector of extinction coefficients
% defined on a grid with the given variable step to the vector of optical 
% thicknesses, relative to the reference point 'Nref' (with positive 
% values below it). 'Nfin' defines the size of the matrix. Optical
% thicknesses below the 'Nbeg' point will be the same and equal to that at
% 'Nbeg'.
function tauMatrix = tauIntegrationMatrix(Nbeg, Nref, Nfin, stepVector)
    
    % Assume that 'stepVector(i)' is the distance between nodes i and i+1.

    % Rectangular rule integration matrix, using left corner function 
    % values. This matrix will produce optical thicknesses relative to the
    % ground. Using a sparse matrix speeds up the calculation a bit.
    direct = tril(ones(Nfin), -1) * sparse(1:Nfin, 1:Nfin, stepVector);
    
    % Adjust optical thicknesses below 'Nbeg'.
    direct(1:Nbeg-1, :) = repmat(direct(Nbeg, :), Nbeg-1, 1);
    
    % Vector to produce optical thickness at the reference point relative
    % to the ground, using the same integration rule as in 'direct'.
    reverseVector = stepVector';
    if Nref <= Nfin
        reverseVector(Nref:Nfin) = 0.0;
    end
    
    % To obtain the relative optical thickness, subtract the absolute one
    % from that in the reference point.
    tauMatrix = repmat(reverseVector, Nfin, 1) - direct;
end

% ------------------------------------------------------------------------
% Second order finite difference operator for concentration profiles.
%
% Returns a matrix that transforms a function defined on a grid with the 
% given step to an approximation to that function's second derivative 
% defined on the same grid.
function diffOp = secondDifferenceOperator(N, step)

    % A matrix with '2's on its main diagonal and '-1's on the diagonals
    % just below and just above the main one.
    %
    % Using sparse matrices speeds up the calculation a bit.
    diffOp = sparse(1:N, 1:N, 2) + ...
        sparse(2:N, 1:N-1, -1, N, N) + sparse(1:N-1, 2:N, -1, N, N);
    
    % Remove constraints for the first and the last nodes, as approximations
    % for the second order derivative cannot be obtained there.
    diffOp(1, 1:N) = 0;
    diffOp(N, 1:N) = 0;
    
    diffOp = diffOp / step^2;
end

% ------------------------------------------------------------------------
% First order finite difference operator for concentration profiles.
%
% Returns a matrix that transforms a function defined on a grid with the 
% given step to an approximation to that function's first derivative 
% defined on the same grid.
function diffOp = firstDifferenceOperator(N, step)

    % A matrix with '2's on its main diagonal and '-1's on the diagonals
    % just below and just above the main one.
    %
    % Using sparse matrices speeds up the calculation a bit.
    diffOp = sparse(1:N, 1:N, -1) + sparse(1:N-1, 2:N, 1, N, N);
    
    % Remove constraints for the last node, as approximation for the first 
    % order derivative cannot be obtained there.
    diffOp(N, N) = 0;
    
    diffOp = diffOp / step;
end

%-------------------------------------------------------------------------
function [F, Jac] = optfun(x)

    global k d;
    global etaA etaM
    global hStep Nbeg Nfin N0 N;
    global betaMref Omega betaM Lmeas;
    global tauAMatrix smoothOp2 smoothOp1;
    global iter;
    global ThetaArr GammaArr RArr Psi1Arr Psi2Arr LmeasRArr LcalcArr;
    global gammaMiddle;
    global stage;

    iter = iter + 1;
    
    % Current backscatter ratio profile.
    Theta = reshape(x(1:N-N0+1), N-N0+1, 1);    
    ThetaArr(N0:N, 1, iter) = Theta;
    
    % Current lidar ratio profile.
    if stage == 0
        Gamma = x((N-N0+1)+1);
        GammaArr(N0:N, 1, iter) = ones(N-N0+1, 1) * Gamma;
    else
        Gamma = reshape(x((N-N0+1)+1 : (N-N0+1)*2), N-N0+1, 1);
        GammaArr(N0:N, 1, iter) = Gamma;
    end
    
    % Current lidar correction coefficients.
    if stage == 0
        R = x((N-N0+1)+1+1 : (N-N0+1)+1+2);
    else
        R = x((N-N0+1)*2+1 : (N-N0+1)*2+2);
    end
    
    RArr(1, 1:2, iter) = R;
    
    % Lidar portion of the equations.
    F1 = zeros(sum(Nfin) - sum(Nbeg) + 2, 1);
    if stage == 0
        Jac1 = zeros(sum(Nfin) - sum(Nbeg) + 2, (N-N0+1)+1 + 2);
    else
        Jac1 = zeros(sum(Nfin) - sum(Nbeg) + 2, (N-N0+1)*2 + 2);
    end
    
    for j = 1:2
        % Equation indices for the current channel.
        eqnRange = sum(Nfin(1:j-1)) - sum(Nbeg(1:j-1)) + j-1 + 1 : ...
            sum(Nfin(1:j)) - sum(Nbeg(1:j)) + j;
        
        if j == 1
            betaMCur = betaM(Nbeg(j):Nfin(j), j);
            etaACur = 1.0;
            betaCoeff = 1.0 + Theta(Nbeg(j)-N0+1:Nfin(j)-N0+1);
            betaCoeffDTheta = 1.0;

        elseif j == 2
            betaMCur = betaM(Nbeg(j):Nfin(j), j) / etaM;
            etaACur = etaA;
            betaCoeff = 1.0;
            betaCoeffDTheta = 0.0;
        end
        
        if stage == 0
            sigmaA = betaMCur .* Theta(Nbeg(j)-N0+1 : Nfin(j)-N0+1) .* Gamma;
        else
            sigmaA = betaMCur .* Theta(Nbeg(j)-N0+1 : Nfin(j)-N0+1) .* ...
                Gamma(Nbeg(j)-N0+1 : Nfin(j)-N0+1);
        end
                
        % Calculated aerosol optical thickness, relative to the reference
        % point.
        tauA = tauAMatrix{j} * sigmaA;
        
        % Measured lidar signal, divided by its average value around the 
        % reference point, with molecular component removed.
        LmeasR = Lmeas(Nbeg(j):Nfin(j), j) * R(j);
        % Calculated lidar signal, divided by its value at the reference 
        % point, with molecular component removed.
        Lcalc = betaM(Nbeg(j):Nfin(j), j) ./ betaMref(j) .* betaCoeff .* ...
            exp((1 + etaACur) * tauA);
        
        % Weighting coefficients for lidar equations.
        coeff = k(j) * sqrt(hStep ./ Omega(Nbeg(j):Nfin(j), j));
        
        F1(eqnRange) = (LmeasR - Lcalc) .* coeff;
        
        Psi1Arr(1, j, iter) = norm(F1(eqnRange));
        
        LmeasRArr(Nbeg(j):Nfin(j), j, iter) = LmeasR;
        LcalcArr(Nbeg(j):Nfin(j), j, iter) = Lcalc;
        
        if stage == 0
            % Derivative with respect to aerosol backscatter ratio profile 
            % Theta. Use 'sparse' for diagonal matrices to improve performance.
            Jac1(eqnRange, Nbeg(j)-N0+1 : Nfin(j)-N0+1) = ...
                -sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                    (1 + etaACur) .* Lcalc .* coeff) * tauAMatrix{j} * ...
                    sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                    betaMCur .* Gamma) - ...
                sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                     betaCoeffDTheta * exp((1 + etaACur) * tauA) .* ...
                     betaM(Nbeg(j):Nfin(j), j) ./ betaMref(j) .* coeff);
            
            % Derivative with respect to aerosol lidar ratio profile Gamma.
            Jac1(eqnRange, (N-N0+1)+1) = ...
                -sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                    (1 + etaACur) .* betaMCur .* Lcalc .* coeff) * ...
                    tauAMatrix{j} * Theta(Nbeg(j)-N0+1:Nfin(j)-N0+1);

            % Derivative with respect to lidar correction coefficient R(j).
            Jac1(eqnRange, (N-N0+1)+1+j) = Lmeas(Nbeg(j):Nfin(j), j) .* coeff;

        else
            % Derivative with respect to aerosol backscatter ratio profile 
            % Theta. Use 'sparse' for diagonal matrices to improve performance.
            Jac1(eqnRange, Nbeg(j)-N0+1 : Nfin(j)-N0+1) = ...
                -sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                    (1 + etaACur) .* Lcalc .* coeff) * tauAMatrix{j} * ...
                    sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                    betaMCur .* Gamma(Nbeg(j)-N0+1:Nfin(j)-N0+1)) - ...
                sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                     betaCoeffDTheta * exp((1 + etaACur) * tauA) .* ...
                     betaM(Nbeg(j):Nfin(j), j) ./ betaMref(j) .* coeff);
            
            % Derivative with respect to aerosol lidar ratio profile Gamma.
            Jac1(eqnRange, (N-N0+1)+Nbeg(j)-N0+1 : (N-N0+1)+Nfin(j)-N0+1) = ...
                -sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                    (1 + etaACur) .* Lcalc .* coeff) * tauAMatrix{j} * ...
                    sparse(1:Nfin(j)-Nbeg(j)+1, 1:Nfin(j)-Nbeg(j)+1, ...
                    betaMCur .* Theta(Nbeg(j)-N0+1:Nfin(j)-N0+1));

            % Derivative with respect to lidar correction coefficient R(j).
            Jac1(eqnRange, (N-N0+1)*2+j) = Lmeas(Nbeg(j):Nfin(j), j) .* coeff;
        end
    end
    
    % Smoothness portion of the equations.
    % The matrices are sparse here, to improve performance.
    if stage == 0
        F2 = zeros((N-N0+1), 1);
        Jac2Matrix = cell(1, 1);
    else
        F2 = zeros((N-N0+1)*3, 1);
        Jac2Matrix = cell(3, 1);
    end
    
    F2(1 : (N-N0+1)) = d(1) * sqrt(hStep) * smoothOp2 * Theta;
    
    Jac2Matrix{1} = d(1) * sqrt(hStep) * smoothOp2;
    
    if stage == 1
        F2((N-N0+1)+1   : (N-N0+1)*2) = d(2) * sqrt(hStep) * smoothOp1 * Gamma;
        F2((N-N0+1)*2+1 : (N-N0+1)*3) = d(3) * sqrt(hStep) * ...
            (Gamma - gammaMiddle);

        Jac2Matrix{2} = d(2) * sqrt(hStep) * smoothOp1;
        Jac2Matrix{3} = d(3) * sqrt(hStep) * speye(N-N0+1);
    end

    for i = 1 : size(Jac2Matrix, 1)
        Psi2Arr(1, i, iter) = norm(F2((N-N0+1)*(i-1)+1 : (N-N0+1)*i));
    end
    
    if stage == 0
        Jac2 = [Jac2Matrix{1}, zeros((N-N0+1), 1+2)];
    else
        Jac2 = [blkdiag(Jac2Matrix{1:2}), zeros((N-N0+1)*2, 2);
            sparse(N-N0+1, N-N0+1), Jac2Matrix{3}, zeros((N-N0+1), 2)];
    end
    
    F = [F1; F2];
    Jac = [Jac1; Jac2];
end