clear; set(0,'defaultaxesfontsize',20); format long
%%% MCMC RWM algorithm for the logistic map
%%% based on Ex 2.4 in: Law, Stuart, Zygalakis (2015), "Data Assimilation"

%% setup

N=50;% number of steps
r=3;% dynamics determined by alpha (vary between r=1 and r=4, with smaller N for larger r)
gamma=0.1;% observational noise variance is gamma^2
C0=0.01;% prior initial condition varianceekk
m0=0.4;% prior initial condition mean (e.g. 0.4,0.5,0.7)
sd=10;rng(sd);% choose random number seed

%% truth

xt=0.1;xx(1)=xt;% truth initial condition
Jdet=1/2/C0*(xt-m0)^2;% background penalization
Phidet=0;% initialization model-data misfit functional
for j=1:N
    % can be replaced by Psi for each problem
    xx(j+1)=r*xx(j)*(1-xx(j));% create truth
    y(j)=xx(j+1)+gamma*randn;% create data
    Phidet=Phidet+1/2/gamma^2*(y(j)-xx(j+1))^2;% misfit functional
end
Idet=Jdet+Phidet;% compute log posterior of the truth

%% solution
% Markov Chain Monte Carlo: M forward steps of the
% Markov Chain on R (with truth initial condition)
M=1e5;% number of samples
X=zeros(M,1);% preallocate space to save time
beta=0.05;% step-size of random walker 
x=0;%xt;% truth initial condition (or else update I0) 
n=1; bb=0; rat(1)=0; 
while n<=M
    w=x+sqrt(2*beta)*randn;% propose sample from random walker
    xx(1)=w;
    Jdetprop=1/2/C0*(w-m0)^2;% background penalization
    Phidetprop=0;	
    for i=1:N
            xx(i+1)=r*xx(i)*(1-xx(i));
            Phidetprop=Phidetprop+1/2/gamma^2*(y(i)-xx(i+1))^2;
    end
    Idetprop=Jdetprop+Phidetprop;% compute log posterior of the proposal 	

    if rand<exp(Idet-Idetprop)% accept or reject proposed sample
        x=w; Idet=Idetprop; bb=bb+1;% update the Markov chain
    end
    rat(n)=bb/n;% running rate of acceptance
    X(n)=x;% store the chain
    n=n+1; 
end
dx=0.0005; x0=[0.01:dx:0.99];
Z=hist(X,x0);% construct the posterior histogram 
figure(1), plot(x0,Z/trapz(x0,Z),'k','Linewidth',2)% visualize the posterior
legend('MC approximation of posterior')

figure(2),plot(rat,'r','Linewidth',2),xlabel('MCMC steps'), ylabel('acceptance rate')
legend('running average')
