<commit_before>package dr.app.beagle.tools;

import java.util.ArrayList;
import java.util.List;

import dr.app.beagle.evomodel.sitemodel.EpochBranchSubstitutionModel;
import dr.app.beagle.evomodel.sitemodel.GammaSiteRateModel;
import dr.app.beagle.evomodel.substmodel.FrequencyModel;
import dr.app.beagle.evomodel.substmodel.HKY;
import dr.app.beagle.evomodel.substmodel.SubstitutionModel;
import dr.app.seqgen.SequenceSimulator;
import dr.evolution.alignment.Alignment;
import dr.evolution.alignment.SimpleAlignment;
import dr.evolution.datatype.Nucleotides;
import dr.evolution.io.NewickImporter;
import dr.evolution.sequence.Sequence;
import dr.evolution.tree.NodeRef;
import dr.evolution.tree.Tree;
import dr.evomodel.branchratemodel.BranchRateModel;
import dr.evomodel.branchratemodel.DefaultBranchRateModel;
import dr.evomodel.sitemodel.SiteModel;
import dr.inference.model.Model;
import dr.inference.model.Parameter;
import dr.math.MathUtils;
import dr.xml.AbstractXMLObjectParser;
import dr.xml.AttributeRule;
import dr.xml.ElementRule;
import dr.xml.XMLObject;
import dr.xml.XMLObjectParser;
import dr.xml.XMLParseException;
import dr.xml.XMLSyntaxRule;

/**
 * Class for performing random sequence generation for a given site model.
 * Sequences for the leave nodes in the tree are returned as an alignment.
 */

public class EpochBranchSubstitutionModelSimulator {

	/** nr of samples to generate **/
	private int nReplications;
	/** tree used for generating samples **/
	private Tree tree;
//	private List<SubstitutionModel> substModelList;
//	private List<FrequencyModel> frequencyModelList;
//	private Parameter epochTimes;
	/** site model used for generating samples **/
	private GammaSiteRateModel siteModel;
	/** branch rate model used for generating samples **/
	private BranchRateModel branchRateModel;
	/** nr of categories in site model **/
	private int categoryCount;
	/** nr of states in site model **/
	private int stateCount;

	private static boolean has_ancestralSequence = false;
	private Sequence ancestralSequence;

	/** an array used to transfer transition probabilities **/
	private double[][] transitionProbabilities;

	
	private EpochBranchSubstitutionModel epochModel;
	
	/**
	 * Constructor
	 * 
	 * @param tree
	 * @param siteModel
	 * @param branchRateModel
	 * @param sequenceLength: nr of sites to generate
	 */
//	EpochBranchSubstitutionModelSimulator(Tree tree,
//			EpochBranchSubstitutionModel epochModel,
//			List<SubstitutionModel> substModelList,
////			List<FrequencyModel> frequencyModelList, 
//			Parameter epochTimes,
//			GammaSiteRateModel siteModel, 
//			BranchRateModel branchRateModel,
//			int sequenceLength) {
//
//		this.tree = tree;
//		this.epochModel = epochModel;
//		this.substModelList = substModelList;
////		this.frequencyModelList = frequencyModelList;
//		this.epochTimes = epochTimes;
//		this.siteModel = siteModel;
//		this.branchRateModel = branchRateModel;
//		this.nReplications = sequenceLength;
//		this.stateCount = substModelList.get(0).getDataType().getStateCount();
//		this.categoryCount = siteModel.getCategoryCount();
//		this.transitionProbabilities = new double[categoryCount][stateCount * stateCount];
//	} // END: Constructor
	
	//suggestion: use this type of constructor
	EpochBranchSubstitutionModelSimulator(Tree tree, EpochBranchSubstitutionModel epochModel, GammaSiteRateModel siteModel, int sequenceLength) {
		this.tree = tree;
		this.siteModel = siteModel;
		this.nReplications = sequenceLength;
		//don't really need a siteModel to get the stateCount
		//too bad the siteModel only functions as a container with respect to the substitutionModel
		this.stateCount = epochModel.getSubstitutionModel(0,0).getDataType().getStateCount();
		//don't like this categoryCount though, does require a siteModel
		this.categoryCount = siteModel.getCategoryCount();
		this.transitionProbabilities = new double[categoryCount][stateCount * stateCount];
	} // END: Constructor
	
	/**
	 * Convert integer representation of sequence into a Sequence
	 * 
	 * @param seq: integer representation of the sequence
	 * @param node: used to determine taxon for sequence
	 * @return Sequence
	 */
	Sequence intArray2Sequence(int[] seq, NodeRef node) {
		
		StringBuilder sSeq = new StringBuilder();
		
		for (int i = 0; i < nReplications; i++) {
			sSeq.append(siteModel.getSubstitutionModel().getDataType().getCode(seq[i]));
		}
		
		return new Sequence(tree.getNodeTaxon(node), sSeq.toString());
	} // END: intArray2Sequence

	void setAncestralSequence(Sequence seq) {
		ancestralSequence = seq;
		has_ancestralSequence = true;
	}

