%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% metro2d_constr_K.m
%
% demonstration: Construct transition probability matrix K(x->y)
%                for 2D metropolis MC simulation in a potential landscape.
%                find exact solution of psi(x), G 
%                and optimal importance function J(x)
%
% Wei Cai
% for LLNL CCMS Summer Institute Lectures Aug.1-3, 2005
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Niter=1000;

T=500; %K kB=8.6174e-5; %eV/K kT=kB*T; %eV
kB=8.6174e-5; %eV/K
kT=kB*T; %eV

x=[-1.5:0.1:1.5]; y=x; 
NGRID=length(x);

xx=ones(size(y'))*x;
yy=y'*ones(size(x));

%construct energy landscape
Z=pot2d(xx,yy)+yy*0.02;

%construct transition probability matrix K
dr = [1 0; -1 0; 0 1; 0 -1];
K=zeros(NGRID,NGRID,4);   %old transition probability matrix
for i=1:NGRID,
   for j=1:NGRID,
      if(i<NGRID)
         K(j,i,1)=exp( (Z(j,i)-Z(j,i+1))/kT );
      end
      if(i>1)
         K(j,i,2)=exp( (Z(j,i)-Z(j,i-1))/kT );
      end
      if(j<NGRID)
         K(j,i,3)=exp( (Z(j,i)-Z(j+1,i))/kT );
      end
      if(j>1)
         K(j,i,4)=exp( (Z(j,i)-Z(j-1,i))/kT );
      end
      K(j,i,:)=min(K(j,i,:),1);
      K(j,i,:)=K(j,i,:)/sum(K(j,i,:));
   end
end

%source function
s=zeros(NGRID,NGRID);
s(16,6)=K(16,5,1);
s(16,4)=K(16,5,2);
s(17,5)=K(16,5,3);
s(15,5)=K(16,5,4);

%gain function
g=zeros(NGRID,NGRID);
g(16,27)=1;

%absorption function
q=zeros(NGRID,NGRID);
q(16,5)=1; q(16,27)=1;
for i=1:NGRID,
   for j=1:NGRID,
      K(j,i,:)=K(j,i,:)*(1-q(j,i));
   end
end

%solve for psi(x,y)
psi=s;
for iter=1:Niter,
   psiK=zeros(size(psi));
   psiK1=psiK; psiK1(:,:)=psi.*K(:,:,1);
   psiK2=psiK; psiK2(:,:)=psi.*K(:,:,2);
   psiK3=psiK; psiK3(:,:)=psi.*K(:,:,3);
   psiK4=psiK; psiK4(:,:)=psi.*K(:,:,4);
   psiK=psiK1(:,[NGRID,1:NGRID-1])+psiK2(:,[2:NGRID,1]) ...
       +psiK3([NGRID,1:NGRID-1],:)+psiK4([2:NGRID,1],:) ;
   psi=s+psiK;
end

G=sum(sum(psi.*g));
disp(sprintf('exact solution G=psi*g= %e',G));

%solve for J(x,y)
J=g;
for iter=1:Niter,
   KJ=zeros(size(J));
   KJ1=KJ; KJ1(:,:)=K(:,:,1).*J(:,[2:NGRID,1]);
   KJ2=KJ; KJ2(:,:)=K(:,:,2).*J(:,[NGRID,1:NGRID-1]);
   KJ3=KJ; KJ3(:,:)=K(:,:,3).*J([2:NGRID,1],:);
   KJ4=KJ; KJ4(:,:)=K(:,:,4).*J([NGRID,1:NGRID-1],:);
   KJ=KJ1+KJ2+KJ3+KJ4;
   J=KJ+g;
end

G1=sum(sum(s.*J));
disp(sprintf('exact solution G=  s*J= %e',G1));


%plot results
figure(1);
contour(x,y,Z,50);
axis equal
hold on
plot(x(5),y(16),'bo',x(27),y(16),'rx');
contour(x,y,log10(psi),[0:0.01:0.2]);
hold off
xlim([min(x), max(x)]);
ylim([min(y), max(y)]);
xlabel('x');
ylabel('y');

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

figure(2);
mesh(x,y,max(log10(psi),-13));
zlim([-15, 0]);
set(gca,'FontSize',fs);
xlabel('x');
ylabel('y');
zlabel('\psi');

figure(3);
mesh(x,y,min(Z,2.1));
zlim([-1,3]);
set(gca,'FontSize',fs);
xlabel('x');
ylabel('y');
zlabel('E');

figure(4);
mesh(x,y,log10(J));

set(gca,'FontSize',fs);
xlabel('x');
ylabel('y');
zlabel('log_{10} J');


