%MATLAB CODE FOR
% AGE-FROM-STAGE ANALYSIS IN MARKOVIAN ENVIRONMENTS
% FOR ROSTOCK AGE DETERMINATION WORKSHOP OCTOBER 2007
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% based on modification of troptrees implementation of code
%  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
% 
%(initial coding was july05 to implement analytical stuff from Tulja's
%     notes 14june05)
%  CAUTION: equation numbers need to be updated to match 
%  Tuljapurkar and Horvitz 2006. Ecology 87: 1497-1509
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear all
close all

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  MODULE 1
%
% in this module, we should 
%  
%  1) load the demographic and habitat trans matrices
%  2) specify the structures
%  3) correct for perfect one-period survival
%  4) choose the subsets of demographic matrices to be analyzed
%  5) choose the habitat dynamics matrix
%  6) get the dimensions of the set you have chosen
%demographic matrices must be in an array
%habitat transition must be in a matrx
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%name the species, set up
names=['Ardisia       '];
codes=['AE'];
   

%Markov chain habitat transition matrix follows 
%import from numeric text file
load ardesc_cmat.txt
cmat=ardesc_cmat
  
sp=1;                             %here there is only 1 sp and 1 hab dynamics
name=names(sp,:);                 %finds the name of the species
code=codes(sp,:);                 %find the two letter code of the species
matfilename=(name);               %uses species name to name output file
nstates=size(cmat,1);             %S in the notes

maxage=300; % pick a max age to which you will estimate survivorship


%import the qs from numeric text files and make into array
for i=1:nstates
eval(['qs(:,:,i)=load (''ardesc_q' num2str(i) '.txt'');' ]);
end
                   
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% letters used correspond to Tulja's 14june05 notes
%qs=qs;                     %the q matrices in an array, one for each state, in the notes

%get the dimension of the sets you have actually chosen
nstates=size(cmat,1);        %K in the paper
nstages=size(qs(:,:,1),1);   %S in the paper

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% output variables of interest from module 1 
%         
%         maxage
%         name, code matfilename
%         nstates, nstages, qs, cmat
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% MODULE 2 
%  
% ANALYZE habitat dynamics matrix
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%get the stationary habitat state distribution, pi in the notes and rho
[wmat,dmat]=eig(cmat);%asks for eigenvectors and eigenvalues
[lmat,j]=max(max(dmat));%gets the largest eigenvalue and its location,j
wmat=wmat(:,j);%gets the eigenvector in the jth location of wmat
pis=wmat./sum(wmat);%rescales the eigenvector to sum to 1

%get rho, the serial autocorrelation %%%%
d=diag(dmat); % the eigenvalues
absd=abs(d);   % gets their absolute magnitudes
[y,index]=sort(absd); %sort by absolute magnitude and get index for each 
rho=(d(index(nstates-1))); %finds the subdominant, which is in position nstates-1 = serial correlation
hab_eigs=d; %to save the eigenvalue spectrum of the habitat process

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% output variables of interest from module 2
%   
%         pis, rho, hab_eigs
%         
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
%MODULE 3  
%Deterministic analyses
%note: this is the original version of this module
%      not the speed version which has a bug
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% The individual environmental states as constant evironments
%%%%        constant environment survivorship, mortality 

ages=1:maxage; %a vector from one to max age
Nzero=zeros(nstages,1);
Nzero(1,1)=1;

qcon_surv=zeros(nstages,maxage,nstates);
for i=1:nstates
 for x=1:maxage %for individual environments
    qcon_surv(:,x,i)=sum(qs(:,:,i)^x)'; %all stages, each row is a stage
 end
end


%set up arrays to receive output
qcon_lx=zeros(nstates,maxage);
qcon_lxplus1=zeros(nstates,maxage-1);
qcon_mu=zeros(nstates,maxage-1);

