%
% This tests the use of multilevel Monte Carlo for an SDE
% which requires adaptive time-stepping
%
% Developed with James Whittle, summer 2014
%

function sde_adapted

close all; clear all;

%
% call mlmc_test to perform MLMC tests
%

nvert = 2;      % plotting option (1 for slides, 2 for papers, 3 for full set)
M     = 2;      % refinement cost factor (2^gamma in general MLMC Thm)

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

N0    = 25;    % initial number of samples on coarsest levels
Eps   = [ 0.0005 0.001 0.002 0.005 0.01 ];  % desired accuracies for MLMC calcs

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

subplot(nvert,2,2*nvert-1)
legend('\epsilon=0.0005','\epsilon=0.001','\epsilon=0.002','\epsilon=0.005','\epsilon=0.01',1)

subplot(nvert,2,2*nvert)
axis([5e-4 0.01 0.1 10])

print('-deps2','sde_adapted.eps')


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

function [sum1 sum2] = sde_l(l,N)

global beta D

D = 2;   % dimension

beta = 4;

M = 2;
T = 1;

sum1(1:4) = 0;
sum2(1:2) = 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)
      t  = tf;
      dW = repmat(sqrt(hf),D,1).*randn(D,N2);

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

  else

    while any(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)).*repmat(hf(ind),D,1) + 2*dWf(:,ind);
        qf(:,ind) = clamp(qf(:,ind));
        hf(ind) = min( dt(qf(:,ind),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)).*repmat(hc(ind),D,1) + 2*dWc(:,ind);
        qc(:,ind) = clamp(qc(:,ind));
        hc(ind) = min( dt(qc(:,ind),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

  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

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

function v = drift(q)

global beta D

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

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

function h = dt(q,l)

global beta D

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

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

function q2 = clamp(q)

global beta D

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