% Written by Vesa Kaarnioja, 2022
% Model solution to exercise 2, week 7

% Precompute the FEM data
level = 5;
[mass,grad,nodes,element,interior,centers,ncoord,nelem] = FEMdata(level);

% Set up the diffusion coefficient
s = 100; % stochastic dimension
decay = 2; % the "decay parameter" \vartheta from the slides
indices = (1:s)';
deterministic = indices.^(-decay) .* sin(pi*indices*centers(:,1)') .* sin(pi* indices * centers(:,2)');
a = @(y) 2+deterministic'*y;

% Use n == 2^maxiter as the "reference solution"
maxiter = 15;

% Initialize the parallel pool (uses the Parallel Computing Toolbox)
poolobj = gcp('nocreate');
if isempty(poolobj)
    parpool;
end

evec_sums = [];
evec_means = [];
eval_sums = [];
eval_means = [];
null = zeros(ncoord,1);

% Range over an increasing number of random samples
for ii = 0:maxiter % n = 2^ii, ii = 0,...,maxiter
    disp(['iteration: ',num2str(ii)]);
    nincr = max(1,2^(ii-1));
    n = 2^ii;
    tmp = zeros(ncoord,1);
    tmp2 = 0;
    % The main loop
    parfor k = 1:nincr
        mcnode = rand(s,1)-.5; % A realization of the parametric vector "y"
        stiffness = UpdateStiffness(grad,a(mcnode)); % Assemble the stiffness matrix corresponding to a(y), y == mcnode
        [evec,eval] = eigs(stiffness(interior,interior),mass(interior,interior),1,'smallestabs');
        sol = null;
        sol(interior) = sign(evec(1))*evec/sqrt(evec'*mass(interior,interior)*evec);
        tmp = tmp + sol; % Sum over Monte Carlo points
        tmp2 = tmp2 + eval;
    end
    evec_sums = [evec_sums,tmp]; % Store the sum
    evec_means = [evec_means,sum(evec_sums(:,1:size(evec_sums,2)),2)/n]; % Compute sample average (this is the Monte Carlo estimate)
    eval_sums = [eval_sums,tmp2];
    eval_means = [eval_means,sum(eval_sums(1:length(eval_sums)))/n];
end

% Use the solution corresponding to n = 2^maxiter as the reference solution
evec_ref = evec_means(:,end);
eval_ref = eval_means(end);

% Compute the L2 errors of solutions corresponding to n = 2^ii, 
% ii = 0,...,maxiter-1, vis-a-vis the reference solution
evec_errors = [];
eval_errors = [];
for ii = 1:size(evec_means,2)-1
    evec_errors = [evec_errors;2^(ii-1),sqrt((evec_means(:,ii)-evec_ref)'*mass*(evec_means(:,ii)-evec_ref))];
    eval_errors = [eval_errors;2^(ii-1),abs(eval_means(ii)-eval_ref)];
end

% Find the least squares fit for the errors
evec_lsq = [ones(size(evec_errors(:,1))),log(evec_errors(:,1))]\log(evec_errors(:,2));
evec_lsq(1) = exp(evec_lsq(1));
eval_lsq = [ones(size(eval_errors(:,1))),log(eval_errors(:,1))]\log(eval_errors(:,2));
eval_lsq(1) = exp(eval_lsq(1));

% Plot the absolute errors as a log-log plot
figure;
subplot(2,1,1);
loglog(eval_errors(:,1),eval_errors(:,2),'r.','MarkerSize',18), hold on
x = 2.^(0:maxiter-1);
loglog(x,eval_lsq(1)*x.^(eval_lsq(2)),'b','LineWidth',2), hold off
set(gca,'FontSize',14,'FontWeight','bold')
legend('data',['slope: ',num2str(eval_lsq(2))])
title(['Eigenvalues / Monte Carlo (s = ',num2str(s),')'])
xlabel('n')
ylabel('absolute error')

% Plot the L2 errors as a log-log plot.
subplot(2,1,2);
loglog(evec_errors(:,1),evec_errors(:,2),'r.','MarkerSize',18), hold on
x = 2.^(0:maxiter-1);
loglog(x,evec_lsq(1)*x.^(evec_lsq(2)),'b','LineWidth',2), hold off
set(gca,'FontSize',14,'FontWeight','bold')
legend('data',['slope: ',num2str(evec_lsq(2))])
title(['Eigenvectors / Monte Carlo (s = ',num2str(s),')'])
xlabel('n')
ylabel('L^2 error')


