clear; set(0,'defaultaxesfontsize',20); format long
%%% smoothing posterior of the deterministic 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 r (vary between 1 and 4, reduce N for larger values of r)
gamma=0.1;% observational noise variance is gamma^2
C0=0.01;% prior initial condition variance
m0=0.5;% prior initial condition mean
sd=1;rng(sd);% choose random number seed

%% truth

x=zeros(N+1,1); y=zeros(N,1);% preallocate space to save time
x(1)=0.1;% truth initial condition
for j=1:N
    % can be replaced by Psi for each problem
    x(j+1)=r*x(j)*(1-x(j));% create truth
    y(j)=x(j+1)+gamma*randn;% create data
end

%% solution

x0=[0.01:0.0005:0.99];% construct vector of different initial data
Phidet=zeros(length(x0),1);Idet=Phidet;Jdet=Phidet;% preallocate space
xx=zeros(N,1);% preallocate space 
% loop through initial conditions vv0, and compute log posterior I0(vv0) 
for j=1:length(x0) 
   xx(1)=x0(j); Jdet(j)=1/2/C0*(x0(j)-m0)^2;% background penalization 
     for i=1:N
        xx(i+1)=r*xx(i)*(1-xx(i)); 
        Phidet(j)=Phidet(j)+1/2/gamma^2*(y(i)-xx(i+1))^2;% misfit functional
     end 
    Idet(j)=Phidet(j)+Jdet(j);
end

logconstant=trapz(x0,Idet);% calculate a log constant shift 
constant=trapz(x0,exp(logconstant-Idet));% calculate normalizing constant
P=exp(logconstant-Idet)/constant;% normalize posterior distribution
prior=exp(-(x0-m0).^2/2/C0)/sqrt(2*pi*C0);% calculate prior distribution

figure(1),plot(x,'k','LineWidth',2)
hold on, plot(y,'r--','LineWidth',2), xlabel 'time',
legend('signal','observation') 
hold off

figure(2),plot(x0,prior,'k','LineWidth',2)
hold on, plot(x0,P,'c--','LineWidth',2), xlabel 'x_0',
legend('prior',['posterior (N=',num2str(N),')']) 
hold off

figure(3),plot(x0,Idet,'LineWidth',2), xlabel('x_0'), ylabel('I_{det}'), legend('potential')

