% Written by Vesa Kaarnioja, 2022
% Model solution to exercise 4, 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

% Use an off-the-shelf generating vector downloaded from
% https://web.maths.unsw.edu.au/~fkuo/lattice/index.html
z = load('~/lattice-39101-1024-1048576.3600');
z = z(1:s,2);

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

% Range over an increasing number of cubature nodes
for ii = 10:maxiter % n = 2^ii
    disp(['iteration: ',num2str(ii)]);
    if ii == 10
        n = 2^ii;
        n2 = n;
        ind = 0:(n-1);
    else
        n = 2^(ii-1);
        n2 = 2^ii;
        ind = 1:2:(n2-1);
    end
    tmp = zeros(ncoord,1);
    tmp2 = 0;
    % The main loop
    parfor k = 1:n
        qmcnode = mod(z*ind(k)/n2,1)-.5; % QMC node
        stiffness = UpdateStiffness(grad,a(qmcnode)); % Assemble the stiffness matrix corresponding to a(y), y == qmcnode
        [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;
        tmp2 = tmp2 + eval; % Sum over quasi-Monte Carlo points
    end
    evec_sums = [evec_sums,tmp]; % Store the sum
    evec_means = [evec_means,sum(evec_sums(:,1:size(evec_sums,2)),2)/n2]; % Compute the mean (this is the quasi-Monte Carlo estimate)
    eval_sums = [eval_sums,tmp2];
    eval_means = [eval_means,sum(eval_sums(1:length(eval_sums)))/n2];
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 = 10,...,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^(9+ii),sqrt((evec_means(:,ii)-evec_ref)'*mass*(evec_means(:,ii)-evec_ref))];
    eval_errors = [eval_errors;2^(9+ii),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.^(10: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 / quasi-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.^(10: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 / quasi-Monte Carlo (s = ',num2str(s),')'])
xlabel('n')
ylabel('L^2 error')



