/* This is an example HY-PHY Batch File.

	In this file, we will illustrate how to define discrete category
	variables to model rate heterogeneity.
	
   Sergei L. Kosakovsky Pond and Spencer V. Muse 
   June 2001. 
*/

/* this will display the category info, along with sample mean and variance */
#include "displayFunction.bf";


/* 1. We begin with the simplest possible category variable:

	it has two rates: 1 and 2 with probability .5 each
	
	The first argument in defining a category variable is the number of rate classes.
	
	The second argument defines the probability of each class. We set it to
	EQUAL.
	
	The third argument is only meaningful for continuously sampled distributions,
	and is essentially ignored.
	
	The fourth argument is the density of the distribution that we are sampling,
	which is again not applicable here.
	
	The fifth argument is the CDF of the distribution that we are sampling,
	which in this case consists of the explicit enumeration of rates.
	
	The sixth and seventh arguments are the bounds of the support of the distribution.
	([1,2]) in this case. For discrete distributions, the bounds are not important.

	The rest of the arguments aren't used for this simple case and are omitted.
*/

category catVar = (2,EQUAL,MEAN, ,{{1}{2}}, 1,2);

fprintf (stdout,"1). 2 fixed equiprobable rates\n");

GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


/* 
   2. We modify the previous example so that rate "1" has probability 1/3 and
	  rate "2" has probability 2/3. To that end we simply replace the second
	  argument with the explicit matrix of frequencies.
*/

category catVar = (2,{{1/3}{2/3}} ,MEAN, ,{{1}{2}}, 1,2);

fprintf (stdout,"2). 2 fixed rates with different probabilities\n");

GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


/* 
   3. Next category variable has 3 equiprobable rates which depend on the parameter R, which
   and maintain the ratios 1:4:9. 
*/

global R;

category catVar = (3,EQUAL,MEAN, ,{{R}{4*R}{9*R}}, 0, 1e25);

R = 1;
fprintf (stdout,"3). 3 proportional equiprobable rates\n R = 1.\n");

GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


R = 5;
fprintf (stdout,"\n R = 5. \n");

GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


/* 
   4. The next example prompts the user to enter the number of categories and
   the ratios between categories.
*/

fprintf (stdout,"4). User defined discrete distribution\n\n Enter the number of categories:");
fscanf (stdin,"Number",rateCount);

if (rateCount<=1)
{
	fprintf (stdout, "\n\n", rateCount, " is an invalid number of categories");
	return;
}

/* define the matrix for the ratios */

multiplierMatrix = {rateCount,1};

fprintf (stdout,"\n");

for (k=0; k<rateCount; k=k+1)
{
	fprintf (stdout, "Multiplier for category ", k+1,":");
	fscanf (stdin,"Number",M);
	multiplierMatrix[k][0]:=M__*R;
	/* the use of M__ forces HyPhy to substitute the numeric value of M into the expression.
	   If we didn't use it, all matrix entries would hold the same expression : "M*R"
	 */
}

category catVar = (rateCount,EQUAL,MEAN, ,multiplierMatrix, 0, 1e25);

R=1;
fprintf (stdout,"\n R = 1. \n");
GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


R=2;
fprintf (stdout,"\n R = 2. \n");
GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


/* 
   5. Here is a way to define the general discrete distribution with 3 rates.
   Note how we avoid the use of P1+P2+P3=1 constraint be an alternative 
   parametrization.
*/

global P  = 1/6;
global P1 = 1/3;

/* constrain the probabilities to be no more than 1 */

P:<1;
P1:<1;

/* probability matrix */

probMatrix = {{P}
			  {(1-P)P1}
			  {(1-P)(1-P1)}};
			  
/* rate matrix is parametrized in such a way that rate1<=rate2<=rate3.
   It is just a matter of preference though. */
   
global  M1 = 2;
global  M2 = 1.5;  
   
M1:>1;
M2:>1;  

rateMatrix = {{R}
			  {R*M1}
			  {R*M1*M2}};
			  
category catVar = (3,probMatrix,MEAN, ,rateMatrix, 0, 1e25);

R=1;
fprintf (stdout,"5). General Discrete Distribution with 3 rate classes.\n");

GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


/* 
   6. Truncated ([0,9]) Poisson with parameter lambda. P{X=k} = lambda^k/k! Exp(-lambda)
*/

global lambda = .9;

rateMatrix = {10,1};
for (k=0; k<10; k=k+1)
{
	rateMatrix[k][0]:=k__;
}

probMatrix = {10,1}; 

probMatrix[0][0]:= Exp(-lambda);
factorial = 1;
for (k=1; k<10; k=k+1)
{
	factorial = factorial*k;
	probMatrix[k][0]:=lambda^(k__)/factorial__*Exp(-lambda);
}

/* HYPHY will automatically normalize the probability matrix to sum to 1 
if it doesn't already */

category catVar = (10,probMatrix,MEAN, ,rateMatrix, 0, 1e25);

fprintf (stdout,"6). Truncated Poisson with 10 rate classes (lambda = .9).\n");
GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);


/* 
   7. Truncated geometric with parameter p. P{X=k} = (1-p)p^k
*/

fprintf (stdout,"7). Truncated geometric (P=.25) with user specified number of classes.\n\n Enter the number of categories:");
fscanf (stdin,"Number",rateCount);

if (rateCount<=1)
{
	fprintf (stdout, "\n\n", rateCount, " is an invalid number of categories");
	return;
}

global P = .25;
P:<1;

rateMatrix = {rateCount,1};
for (k=0; k<rateCount; k=k+1)
{
	rateMatrix[k][0]:=k__;
}
probMatrix = {rateCount,1}; 
probMatrix[0][0]:= 1-P;
for (k=1; k<rateCount; k=k+1)
{
	probMatrix[k][0]:=(1-P)P^(k__);
}

/* HYPHY will automatically normalize the probability matrix to sum to 1 
if it doesn't already */

category catVar = (rateCount,probMatrix,MEAN, ,rateMatrix, 0, 1e25);

GetInformation (catInfo,catVar);
catInfo = echoCatVar (catInfo);