for i=1:nstates
 qcon_lx(i,:)=qcon_surv(1,:,i); %take the stage 1 survivorships only
 qcon_lxplus1(i,:)=qcon_lx(i,2:maxage);
 qcon_mu(i,:) = log(qcon_lx(i,1:maxage-1)) - log(qcon_lxplus1(i,:));
 qcon_plateau(i)=-log(max(eig(qs(:,:,i))));
 qcon_oneperiodsurvival(i,:)=sum(qs(:,:,i)); %stage-specific survival
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%The Average q: the so-called "deterministic" case
% q's weighted by their frequency in habitat

qbar=zeros(nstages);
for i=1:nstates;
qbar=qbar+pis(i)*qs(:,:,i); 
end;      

%%%% "deterministic" survivorship, mortality, 
%%%%        and stage structure at each age %%%%%%%%%%%%%

ages=1:maxage; %a vector from one to max age
Nzero=zeros(nstages,1);
Nzero(1,1)=1;

for x=1:maxage 
    qbar_surv(x,:)= sum(qbar^x);
    qbar_trans(:,x)= qbar^x*Nzero/sum(qbar^x*Nzero); %stage structure at each age
end
qbar_lx=qbar_surv(:,1)'; %survivorship to age x, conditional on starting in stage 1,
                  % the other columns are conditional on starting in stages
                  % 2, 3, 4 etc
                  
qbar_lxplus1=qbar_lx(2:maxage);
qbar_mu = log(qbar_lx(1:maxage-1)) - log(qbar_lxplus1);

qbar_plateau=-log(max(eig(qbar)));

qbar_oneperiodsurvival=sum(qbar); %stage-specific survival

% distribution of age at death
qbar_deaths = qbar_lx(1:maxage-1)-qbar_lxplus1;
qbar_deaths=qbar_deaths/sum(qbar_deaths);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% output variables of interest from module 3
%    ages                 % a vector of ages from 1 to maxage
%    
%    qbar_lx
%    qbar_lxplus1
%    qbar_mu
%    qbar_plateau
%    qbar_deaths
%    qbar_oneperiodsurvival
%    qbar_trans           % the stage sructure at each age, transient dynamics
%
%    qcon_lx
%    qcon_lxplus1
%    qcon_mu
%    qcon_plateau
%    qcon_deaths
%    qcon_oneperiodsurvival
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%MODULE 4  
%  MEGAMATRIX FOR FUNDAMENTAL MATRIX FOR DOUBLY MARKOV PROCESS
% note: this is the speed version
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
ckron = kron(cmat, eye(nstages));
%eye(nstages) is an identity matrix, S X S
%ckron is a matrix we have defined that has dimension SK x SK
% the Kronecker Product of the c matrix with an S X S identity matrix.
% ckron looks like: 
%      it is subdivided into K X K blocks
%       each of which is the same size as the q matrices (S X S); 
%       in each block, all are zeroes, except the diagonal,
%       which has the same value throughout the block, 
%             one particular entry of the c-matrix
%  SO...  in a case of 4 environmental states and 8 stages 
%     define an 8 x8 matrix that is zeros everywhere except the diagonal
%         which is one c-matrix parameter
% 
%C11=[c11 0   0   0   0   0   0   0 
%       0   c11 0   0   0   0   0   0 
%       0   0   c11 0   0   0   0   0 
%       0   0   0   c11 0   0   0   0 
%       0   0   0   0   c11 0   0   0 
%       0   0   0   0   0   c11 0   0 
%       0   0   0   0   0   0   c11 0
%       0   0   0   0   0   0   0   c11 ]
%
% then 
%ckron = [C11 C12 C13 C14 
%         C21 C22 C23 C24  
%         C31 C32 C33 C34  
%         C41 C42 C43 C44] 
%        
% where each of these is an 8x8 matrix, as is C11

%the following loop creates a large matrix 
% that is also of the full SK x SK dimension

