%1D FSP, Brian munsky, qbss 2019
%Naomi Ziv

%simple birth/death model

%loop on size of time (only done after working through some examples) 
tarr=logspace(-2,4,100);
for j=1:length(tarr);

    %N is the size of the approximation
    %100 is pretty big, error will be small
    N=100;
    %50 is small, error will be big
    %N=50;
    %K is the birth rate
    K=5*ones(1,N+1);
    %G is the death rate
    G=0.1*(0:N);
    
    %intialise A matix 
    A=zeros(N+1);
    %intialise B, the last column
    B=zeros(1,N+1);
    
    %loop on col/row of A
    for n=0:N
        
        %fill diagonal, you leave state with birth and death
        A(n+1,n+1)=-K(n+1)-G(n+1);
        
        if n<N %not last step
            %off diagonal, birth leads to +1 state
            A(n+2,n+1)=K(n+1);
        else
            %B is beyond approximation
            B(1,n+1)=K(n+1);
        end
        
        if n>0 %not first step
            %off diagonal, death leads to -1 state
            A(n,n+1)=G(n+1);
        end
        
    end
    
    %make big matrix
    A=[A,zeros(N+1,1);...
        B,0];
    
    %%%%%%%%%%%%%%%%%%%
    %intial condidtions:
    X0=3;
    %X0=0;
    
    %first example for P0, gives poission for X0=0 not X0=3, at t=5
    P0=zeros(N+2,1);
    P0(X0+1)=1;
    
    %second example for P0, poission distribution
    %P0=poisspdf([0:N].',X0);
    %P0(end+1)=0;
    
    %time
    %t=5;
    %better time scale for capturing steady state, even with non-poission
    %intial conditions, steady state will be poission
    %t=50;
    
    %time when looping on time
    t=tarr(j);
    %%%%%%%%%%%%%%%%%%%
    
    %solve
    Pt=expm(A*t)*P0;
    
    %ploting probabilities (used for specfic examples but also nice across
    %iterations if you look at all distributions on same plot)
    %plot([0:N],Pt(1:end-1))
    
    %save error for each iteration on time
    C(j)=Pt(end);
end

%this plot is for example of looping on time durations
%sould be cumalative distribution (to really see, have N=50) 
%it is distribution of "first passage times"
%plot(tarr,C)
plot(log(tarr),C)


%these calculations are for demonstration during specfic examples

%calculate moments
%mean
mu=[0:N]*Pt(1:N+1)
%variance
mux2=[0:N].^2*Pt(1:N+1);
s2=mux2-mu^2

%calculate analytically (5, 0.1 are parameters of model)
Z=(5/0.1)*(1-exp(-0.1*t))
%calculate analytically when intial conditions are distribution
Z2=(5/0.1)*(1-exp(-0.1*t))+exp(-0.1*t)*X0

%compare to theoretical distribution (you should have the plot of
%probabilties up)
Q=poisspdf([0:N],Z);
hold on;
plot([0:N],Q,'--')

%error, prob leaving approximation
Pt(end)