%
% These are some MLMC tests for research with Lukas Szpruch and Arnulf Jentzen
%
% SPDE is du = u_{xx} dt + dW
% subject to u(x,0)=0, u(0,t)=u(1,t)=0
%
% dW represents additive white noise
%

function white_noise

close all; clear all;

nvert = 3;   % nvert=1 for slides; nvert=2 for papers

N  = 1000;   % samples for convergence tests
L  = 6;      % levels for convergence tests 

M  = 8;      % refinement cost factor
N0 = 50;     % initial samples on coarsest levels

file = 'white_noise';
Eps  = [ 0.1 0.2 0.5 1 ];

figs = mlmc_test(@spde_l, M, N,L, N0,Eps, nvert);

figure(figs(1))

if(nvert==1)
  print('-deps2',[file 'a.eps'])
  figure(figs(2))
end

subplot(nvert,2,2*nvert-1);
axis([0 6 5 1e6])
legend('\epsilon=0.1','\epsilon=0.2','\epsilon=0.5','\epsilon=1',1)

subplot(nvert,2,2*nvert); 
axis([0.1 1.0 1e4 1e9])

if(nvert==1)
  print('-deps2',[file 'b.eps'])
else
  print('-deps2',[file '.eps'])
end

%
% Code to solve SPDE using a multilevel approach
%

function [sum1, sum2] = spde_l(l,M)

%
% grid parameters
%

X = 1.0;
T = 1.0;

J0 = 2;
N0 = 10;

J  = J0 * 2^l;
N  = N0 * 4^l;

xf  = linspace(0,X,J+1)';
dxf = X / J;
dtf = T / N;
jf  = 2:J;

xc  = linspace(0,X,J/2+1)';
dxc = 2*dxf;
dtc = 4*dtf;
jc  = 2:J/2;

%
% main loop, doing M2 at a time
%

sum1 = zeros(1,4);
sum2 = zeros(1,2);
M2   = 10;     % max number of samples at a time

for m = 1:M2:M
  m2 = min(M2,M-m+1);

  vf = zeros(J+1,m2);
  vc = zeros(J/2+1,m2);

  dMc = 0;

  for n = 1:N
    dMf = sqrt(dtf/dxf)*randn(J,m2);
    dMc = dMc + 0.5*(dMf(1:2:end,:) + dMf(2:2:end,:));

    d2vdx2f = vf(jf+1,:) - 2*vf(jf,:) + vf(jf-1,:);
    dMf     = 0.5*(dMf(jf-1,:) + dMf(jf,:));
    vf(jf,:) = vf(jf,:) + (dtf/dxf^2)*d2vdx2f + dMf;

    if (mod(n,4)==0)
      d2vdx2c  = vc(jc+1,:) - 2*vc(jc,:) + vc(jc-1,:);
      dMc      = 0.5*(dMc(jc-1,:) + dMc(jc,:));
      vc(jc,:) = vc(jc,:) + (dtc/dxc^2)*d2vdx2c + dMc;
      dMc = 0;
    end
  end

%
% calculate L2 norm as output
%
  Pf = 100 * sum(vf.^2,1)*dxf;
  if l==0
    Pc = 0;
  else
    Pc = 100 * sum(vc.^2,1)*dxc;
  end

  sum1(1) = sum1(1) + sum(Pf-Pc);
  sum1(2) = sum1(2) + sum((Pf-Pc).^2);
  sum1(3) = sum1(3) + sum((Pf-Pc).^3);
  sum1(4) = sum1(4) + sum((Pf-Pc).^4);
  sum2(1) = sum2(1) + sum(Pf);
  sum2(2) = sum2(2) + sum(Pf.^2);
end
