%
% Tests on reflected diffusion, based on testcase in 
% "Symmetrized Euler scheme ...." by Bossy, Gobet and Talay
%

function reflected

close all; clear all;

addpath('..');

%
% call mlmc_test to perform MLMC tests
%

for pass = 0:2
  nvert = 3;      % plotting option (1 for slides, 2 for papers, 3 for full set)

  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

  if pass==0
    name = 'reflected_ref';
    fprintf(1,'\n*** Tests using reflection scheme ***\n');
    L = 4;      % levels for convergence tests 
    N = 100000; % samples for convergence tests
  elseif pass==1 
    name = 'reflected_proj';
    fprintf(1,'\n*** Tests using projection scheme ***\n');
    L = 8;      % levels for convergence tests 
    N = 10000;  % samples for convergence tests
  else
    name = 'reflected_pen';
    fprintf(1,'\n*** Tests using penalisation scheme ***\n');
    L = 8;      % levels for convergence tests 
    N = 10000;  % samples for convergence tests
  end

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

  mlmc_plot(name, nvert);

  if(nvert==1)
    figure(1)
    print('-deps2',[name '_a.eps'])
    figure(2)
    print('-deps2',[name '_b.eps'])
  else
    print('-deps2',[name '.eps'])
  end

end

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

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

D = 3;   % dimension

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;

  xf = zeros(D,N2);
  xc = xf;

  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);

      xf = xf + drift_dt(xf) + dW;
      xf = reflect(xf, D,pass);

      hf = min( dt(xf,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
        xf(:,ind) = xf(:,ind) + drift_dt(xf(:,ind)) + dWf(:,ind);
        xf(:,ind) = reflect(xf(:,ind), D,pass);

        hf(ind) = min( dt(xf(:,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
        xc(:,ind) = xc(:,ind) + drift_dt(xc(:,ind)) + dWc(:,ind);
        xc(:,ind) = reflect(xc(:,ind), D,pass);

        hc(ind) = min( dt(xc(:,ind),l-1) , T-tc(ind));
        tc(ind) = min(tc(ind) + hc(ind), T);
        dWc(:,ind) = 0;
      end
    end
  end

  Pf = sum(xf.^2,1);
  if (l>0)
    Pc = sum(xc.^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_dt(x,h)

D = size(x,1);
r = max(1, sqrt(sum(x.^2,1)));
v = repmat((1-r)./r,D,1).*x;

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

function h = dt(x,l)

r = sqrt(sum(x.^2,1));
d = max(0, 1-r);
h = min(0.5^(l+3), max(0.25^(l+3), (d/4).^2));

%
%------ reflect at r=1 ---------
%

function x = reflect(x, D,pass)

if pass==2
  return
end

r2 = sum(x.^2,1);

if (max(r2)<=1)
  return
end

i   = find(r2>1);

if pass==0
  fac = 2./sqrt(r2(i)) - 1;   % reflection
else
  fac = 1./sqrt(r2(i));       % projection
end

x(:,i) = repmat(fac,D,1) .* x(:,i);
