%
% This tests the use of multilevel Monte Carlo for an SDE
% which requires adaptive time-stepping
%
% Developed with James Whittle, summer 2014
%
% M.B. Giles, C. Lester, J. Whittle.
% 'Non-nested adaptive timesteps in multilevel Monte Carlo computations'.
% Monte Carlo and Quasi-Monte Carlo Methods 2014, Springer, 2015.
%

function adapted

close all; clear all;
addpath('..');

%
% call mlmc_test to perform MLMC tests
%

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

N0    = 100;    % initial number of samples on coarsest levels
Lmin  = 2;      % minimum refinement level
Lmax  = 10;     % maximum refinement level

Eps   = [ 0.0005 0.001 0.002 0.005 0.01 ];  % desired accuracies for MLMC calcs

fp = fopen('adapted.txt','w');
mlmc_test(@sde_l, N,L, N0,Eps,Lmin,Lmax, fp);
fclose(fp);

nvert = 3;      % plotting option (1 for slides, 2 for papers, 3 for full set)
mlmc_plot('adapted', nvert);

if(nvert==1)
  figure(1)
  print('-deps2','adapted_a.eps')
  figure(2)
  print('-deps2','adapted_b.eps')
else
  print('-deps2','adapted.eps')
end

%-------------------------------------------------------
%
% level l estimator for SDE with factor 2 refinement
%

function [sums cost] = sde_l(l,N)

D = 2;   % dimension

beta = 4;

M = 2;
T = 1;

sums(1:6) = 0;
cost      = 0;

for N1 = 1:10000:N
  N2 = min(10000,N-N1+1);

  t  = zeros(1,N2);
  tf = t;
  tc = t;

  hf = zeros(1,N2);
  hc = hf;

  qf = zeros(D,N2);
  qc = qf;

  dWf = zeros(D,N2);
  dWc = dWf;

  if l==0

    while any(t<T)
      cost = cost + sum(t<T);

      t  = tf;
      dW = repmat(sqrt(hf),D,1).*randn(D,N2);

      qf = qf + drift(qf,beta).*repmat(hf,D,1) + 2*dW;
      qf = clamp(qf,beta);
      hf = min( dt(qf,beta,l), T-tf);
      tf = min(tf + hf, T);
    end

  else

    while any(t<T)
      cost = cost + sum(t<T);

      tm = t;
      t  = min(tf,tc);

      dW = repmat(sqrt(t-tm),D,1).*randn(D,N2);

      dWf = dWf + dW;
      dWc = dWc + dW;

      ind = find(tf==t);
      if size(ind)>0
        qf(:,ind) = qf(:,ind) + drift(qf(:,ind),beta).*repmat(hf(ind),D,1) + 2*dWf(:,ind);
        qf(:,ind) = clamp(qf(:,ind),beta);
        hf(ind) = min( dt(qf(:,ind),beta,l) , T-tf(ind));
        tf(ind) = min(tf(ind) + hf(ind), T);
        dWf(:,ind) = 0;
      end

      ind = find(tc==t);
      if size(ind)>0
        qc(:,ind) = qc(:,ind) + drift(qc(:,ind),beta).*repmat(hc(ind),D,1) + 2*dWc(:,ind);
        qc(:,ind) = clamp(qc(:,ind),beta);
        hc(ind) = min( dt(qc(:,ind),beta,l-1) , T-tc(ind));
        tc(ind) = min(tc(ind) + hc(ind), T);
        dWc(:,ind) = 0;
      end
    end
  end

  Pf = sum(qf.^2,1);
  if (l>0)
    Pc = sum(qc.^2,1);
  else
    Pc = zeros(1,N2);
  end

  sums(1) = sums(1) + sum(Pf-Pc);
  sums(2) = sums(2) + sum((Pf-Pc).^2);
  sums(3) = sums(3) + sum((Pf-Pc).^3);
  sums(4) = sums(4) + sum((Pf-Pc).^4);
  sums(5) = sums(5) + sum(Pf);
  sums(6) = sums(6) + sum(Pf.^2);
end

%
%------ drift function ---------
%

function v = drift(q,beta)
D = size(q,1);
v = - 4*beta * q ./ repmat(1-sum(q.^2,1),D,1);

%
%------ adaptive timestep ---------
%

function h = dt(q,beta,l)
h = 0.5^l * (1 - sqrt(sum(q.^2,1))).^2 / max(2*beta,50);

%
%------ clamp ---------
%

function q2 = clamp(q,beta)
D  = size(q,1);
q2 = q .* repmat( min(1, (1 - 1e-5)./(1e-100 + sqrt(sum(q.^2,1)))), D,1);