%OLD slow way:
% for i = 1:nstates
% mat1 = qs(:,:,i); %choose the ith q matrix
% rowindex1=nstages*(i-1)+1; rowindex2=nstages*i; %gets you into the block
% matrix2(rowindex1:rowindex2,rowindex1:rowindex2) = mat1;
% end

%NEW speedy way: making a matrix with the qs on the diagonal
%this diag matrix is just "matrix2" in the original program
matrix2 = sparse(nstages*nstates, nstages*nstates);%%prellocate for speed
indxvec = nstages.*(0: nstates); %this marks the row & col locations of the blocks
for i=1:nstates
    matrix2(indxvec(i)+1:indxvec(i+1), indxvec(i)+1:indxvec(i+1)) = qs(:,:,i);;
end

%this matrix2 contains within it all the q matrices;
%     is also subdivided into S x S blocks, all of these blocks are
%      zeros everywhere except the blocks along the diagonal..
%        and these are the q matrices, one for each state 

% an alternative way to get this matrix for our example
%    follows...it might be more intuitive to some
% of course one would have to enter it by hand 
%    if one changed dimensions of either demography or environment, 
%    whereas the above loop is more general

%z8=zeros(8)
%matrix2a=[m1 z8 z8 z8 
%          z8 m2 z8 z8  
%          z8 z8 m3 z8 
%          z8 z8 z8 m4] 
%    

m = matrix2 * ckron;           %the m matrix, 
                               % eqn 29 in Tuljapurkar and Horvitz 2006
                               %   there is a typo in the paper 
                               %    the second c21 should be c22
clear ckron

% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% chekm=(qs(:,:,2)*cmat(2,1))-m(nstages+1:2*nstages,1:nstages); %
% %this should be exactly 0's, it is
% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Itilda=eye(nstages*nstates);    %see Table 1, SK by SK identity matrix

% OLD: slow way
% eatildas=zeros(nstates*nstages,nstages,nstates); % set up array to receive eatildas
% for i=1:nstates
% eatildas(:,:,i)=Itilda(:,(i-1)*nstages+1:nstages*i);
% end;                           %array of eatilidas, eqn 25
%                                % all rows, columns read from Itilda 1-8, 9-16,
%                                %                                    17-24, etc

qatildas=zeros(nstates*nstages,nstages,nstates); % set up array to receive qatildas
for i=1:nstates
qatildas(:,:,i)=matrix2(:,(i-1)*nstages+1:nstages*i);
end;                            %see Table 1, array of qatilda's
                                % each one is an SK x S matrix [0 0 0..qa ..0 0]'

etilda=zeros(nstates*nstages,nstages); %see Table 1, etilda, only one not an array
                                       % SK x S matrix [I I I...I I]'
for i=1:nstates
etilda((i-1)*nstages+1:nstages*i,:)=eye(nstages);
end;                            %etilda

I=eye(nstages);                 %identity matrix of dimension S x S

