%% SENSITIVITY: COMPARISON OF MC- VS. PCE- AND LRA-BASED SOBOL' INDICES
%
% In this example, Sobol' sensitivity indices for the 
% <https://uqworld.org/t/borehole-function/ borehole function>
% are calculated with three different methods:
% Monte Carlo (MC) simulation, polynomial chaos expansion (PCE), 
% and canonical low-rank approximation (LRA).

%% 1 - INITIALIZE UQLAB
%
% Clear variables from the workspace,
% set random number generator for reproducible results,
% and initialize the UQLab framework:
clearvars
rng(100,'twister')
uqlab

%% 2 - COMPUTATIONAL MODEL
%
% and initialize the UQLab framework:and initialize the UQLab framework:
%
% Create a MODEL object from the function file:
ModelOpts.mFile = 'uq_ABC';

myModel = uq_createModel(ModelOpts);

%%
% Type |help uq_borehole| for information on the model structure as well as
% the description of each variable.

%% 3 - PROBABILISTIC INPUT MODEL
%
% The probabilistic input model consists of eight independent random 
% variables.
%
% Specify the marginals as follows:
InputOpts.Marginals(1).Name = 'X1';
InputOpts.Marginals(1).Type = 'uniform';
InputOpts.Marginals(1).Parameters = [0 0.0226];
InputOpts.Marginals(2).Name = 'X2';
InputOpts.Marginals(2).Type = 'uniform';
InputOpts.Marginals(2).Parameters = [0 0.0456];
InputOpts.Marginals(3).Name = 'X3';
InputOpts.Marginals(3).Type = 'uniform';
InputOpts.Marginals(3).Parameters = [0 0.0209];
InputOpts.Marginals(4).Name = 'X4';
InputOpts.Marginals(4).Type = 'uniform';
InputOpts.Marginals(4).Parameters = [0 3.6101E-5];
InputOpts.Marginals(5).Name = 'X5';
InputOpts.Marginals(5).Type = 'uniform';
InputOpts.Marginals(5).Parameters = [0 2.0549E-5];
InputOpts.Marginals(6).Name = 'X6';
InputOpts.Marginals(6).Type = 'uniform';
InputOpts.Marginals(6).Parameters = [0 4.8597E-5];
InputOpts.Marginals(7).Name = 'X7';
InputOpts.Marginals(7).Type = 'uniform';
InputOpts.Marginals(7).Parameters = [0 0.0323];
InputOpts.Marginals(8).Name = 'X8';
InputOpts.Marginals(8).Type = 'uniform';
InputOpts.Marginals(8).Parameters = [0 0.0314];
InputOpts.Marginals(9).Name = 'X9';
InputOpts.Marginals(9).Type = 'uniform';
InputOpts.Marginals(9).Parameters = [0 0.0209];
InputOpts.Marginals(10).Name = 'X10';
InputOpts.Marginals(10).Type = 'uniform';
InputOpts.Marginals(10).Parameters = [0 3.4571E-5];
InputOpts.Marginals(11).Name = 'X11';
InputOpts.Marginals(11).Type = 'uniform';
InputOpts.Marginals(11).Parameters = [0 2.8437E-5];
InputOpts.Marginals(12).Name = 'X12';
InputOpts.Marginals(12).Type = 'uniform';
InputOpts.Marginals(12).Parameters = [0 2.8991E-5];
InputOpts.Marginals(13).Name = 'X13';
InputOpts.Marginals(13).Type = 'uniform';
InputOpts.Marginals(13).Parameters = [0 0.0105];
InputOpts.Marginals(14).Name = 'X14';
InputOpts.Marginals(14).Type = 'uniform';
InputOpts.Marginals(14).Parameters = [0 0.0218];
InputOpts.Marginals(15).Name = 'X15';
InputOpts.Marginals(15).Type = 'uniform';
InputOpts.Marginals(15).Parameters = [0 0.0103];
InputOpts.Marginals(16).Name = 'X16';
InputOpts.Marginals(16).Type = 'uniform';
InputOpts.Marginals(16).Parameters = [0 2.8560E-5];
InputOpts.Marginals(17).Name = 'X17';
InputOpts.Marginals(17).Type = 'uniform';
InputOpts.Marginals(17).Parameters = [0 2.2815E-5];
InputOpts.Marginals(18).Name = 'X18';
InputOpts.Marginals(18).Type = 'uniform';
InputOpts.Marginals(18).Parameters = [0 3.5323E-5];
InputOpts.Marginals(19).Name = 'X19';
InputOpts.Marginals(19).Type = 'uniform';
InputOpts.Marginals(19).Parameters = [0 0.0485];
InputOpts.Marginals(20).Name = 'X20';
InputOpts.Marginals(20).Type = 'uniform';
InputOpts.Marginals(20).Parameters = [0 0.0231];
InputOpts.Marginals(21).Name = 'X21';
InputOpts.Marginals(21).Type = 'uniform';
InputOpts.Marginals(21).Parameters = [0 0.0599];
InputOpts.Marginals(22).Name = 'X22';
InputOpts.Marginals(22).Type = 'uniform';
InputOpts.Marginals(22).Parameters = [0 2.9510E-5];
InputOpts.Marginals(23).Name = 'X23';
InputOpts.Marginals(23).Type = 'uniform';
InputOpts.Marginals(23).Parameters = [0 2.9021E-5];
InputOpts.Marginals(24).Name = 'X24';
InputOpts.Marginals(24).Type = 'uniform';
InputOpts.Marginals(24).Parameters = [0 2.1696E-5]; 
%%
% Create an INPUT object based on the specified marginals:

