client = jcmoptimizer.Client(); 

% Definition of the search domain
design_space = { ...
  struct('name', 'b1', 'type', 'continuous', 'domain', [0,10]), ... 
  struct('name', 'b2', 'type', 'continuous', 'domain', [0.1,4]), ...
  struct('name', 'b3', 'type', 'continuous', 'domain', [-4,-0.1]), ...
  struct('name', 'b4', 'type', 'continuous', 'domain', [0.05,1]), ...
  struct('name', 'b5', 'type', 'continuous', 'domain', [0.05,1]) ...
};

 % Creation of the study object with study_id 'sensitivity_analysis'
study = client.create_study( ...
    'design_space', design_space, ...
    'driver','ActiveLearning',...
    'study_name','Variance-based sensitivity analysis for parameter reconstruction',...
    'study_id', 'sensitivity_analysis');
% The vectorial model function of the MGH17 problem
function val = model(x)
    s = 0:32;
    val = x(1) + x(2)*exp(-s.*x(4)) + x(3)*exp(-s.*x(5));
end

%variable definition for average of model vector
avg = struct();
avg.type = "LinearCombination";
avg.name = "model_average";
avg.inputs = {"model_vector"};
 
study.configure( ...
    'max_iter', 100, ...
    'surrogates', { ...
        ...A multi-output Gaussian process that learns the dependence of
        ...the model on the design parameters.
        struct('type', "GP", 'name', "model_vector", 'output_dim', 33, ...
             'correlate_outputs', false)
    }, ...
    'variables', {avg}, ...
    'objectives', { ...
        ...The objective is to sample the model at positions of maximal
        ...uncertainty of the model average.
        struct('type', "Explorer", 'variable', "model_average", ...
             'penalize_boundaries', true, 'min_uncertainty', 1e-3) ...
    } ...
);

% Evaluation of the black-box function for specified design parameters
function observation = evaluate(study, sample)

    pause(2); % make objective expensive
    observation = study.new_observation();
    %array of design values
    x = [sample.b1, sample.b2, sample.b3, sample.b4, sample.b5];    
    observation.add(model(x));
    
end  

% Run the minimization
study.set_evaluator(@evaluate);
study.run(); 

sobol_indices = study.driver.get_sobol_indices( ...
    'object_type', "surrogate", ...
    'name', "model_vector", ...
    'max_uncertainty', 0.001 ...
)
variances = sobol_indices.variance;
sobol_values = sobol_indices.first_order;


fig = figure('Position', [0, 0, 1000, 500]);

subplot(2, 1, 1)
for i = 1:length(design_space)
    plot(sobol_values(i,:), '-o', 'DisplayName', design_space{i}.name);
    hold on
end
ylabel("First-order Sobol' index")
legend()
grid()

subplot(2, 1, 2)    
for i = 1:length(design_space)
    var_p = (design_space{i}.domain(2)-design_space{i}.domain(1))^2/12;
    scaled_var = variances.*sobol_values(i,:)/var_p;
    plot(scaled_var, '-o', 'DisplayName', design_space{i}.name);
end
legend()
xlabel("Model vector index")
ylabel("Scaled variance")
grid()
saveas(fig, "variance_sobol.png")


client.shutdown_server();