Ns_markov=zeros(nstages,nstages,nstates); %set up for array of Ns
%junk = ((Itilda-m)')\etilda;
for i=1:nstates;
Ns_markov(:,:,i)=I + etilda'*(Itilda-m)^-1*qatildas(:,:,i); %eqn 30 in Tulja and Carol 2006
                                                            %conditional on initial state 1
%Ns_markov(:,:,i)=I + qatildas(:,:,i)'*junk;
end;

clear matrix2 Itilda
%%%%average over all initial states %%%%%

Nbar_markov=zeros(nstages);
for i=1:nstates;
Nbar_markov=  Nbar_markov+pis(i)*Ns_markov(:,:,i); 
end;                              % 
                                  % fundamental matrix by expected frequency 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%lifeexp, column sums
lifeexp_markov=zeros(nstates,nstages); %set up array
for i=1:nstates
lifeexp_markov(i,:)=sum(Ns_markov(:,:,i));
end;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% output variables of interest from module 4
%
%       m
%       Ns_markov
%       Nbar_markov
%       lifeexp_markov
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% MODULE 5  the variance  of the FUNDAMENTAL MATRIX
% not included in the brief workshop at Rostock
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% (design a save statement for the stuff you want)
save AE_base name code qs nstates nstages cmat maxage ages qcon* qbar*
save AE_markov name code m Ns_markov Nbar_markov lifeexp_markov
save AE_forThatsonly name code m etilda qatildas nstages nstates maxage ages

clear all
load AE_forThatsonly

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Survivorship and hazard rates
% MODULE 6
% For workshop: OLD WAY is used as it corresponds to the format in the paper
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

ages=1:maxage; %a vector from one to max age
Thats=zeros(nstages,nstages,nstates,maxage);

for a=1:nstates;
    
for i=1:maxage;
       Thats(:,:,a,i)= etilda'*m^(i-1)*qatildas(:,:,a); %eqn 34 in Tulja and Carol 2006
    end;

end;   


surv_allstages=sum(Thats,1);%survivorship of each stage to each age
%chek=surv_allstages(:,:,2,10)-sum(Thats(:,:,2,10))
surv_1=surv_allstages(1,1,:,:); %get the element for those that start in stage 1
surv=squeeze(surv_1); %age-specific survivorships 
                        %rows are the envstates and the columns are the ages
                        %so to reference surv for envstate 1...surv(1,:)
                        %                     for envstate 2...surv(2,:)

shift_surv=surv(:,2:size(surv,2)); %lxplus1
mu = log(shift_surv) - log(surv(:,1:size(shift_surv,2))); % see next step, too!
mu_markov = -mu;
        %mu for state 1 ...mu(1,:)
        %mu for state 2 ...mu(2,:)
surv_markov = surv;      
 
% distribution of age at death
deaths = surv(:,1:size(shift_surv,2))-shift_surv;
sumdeath=sum(deaths,2); %sum for each row
for i=1:nstates
    deaths_markov(i,:)=deaths(i,:)/sumdeath(i);
end;   %to divide each element by the corresponding row sum

 %%%%     hazard  
 %%%%     and "quasistationary distribtuion" 
lams_markov=zeros(1,1);
slope_markov=zeros(1,1);
us_markov=zeros(nstages,nstates);
vs_markov=zeros(nstages,nstates);
zs_markov=zeros(1,nstates); 

 
[wmat,dmat]=eig(m);%asks for eigenvectors and eigenvalues of megamatrix
[lmat,j]=max(max(dmat));%gets the largest eigenvalue and its location,j
wmat=wmat(:,j);%gets the eigenvector in the jth location of wmat
u=wmat./sum(wmat);%rescales the eigenvector to sum to 1
lams_markov=lmat;

[vmat,dmat]=eig(m');%asks for eigenvectors and eigenvalues
[lmat,j]=max(max(dmat));%gets the largest eigenvalue and its location,j
vmat=vmat(:,j);%gets the eigenvector in the jth location of vmat
v=vmat./sum(vmat);%rescales the eigenvector to sum to 1

us_markov=reshape(u,nstages,nstates);      
u_markov=sum(us_markov,2);%row sums
vs_markov=reshape(v,nstages,nstates);

slope_markov = -log(lams_markov); %%%%% - slope of log(survivorship) and zs
plateau_markov = slope_markov;



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% output variables of interest from module 6
%    
%    Thats
%    surv_markov
%    shift_surv
%    mu_markov
%    lams_markov
%    plateau_markov
%    deaths_markov
%    us_markov
%    vs_markov
%    zs_markov
%    
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
save AE_survmarkov name code m nstages nstates surv_markov mu_markov plateau_markov deaths_markov us_markov vs_markov


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% % MODULE 7
% % SURVIVORSHIP VARIANCE
%  not presented at Rostock workshop
% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% NO MODULE 8 
% figures are now made in a separate file  
%   figure making loop
% this code is standardized for up to 12 environmental states, any number of stages
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%