function varargout= Kimura2PM(w, tt_ratio)
% Kimura2PM  -  Kimura 2 parameter model
% Computes, for a specified edge length and transition-transversion ratio,
% the nucleotide transition matrix.
% Also, the derivative of this matrix with respect to w can be
% computed (optional).
% See, for instance, Durbin et al., "Biological sequence analysis", 
% p.196, where w= 4*beta*t and tt_ratio= alpha/beta.
%
% INPUT
% w 		->	edge length
% tt_ratio	-> 	transition-transversion ratio
%
% OUTPUT
% nuc_trans_prob(1:4,1:4)  ->  Matrix of nucleotide transition 
%			       probabilities.
% grad_nuc_trans_prob(1:4,1:4)  ->  Derivative of the matrix of
%				    nucleotide transition probabilities
%				    with respect to w [optional].
% I use the convention that nuc_trans_prob(i,j) is the
% probability that a nucleotide initially in state j mutates into
% state i:  P(j -->i).
% (Note that this is the transpose of the matrix you find in
% most textbooks, although for the Kimura model this distinction
% is unimportant as the transition matrix is symmetric).
%
% USAGE
% nuc_trans_prob= Kimura2PM(w, tt_ratio)
% OR
% [nuc_trans_prob, grad_nuc_trans_prob]= Kimura2PM(w, tt_ratio)

if nargin~=2
	error('Wrong number of input arguments')
end

if w<0 | tt_ratio<=0
	error('Parameters must be positive')
end

P_transversion= 1-exp(-w);
P_transversion= 0.25*P_transversion;
P_transition= 1+exp(-w)-2*exp(-0.5*(1+tt_ratio)*w);	
P_transition= 0.25*P_transition;
P_diagonal= 1-P_transition-2*P_transversion;


for i=1:4
	for j=1:4
		if i==j
			nuc_trans_prob(i,j)= P_diagonal;
		elseif mod((i+j),2)==0
			nuc_trans_prob(i,j)= P_transition;
		else
			nuc_trans_prob(i,j)= P_transversion;
		end
	end
end

varargout{1}=nuc_trans_prob;


% --- Optional computation of the derivative with respect to w ----
if nargout>1
	
grad_P_transversion= exp(-w);
grad_P_transversion= 0.25*grad_P_transversion;
grad_P_transition= (1+tt_ratio)*exp(-0.5*(1+tt_ratio)*w) - exp(-w);
grad_P_transition= 0.25*grad_P_transition;
grad_P_diagonal= -grad_P_transition-2*grad_P_transversion;

for i=1:4
	for j=1:4
		if i==j
			grad_nuc_trans_prob(i,j)= grad_P_diagonal;
		elseif mod((i+j),2)==0
			grad_nuc_trans_prob(i,j)= grad_P_transition;
		else
			grad_nuc_trans_prob(i,j)= grad_P_transversion;
		end
	end
end

varargout{2}=grad_nuc_trans_prob;

end