%% Exercise 3: Stochastic Collocation for PDE with random data
clear all; clc; close all;

%% (0) Preliminaries
% Set domain
L = 1; H = 1;

% Diffusion term
a = 'exp(0*x)';

% Random source term
f = @(xi) sprintf('-exp(%.16f) *  (1 + sin(pi * x )*%.16f + cos(pi * y)*%.16f)',xi(1),xi(2),xi(3));
    
% Random boundary data
gs = @(xi) sprintf('sin(%.16f + x + y)',xi);

% Set number of elements and generate mesh
nx = 64; ny=64;
[pt, el] = fem2d_lin_mesh(L,H,1+nx,1+ny);
x = pt(1,:)'; y = pt(2,:)';

% Plot certain realisations of u
for i=1:4,
    gs_sample = gs(rand(1));
    g = {gs_sample, gs_sample, gs_sample, gs_sample};

    u = fem2d_lin(L,H,1+nx,1+ny,a,f(randn(1,3)),g);
    subplot(2,2,i)
    fem2d_lin_plot(u,pt,el);
    title(sprintf('realization %i of u',i),'FontSize',22)
    set(gca,'FontSize',22)
end

%% (a) Compute mean field of u with 1 variational problem
% mean field of f
fmean = '-exp(1/2) *exp(0*x)'; 

% mean field of g
gsmean = 'cos(x+y) - cos(1+x+y)';
for i = 1:4, gmean{i} = gsmean; end

% mean field of u
umean = fem2d_lin(L,H,1+nx,1+ny,a,fmean,gmean);

% Plot mean of u
figure(); 
fem2d_lin_plot(umean, pt, el); xlabel('x'); ylabel('y');
title('mean of u','FontSize',22)
set(gca,'FontSize',22)

%% (b) Stochastic Collocation

% Number of stochastic variables
d = 4;

% And their distribution (uniform = legendre, normal = hermite)
s{1} = 'legendre'; for i=2:4, s{i} = 'hermite'; end;

tic;
for p=0:4,
    % Collocation points and weights
    [X,W] = FullGridSC(d, p*ones(1,d), s);
    nX = size(X,1);
    
    % Solve PDE at collocation points
    Us = zeros((1+nx)*(1+ny), nX);
    for m = 1:nX,
        % Evaluate g at collocation point m
        g = {gs(X(m,1)), gs(X(m,1)), gs(X(m,1)), gs(X(m,1))};
        
        % Solve PDE at collocation point
        Us(:,m) = fem2d_lin(L,H,1+nx,1+ny,a,f(X(m,2:4)),g);
    end
    
    % Compute Mean via quadrature weights
    umeanSC = Us*W;
    
    % Compute H1-error
    [L2_norm, H1_norm] = fem2d_lin_norm(umean-umeanSC, pt, el);
    err_SC(1+p) = H1_norm;
end
toc;

% Plot
figure();
loglog((1:5).^4,err_SC,'ko-','LineWidth',2)
title('Decay of SC integration error','FontSize',22)
set(gca,'FontSize',22)

%% (c) Monte Carlo
tic;
% Loop over 5 different MC runs for each M
for j = 1:5,
    % Loop over different Ms
    for p=0:4,
        M = (1+p)^4;
        Us = zeros((nx+1)*(ny+1), M);
        for m = 1:M,
            % Sample realizations of g
            xi1 = rand(1);
            g = {gs(xi1), gs(xi1), gs(xi1), gs(xi1)};
            
            % Solve PDE
            Us(:,m) = fem2d_lin(L,H,1+nx,1+ny,a,f(randn(3,1)),g);
        end
        % Compute mean by averaging
       umeanMC = mean(Us,2);
        
        % Compute H1-error
        [L2_norm, H1_norm] = fem2d_lin_norm(umean - umeanMC, pt, el);
        err_MC(1+p) = H1_norm;
    end
    Err_MC(j,:) = err_MC;
end
toc;

% Plot
hold on;
loglog((1:5).^4,Err_MC','+-','LineWidth',2)
hold off;
legend('SC','MC 1','MC 2','MC 3','MC 4','MC 5')
xlabel('Number of PDE solves')
ylabel('H_0^1-error in mean')
title('Convergence of mean estimate of MC and SC','FontSize',22)
set(gca,'FontSize',22)