	int[] sequence2intArray(Sequence seq) {

		if (seq.getLength() != nReplications) {
			throw new RuntimeException("Ancestral sequence length has "
					+ seq.getLength() + " characters " + "expecting "
					+ nReplications + " characters");
		}

		int array[] = new int[nReplications];
		for (int i = 0; i < nReplications; i++) {
			array[i] = siteModel.getSubstitutionModel().getDataType().getState(
					seq.getChar(i));
		}
		
		return array;
	}// END: sequence2intArray

	/**
	 * perform the actual sequence generation
	 * 
	 * @return alignment containing randomly generated sequences for the nodes
	 *         in the leaves of the tree
	 */
	public Alignment simulate() {

		NodeRef root = tree.getRoot();

		double[] categoryProbs = siteModel.getCategoryProportions();
		int[] category = new int[nReplications];
		for (int i = 0; i < nReplications; i++) {
			category[i] = MathUtils.randomChoicePDF(categoryProbs);
		}

		int[] seq = new int[nReplications];

		if (has_ancestralSequence) {

			seq = sequence2intArray(ancestralSequence);

		} else {

			FrequencyModel frequencyModel = siteModel.getSubstitutionModel()
					.getFrequencyModel();
			for (int i = 0; i < nReplications; i++) {
				seq[i] = MathUtils.randomChoicePDF(frequencyModel
						.getFrequencies());
			}

		}// END: ancestral sequence check

		SimpleAlignment alignment = new SimpleAlignment();
		alignment.setReportCountStatistics(true);
		alignment.setDataType(siteModel.getSubstitutionModel().getDataType());

		traverse(root, seq, category, alignment);

		return alignment;
	} // END: simulate