myInput = uq_createInput(InputOpts);

%% 4 - SENSITIVITY ANALYSIS
%
% Sobol' indices are calculated first with a direct MC simulation 
% of the model and subsequently through post-processing of the
% coefficients of its PCE and LRA approximation.

%% 4.1 MC-based Sobol' indices
%
% Select the sensitivity analysis module in UQLab
% and the Sobol' analysis method:
SobolOpts.Type = 'Sensitivity';
SobolOpts.Method = 'Sobol';

%%
% Specify the maximum order of the Sobol' indices to be calculated:
SobolOpts.Sobol.Order = 1;

%%
% Specify the sample size for the MC simulation:

SobolOpts.Sobol.SampleSize = 100000;
%%
% Note that the total cost of computation is $(M+2) \times N$,
% where $M$ is the input dimension and $N$ is the sample size.
% Therefore, the total cost for the current setup is
% $(8+2) \times 10^5 = 10^6$ evaluations of the full computational model.

%%
% Run the sensitivity analysis:
mySobolAnalysisMC = uq_createAnalysis(SobolOpts);

%%
% Retrieve the analysis results for comparison:
mySobolResultsMC = mySobolAnalysisMC.Results;


%% 4.2 PCE-based Sobol' indices
%
% Select the metamodeling tool in UQLab
% and the polynomial chaos expansion (PCE) type:
PCEOpts.Type = 'Metamodel';
PCEOpts.MetaType = 'PCE';

%%
% Assign the full computational model:
PCEOpts.FullModel = myModel;
%%
% The full model is used to generate an experimental design for the PCE. 

%%
% Specify the maximum polynomial degree (*default*: sparse PCE):
PCEOpts.Degree = 5;

%%
% Specify the size of the experimental design
% (i.e., the total computational cost of constructing the metamodel):
PCEOpts.ExpDesign.NSamples = 200;

%%
% Calculate the PCE:
myPCE = uq_createModel(PCEOpts);

%%
% The same options structure |SobolOpts| can be re-used
% to create a new analysis on the PCE model:
mySobolAnalysisPCE = uq_createAnalysis(SobolOpts);

%%
% Retrieve the results for comparison:
mySobolResultsPCE = mySobolAnalysisPCE.Results;

