%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% mcdiffuse2d.m
%
% demonstration: Monte Carlo simulation of reaction-diffusion in 2D with importance sampling
%                Compute the probability G of a walker (starting from source)
%                reaching the domain boundary before being absorbed.
%
% differences from mcdiffuse2d.m :
%             1. simultaneously simulate all random walkers
%             2. importance function gives new transition probability (Knew, qnew)
%             3. allow walkers to duplicate if weight exceeds one
%
% Wei Cai
% for LLNL CCMS Summer Institute Lectures Aug.1-3, 2005
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

NGRID = 40; %number of grid points in both directions
r = 0.05;   %absorption rate inside domain
NWALKER = 1000;
NROUNDS = 10;

x=[1:NGRID]; y=[1:NGRID];
xx=ones(NGRID,1)*[1:NGRID];
yy=[1:NGRID]'*ones(1,NGRID);

xs=ceil(NGRID*0.25);
ys=ceil(NGRID*0.4);

%absorption function
q=max((abs(xx/NGRID-0.5)>0.42)|(abs(yy/NGRID-0.5)>0.35), r);

%source function
s=zeros(NGRID,NGRID);
s(ys,xs)=1;

%gain function (at domain boundary)
g=(q==1);

%importance function
if(~exist('mu'))
   mu=0.74; 
end
if(~exist('alpha'))
   alpha=0.03;
end
%Imp=ones(NGRID,NGRID);
%Imp=min(exp((sqrt((xx-xs).^2+(yy-ys).^2)-a)*mu),1); %Imp(ys,xs)=Imp(ys+1,xs)*0.85;
Imp=exp(-mu*exp(-alpha*((xx-xs).^2+(yy-ys).^2)));
figure(3);
mesh(log10(Imp));
xlabel('x');
ylabel('y');
zlabel('log_{10} ( I )');


S0=sum(sum(Imp.*s));
gnew=S0*g./Imp;

%compute new transition probability in importance sampling
K=ones(NGRID,NGRID,4)/4;   %old transition probability matrix
for i=1:NGRID,
   for j=1:NGRID,
      K(j,i,:)=K(j,i,:)*(1-q(j,i));
   end
end

Knew=zeros(size(K));
for i=1:NGRID,
   for j=1:NGRID,
      ip=mod(i,NGRID)+1;   %assuming periodic boundary condition
      im=mod(i-2,NGRID)+1;
      jp=mod(j,NGRID)+1;
      jm=mod(j-2,NGRID)+1;
      Knew(j,i,1)=K(j,i,1)*Imp(j,ip)/Imp(j,i);
      Knew(j,i,2)=K(j,i,2)*Imp(j,im)/Imp(j,i);
      Knew(j,i,3)=K(j,i,3)*Imp(jp,i)/Imp(j,i);
      Knew(j,i,4)=K(j,i,4)*Imp(jm,i)/Imp(j,i);      
   end
end

Nnew=ones(NGRID,NGRID);    %normalization factor
for i=1:NGRID,
   for j=1:NGRID,
      Nnew(j,i)=sum(Knew(j,i,:));
   end
end
if(max(max(Nnew))>2)
   disp(sprintf('max(Nnew)=%e > 2 (not allowed!) - abort',max(max(Nnew))));
   return;
end
qnew=max(1-Nnew,0);
      
figure(1);
p1=contour(x,y,q,[0.5 0.5]*(r+1));
axis equal
hold on
plot(xs,ys,'ro');
hold off
xlim([0 NGRID]);
ylim([0 NGRID]);
xlabel('x');
ylabel('y');

figure(4);
mesh(x,y,Nnew);

%Start simulation

%number of visits
dr = [1 0; -1 0; 0 1; 0 -1];

G=zeros(NROUNDS,1);
for iter=1:NROUNDS,
   nvisit=zeros(NGRID,NGRID);
   %simultaenously simualte all walkers
   POS=ones(NWALKER,1)*[xs,ys]; %position of all walkers
   nwalker=length(POS(:,1));
   while nwalker>0,
    m=1;
    while m<=nwalker,
      pos=POS(m,:);              %position of current walker
      nvisit(pos(2),pos(1))=nvisit(pos(2),pos(1))+1;
      termprob = qnew(pos(2),pos(1));
      copyprob = Nnew(pos(2),pos(1))-1;
      if(rand<termprob)
         %path terminated, delete current walker (m)
         POS(m,:)=POS(nwalker,:);
         nwalker=nwalker-1;
         if(nwalker>0)
            POS=POS(1:nwalker,:);
         else
            POS=[];
         end
      else
         probs=zeros(1,4);
         probs(:)=Knew(pos(2),pos(1),:);
         probs=cumsum(probs);
         xi=rand*probs(4);
         id=sum(probs<xi)+1;
         POS(m,:)=pos+dr(id,:);
         
         if(copyprob>0)
            if(rand<copyprob)
               %insert a new walker at the end of list
               nwalker=nwalker+1;
               xi=rand*probs(4);
               id=sum(probs<xi)+1;
               POS(nwalker,:)=pos+dr(id,:);
            end
         end         
      end      
      m=m+1;
    end
   end
   G(iter)=sum(sum(nvisit.*gnew))/NWALKER;
   if(G(iter)>1)
      break;
   end
end
G
disp(sprintf('G = %e +/- %e  (mu=%e)',mean(G),std(G)/sqrt(NROUNDS),mu));

figure(1);
hold on
contour(nvisit/NWALKER,[0:0.01:0.2])
hold off

fs=18;
set(gca,'FontSize',fs);
xlabel('x');
ylabel('y');

figure(2);
mesh(nvisit/NWALKER);

set(gca,'FontSize',fs);
xlabel('x');
ylabel('y');
zlabel('\psi');