	// TODO: traverse for Epoch model
	void traverse(NodeRef node, int[] parentSequence, int[] category,
			SimpleAlignment alignment) {

		for (int iChild = 0; iChild < tree.getChildCount(node); iChild++) {

			NodeRef child = tree.getChild(node, iChild);
			for (int i = 0; i < categoryCount; i++) {
				getTransitionProbabilities(tree, child, i, transitionProbabilities[i]);
			}

			int[] seq = new int[nReplications];
			double[] cProb = new double[stateCount];
			for (int i = 0; i < nReplications; i++) {
				System.arraycopy(transitionProbabilities[category[i]],
						parentSequence[i] * stateCount, cProb, 0, stateCount);<commit_msg>Parsing from XML should be fine (in theory). Now we need to fix the getTransitionProbabilities method ...<commit_after>package dr.app.beagle.tools;

import java.util.ArrayList;
import java.util.List;

import dr.app.beagle.evomodel.sitemodel.EpochBranchSubstitutionModel;
import dr.app.beagle.evomodel.sitemodel.GammaSiteRateModel;
import dr.app.beagle.evomodel.substmodel.FrequencyModel;
import dr.app.beagle.evomodel.substmodel.HKY;
import dr.app.beagle.evomodel.substmodel.SubstitutionModel;
import dr.app.seqgen.SequenceSimulator;
import dr.evolution.alignment.Alignment;
import dr.evolution.alignment.SimpleAlignment;
import dr.evolution.datatype.Nucleotides;
import dr.evolution.io.NewickImporter;
import dr.evolution.sequence.Sequence;
import dr.evolution.tree.NodeRef;
import dr.evolution.tree.Tree;
import dr.evomodel.branchratemodel.BranchRateModel;
import dr.evomodel.branchratemodel.DefaultBranchRateModel;
import dr.evomodel.sitemodel.SiteModel;
import dr.inference.model.Model;
import dr.inference.model.Parameter;
import dr.math.MathUtils;
import dr.xml.AbstractXMLObjectParser;
import dr.xml.AttributeRule;
import dr.xml.ElementRule;
import dr.xml.XMLObject;
import dr.xml.XMLObjectParser;
import dr.xml.XMLParseException;
import dr.xml.XMLSyntaxRule;

/**
 * Class for performing random sequence generation for a given site model.
 * Sequences for the leave nodes in the tree are returned as an alignment.
 */

public class EpochBranchSubstitutionModelSimulator {

	/** nr of samples to generate **/
	private int nReplications;
	/** tree used for generating samples **/
	private Tree tree;
//	private List<SubstitutionModel> substModelList;
//	private List<FrequencyModel> frequencyModelList;
//	private Parameter epochTimes;
	/** site model used for generating samples **/
	private GammaSiteRateModel siteModel;
	/** branch rate model used for generating samples **/
	private BranchRateModel branchRateModel;
	/** nr of categories in site model **/
	private int categoryCount;
	/** nr of states in site model **/
	private int stateCount;

	private static boolean has_ancestralSequence = false;
	private Sequence ancestralSequence;

	/** an array used to transfer transition probabilities **/
	private double[][] transitionProbabilities;

	
	private EpochBranchSubstitutionModel epochModel;
	
	/**
	 * Constructor
	 * 
	 * @param tree
	 * @param siteModel
	 * @param branchRateModel
	 * @param sequenceLength: nr of sites to generate
	 */
//	EpochBranchSubstitutionModelSimulator(Tree tree,
//			EpochBranchSubstitutionModel epochModel,
//			List<SubstitutionModel> substModelList,
////			List<FrequencyModel> frequencyModelList, 
//			Parameter epochTimes,
//			GammaSiteRateModel siteModel, 
//			BranchRateModel branchRateModel,
//			int sequenceLength) {
//
//		this.tree = tree;
//		this.epochModel = epochModel;
//		this.substModelList = substModelList;
////		this.frequencyModelList = frequencyModelList;
//		this.epochTimes = epochTimes;
//		this.siteModel = siteModel;
//		this.branchRateModel = branchRateModel;
//		this.nReplications = sequenceLength;
//		this.stateCount = substModelList.get(0).getDataType().getStateCount();
//		this.categoryCount = siteModel.getCategoryCount();
//		this.transitionProbabilities = new double[categoryCount][stateCount * stateCount];
//	} // END: Constructor
	
	//suggestion: use this type of constructor
	EpochBranchSubstitutionModelSimulator(Tree tree, EpochBranchSubstitutionModel epochModel, GammaSiteRateModel siteModel, int sequenceLength) {
		this.tree = tree;
		this.siteModel = siteModel;
		this.nReplications = sequenceLength;
		this.epochModel = epochModel;
		//don't really need a siteModel to get the stateCount
		//too bad the siteModel only functions as a container with respect to the substitutionModel
		this.stateCount = epochModel.getSubstitutionModel(0,0).getDataType().getStateCount();
		//don't like this categoryCount though, does require a siteModel
		this.categoryCount = siteModel.getCategoryCount();
		this.transitionProbabilities = new double[categoryCount][stateCount * stateCount];
	} // END: Constructor
	
	/**
	 * Convert integer representation of sequence into a Sequence
	 * 
	 * @param seq: integer representation of the sequence
	 * @param node: used to determine taxon for sequence
	 * @return Sequence
	 */
	Sequence intArray2Sequence(int[] seq, NodeRef node) {
		
		StringBuilder sSeq = new StringBuilder();
		
		for (int i = 0; i < nReplications; i++) {
			sSeq.append(siteModel.getSubstitutionModel().getDataType().getCode(seq[i]));
		}
		
		return new Sequence(tree.getNodeTaxon(node), sSeq.toString());
	} // END: intArray2Sequence

	void setAncestralSequence(Sequence seq) {
		ancestralSequence = seq;
		has_ancestralSequence = true;
	}

	int[] sequence2intArray(Sequence seq) {

		if (seq.getLength() != nReplications) {
			throw new RuntimeException("Ancestral sequence length has "
					+ seq.getLength() + " characters " + "expecting "
					+ nReplications + " characters");
		}

		int array[] = new int[nReplications];
		for (int i = 0; i < nReplications; i++) {
			array[i] = siteModel.getSubstitutionModel().getDataType().getState(
					seq.getChar(i));
		}
		
		return array;
	}// END: sequence2intArray

	/**
	 * perform the actual sequence generation
	 * 
	 * @return alignment containing randomly generated sequences for the nodes
	 *         in the leaves of the tree
	 */
	public Alignment simulate() {

		NodeRef root = tree.getRoot();

		double[] categoryProbs = siteModel.getCategoryProportions();
		int[] category = new int[nReplications];
		for (int i = 0; i < nReplications; i++) {
			category[i] = MathUtils.randomChoicePDF(categoryProbs);
		}

		int[] seq = new int[nReplications];

		if (has_ancestralSequence) {

			seq = sequence2intArray(ancestralSequence);

		} else {

			//FrequencyModel frequencyModel = siteModel.getSubstitutionModel().getFrequencyModel();
			FrequencyModel frequencyModel = epochModel.getSubstitutionModel(0,0).getFrequencyModel();
			for (int i = 0; i < nReplications; i++) {
				seq[i] = MathUtils.randomChoicePDF(frequencyModel.getFrequencies());
			}

		}// END: ancestral sequence check

		SimpleAlignment alignment = new SimpleAlignment();
		alignment.setReportCountStatistics(true);
		alignment.setDataType(epochModel.getSubstitutionModel(0,0).getDataType());

		traverse(root, seq, category, alignment);

		return alignment;
	} // END: simulate

	// TODO: traverse for Epoch model
	void traverse(NodeRef node, int[] parentSequence, int[] category,
			SimpleAlignment alignment) {

		for (int iChild = 0; iChild < tree.getChildCount(node); iChild++) {

			NodeRef child = tree.getChild(node, iChild);
			for (int i = 0; i < categoryCount; i++) {
				getTransitionProbabilities(tree, child, i, transitionProbabilities[i]);
			}

			int[] seq = new int[nReplications];
			double[] cProb = new double[stateCount];
			for (int i = 0; i < nReplications; i++) {
				System.arraycopy(transitionProbabilities[category[i]],
						parentSequence[i] * stateCount, cProb, 0, stateCount);
{'fraction_non_alphanumeric': 0.1079930745924109, 'fraction_numerical': 0.0025970278459096813, 'mean_word_length': 4.565234845443597, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 14}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6541301', 'n_tokens_mistral': 4514, 'n_tokens_neox': 4059, 'n_words': 1050}