%
% This tests the use of multilevel Monte Carlo for a
% continuous-time Markov process
%
% The testcase is from the paper by David Anderson and Des Higham:
% 'Multi-level Monte Carlo for continuous time Markov chains,
%  with applications in biochemical kinetics',
% SIAM Multiscale Modelling and Simulation, 10(1):146-179, 2012.
%

function ctmp

close all; clear all;

addpath('..');

%
% call mlmc_test to perform MLMC tests
%

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

Lmin  = 2;      % minimum refinement level
Lmax  = 6;      % maximum refinement level
N0    = 100;    % initial number of samples on coarsest levels
Eps   = [ 1 2 5 10 20 ];  % desired accuracies for MLMC calcs

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

%
% plot results
%

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

if(nvert==1)
  figure(1)
  print('-deps2','ctmp_a.eps')
  figure(2)
  print('-deps2','ctmp_b.eps')
else
  print('-deps2','ctmp.eps')
end

end


%-------------------------------------------------------
%
% level l estimator for continuous-time Markov process with factor 4 refinement
%

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

nu     = [  1    0    0   -1    0 ; ...  % stoichiometric matrix
            0    1   -2    0   -1 ; ...
            0    0    1    0    0 ];

M = 4;   % refinement factor on each successive level
T = 1;   % simulation time interval

nf = 8 * M^l;  % 8 timesteps on coarsest level
nc = nf/M;

hf = T/nf;
hc = T/nc;

sums(1:6) = 0;

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

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

  lamf = propensity(qf);
  lamc = lamf;

  if l==0
    for n = 1:nf
      p1 = poissrnd(hf*lamf);
      qf = qf + nu*p1;
      qf = max(qf,0);
      lamf = propensity(qf);
    end
    
  else
    for n = 1:nc
      for m = 1:M
        lam_min  = min(lamf,lamc);
        lam_diff = abs(lamf-lamc);

        p1 = poissrnd(hf*lam_min);
        p2 = poissrnd(hf*lam_diff);

        qf = qf + nu*(p1 + (lamf>lamc).*p2);
        qc = qc + nu*(p1 + (lamc>lamf).*p2);
        
        qf = max(qf,0);
        lamf = propensity(qf);
      end
        
     qc = max(qc,0);
     lamc = propensity(qc);
    end
  end
  
  Pf = qf(3,:);
  Pc = qc(3,:);

  if l==0
    Pc = 0;
  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

cost = N*nf;   % cost defined as number of fine timesteps

end

%
%------ propensity function ---------
%

function lam = propensity(q)

lam = [ 25*ones(1,size(q,2));     ...
        1000*q(1,:);              ...
        0.001*q(2,:).*(q(2,:)-1); ...
        0.1*q(1,:);               ...
        q(2,:)                    ];
end