%%
% *Note*:
% The PCE metamodel does not need to be explicitly specified as
% the last created model will be used in the analysis.
% The sensitivity analysis module in UQLab recognizes PCE models and 
% calculates the Sobol' indices accordingly by post-processing the
% coefficients of a PCE model without sampling.

%% 4.3 LRA-based Sobol' indices
%
% Select the metamodeling tool in UQLab
% and the canonical low-rank approximation (LRA) type:
LRAOpts.Type = 'Metamodel';
LRAOpts.MetaType = 'LRA';

%%
% Assign the full computational model:
LRAOpts.FullModel = myModel;

%%
% Specify the rank range:
LRAOpts.Rank = 1:50;

%%
% Specify the degree range:
LRAOpts.Degree = 1:50;

%%
% Specify the size of the experimental design
% (i.e., the total computation cost of the metamodel):

LRAOpts.ExpDesign.NSamples = 200;

%% 
% Calculate the LRA:
myLRA = uq_createModel(LRAOpts);

%%
% As before, the same options structure |SobolOpts| can be re-used
% to create a new analysis on the LRA model:
mySobolAnalysisLRA = uq_createAnalysis(SobolOpts);


%%
% Retrieve the results for comparison:
mySobolResultsLRA = mySobolAnalysisLRA.Results;

%%
% *Note*:
% The LRA metamodel does not need to be explicitly specified as
% the last created model will be used in the analysis.
% The sensitivity analysis module in UQLab recognizes LRA models and 
% calculates the Sobol' indices accordingly by post-processing the
% coefficients of a LRA model without sampling.

%% 5 - COMPARISON OF THE RESULTS
%
% Print the results of the Sobol' indices calculation
% based on the MC simulation:
uq_print(mySobolAnalysisMC)

%%
% Print the results of the Sobol' indices calculation based on PCE:
uq_print(mySobolAnalysisPCE)

%%
% Print the results of the Sobol' indices calculation based on LRA:
uq_print(mySobolAnalysisLRA)


%%
% Compile the relevant results for comparison:
SobolTotal = [...
    mySobolResultsMC.Total ...
    mySobolResultsPCE.Total ...
    mySobolResultsLRA.Total];
SobolFirstOrder = [...
    mySobolResultsMC.FirstOrder ...
    mySobolResultsPCE.FirstOrder ...
    mySobolResultsLRA.FirstOrder];

%% 
% Create a bar plot to compare the total Sobol' indices:

% Create the plot
uq_figure('Name', 'Total Sobol'' Indices')
barWidth = 1;
uq_bar(1:24, SobolTotal, barWidth)
% Set axes limits
ylim([0 1])
xlim([0 25])
% Set labels
xlabel('MAs')
ylabel('Total Sobol'' indices')
set(...
    gca,...
    'XTick', 1:length(InputOpts.Marginals),...
    'XTickLabel', mySobolResultsPCE.VariableNames)
% Set legend
uq_legend({...
    sprintf('MC-based (%.0E simulations)', mySobolResultsMC.Cost),...
    sprintf('PCE-based (%d simulations)', myPCE.ExpDesign.NSamples),...
    sprintf('LRA-based (%d simulations)', myLRA.ExpDesign.NSamples)},...
    'Location', 'northeast')

%% 
% Create a bar plot to compare the first-order Sobol' indices:

% Create the plot
uq_figure('Name', 'First-order Sobol'' Indices')
uq_bar(1:24, SobolFirstOrder, barWidth)
% Set axes limits
xlim([0 25])
ylim([0 1])
% Set labels
xlabel('MAs')
ylabel('First-order Sobol'' indices')
set(...
    gca,...
    'XTick', 1:length(InputOpts.Marginals),...
    'XTickLabel', mySobolResultsPCE.VariableNames)
% Set legend
uq_legend({...
    sprintf('MC-based (%.0E simulations)', mySobolResultsMC.Cost),...
    sprintf('PCE-based (%d simulations)', myPCE.ExpDesign.NSamples),...
    sprintf('LRA-based (%d simulations)', myLRA.ExpDesign.NSamples)},...
    'Location', 'northeast')
