/*
* InitialTreeGenerator.java
*
* Copyright (c) 2002-2015 Alexei Drummond, Andrew Rambaut and Marc Suchard
*
* This file is part of BEAST.
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership and licensing.
*
* BEAST is free software; you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as
* published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* BEAST is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with BEAST; if not, write to the
* Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
* Boston, MA 02110-1301 USA
*/
package dr.app.beauti.generator;
import dr.app.beauti.components.ComponentFactory;
import dr.app.beauti.options.*;
import dr.app.beauti.types.PriorType;
import dr.app.beauti.types.TreePriorType;
import dr.app.beauti.util.XMLWriter;
import dr.evolution.tree.NodeRef;
import dr.evolution.tree.Tree;
import dr.evolution.tree.TreeUtils;
import dr.evolution.util.Taxa;
import dr.evomodelxml.coalescent.CoalescentSimulatorParser;
import dr.evomodelxml.coalescent.OldCoalescentSimulatorParser;
import dr.evomodelxml.coalescent.ConstantPopulationModelParser;
import dr.evomodelxml.coalescent.ExponentialGrowthModelParser;
import dr.evoxml.*;
import dr.util.Attribute;
import dr.xml.XMLParser;
import java.util.ArrayList;
import java.util.List;
/**
* @author Alexei Drummond
* @author Andrew Rambaut
* @author Walter Xie
*/
public class InitialTreeGenerator extends Generator {
final static public String STARTING_TREE = "startingTree";
public InitialTreeGenerator(BeautiOptions options, ComponentFactory[] components) {
super(options, components);
}
/**
* Generate XML for the starting tree
* @param model PartitionTreeModel
*
* @param writer the writer
*/
public void writeStartingTree(PartitionTreeModel model, XMLWriter writer) {
setModelPrefix(model.getPrefix()); // only has prefix, if (options.getPartitionTreeModels().size() > 1)
switch (model.getStartingTreeType()) {
case USER:
case UPGMA:
Parameter rootHeight = model.getParameter("treeModel.rootHeight");
// generate a rescaled starting tree
writer.writeComment("Construct a starting tree that is compatible with specified clade heights");
Attribute[] attributes = (rootHeight.priorType != PriorType.NONE_TREE_PRIOR ?
new Attribute[] {
new Attribute.Default<String>(XMLParser.ID, modelPrefix + STARTING_TREE),
new Attribute.Default<String>(RescaledTreeParser.HEIGHT, "" + rootHeight.getInitial())
} :
new Attribute[] {
new Attribute.Default<String>(XMLParser.ID, modelPrefix + STARTING_TREE)
});
writer.writeOpenTag(RescaledTreeParser.RESCALED_TREE, attributes);
writeSourceTree(model, writer);
if (options.taxonSets != null && options.taxonSets.size() > 0 && !options.useStarBEAST) {
for (Taxa taxa : options.taxonSets) {
Double height = options.taxonSetsHeights.get(taxa);
if (height != null) {
writer.writeOpenTag(RescaledTreeParser.CLADE, new Attribute.Default<String>(RescaledTreeParser.HEIGHT, height.toString()));
writer.writeTag("taxa", new Attribute.Default<String>(XMLParser.IDREF, taxa.getId()), true);
writer.writeCloseTag(RescaledTreeParser.CLADE);
} else if (options.taxonSetsMono.get(taxa)) {
// if monophyly is enforced then placing this clade element here will force BEAST to check
// the clade exists in the tree.
writer.writeOpenTag(RescaledTreeParser.CLADE);
writer.writeTag("taxa", new Attribute.Default<String>(XMLParser.IDREF, taxa.getId()), true);
writer.writeCloseTag(RescaledTreeParser.CLADE);
}
}
}
writer.writeCloseTag(RescaledTreeParser.RESCALED_TREE);
break;
case RANDOM:
// generate a coalescent tree
String simulatorId = modelPrefix + STARTING_TREE;
String taxaId = TaxaParser.TAXA;
AbstractPartitionData partition = options.getDataPartitions(model).get(0);
if (!options.hasIdenticalTaxa()) {
taxaId = partition.getPartitionTreeModel().getPrefix() + TaxaParser.TAXA;
}
if (partition instanceof PartitionPattern && ((PartitionPattern) partition).getPatterns().hasMask()) {
taxaId = partition.getPrefix() + TaxaParser.TAXA;
}
writer.writeComment("Generate a random starting tree under the coalescent process");
if (options.taxonSets != null && options.taxonSets.size() > 0 && !options.useStarBEAST) { // need !options.useStarBEAST,
writeSubTree(simulatorId, taxaId, options.taxonList, model, writer);
} else {
writer.writeOpenTag(
CoalescentSimulatorParser.COALESCENT_SIMULATOR,
new Attribute[]{
new Attribute.Default<String>(XMLParser.ID, simulatorId)
}
);
writeTaxaRef(taxaId, model, writer);
writeInitialDemoModelRef(model, writer);
writer.writeCloseTag(CoalescentSimulatorParser.COALESCENT_SIMULATOR);
}
break;
default:
throw new IllegalArgumentException("Unknown StartingTreeType");
}
}
public void writeSourceTree(PartitionTreeModel model, XMLWriter writer) {
switch (model.getStartingTreeType()) {
case USER:
if (model.isNewick()) {
writeNewickTree(model.getUserStartingTree(), writer);
} else {
writeSimpleTree(model.getUserStartingTree(), writer);
}
break;
case UPGMA:
// generate a upgma starting tree
writer.writeComment("Construct a rough-and-ready UPGMA tree as an starting tree");
writer.writeOpenTag(UPGMATreeParser.UPGMA_TREE);
writer.writeOpenTag(
DistanceMatrixParser.DISTANCE_MATRIX,
new Attribute[]{
new Attribute.Default<String>(DistanceMatrixParser.CORRECTION, "JC")
}
);
writer.writeOpenTag(SitePatternsParser.PATTERNS);
writer.writeComment("To generate UPGMA starting tree, only use the 1st aligment, "
+ "which may be 1 of many aligments using this tree.");
writer.writeIDref(AlignmentParser.ALIGNMENT, options.getDataPartitions(model).get(0).getTaxonList().getId());
// alignment has no gene prefix
writer.writeCloseTag(SitePatternsParser.PATTERNS);
writer.writeCloseTag(DistanceMatrixParser.DISTANCE_MATRIX);
writer.writeCloseTag(UPGMATreeParser.UPGMA_TREE);
break;
case RANDOM:
throw new IllegalArgumentException("Shouldn't be here");
default:
throw new IllegalArgumentException("Unknown StartingTreeType");
}
}
private void writeTaxaRef(String taxaId, PartitionTreeModel model, XMLWriter writer) {
Attribute[] taxaAttribute = {new Attribute.Default<String>(XMLParser.IDREF, taxaId)};
if (options.taxonSets != null && options.taxonSets.size() > 0 && !options.useStarBEAST) { // need !options.useStarBEAST,
// *BEAST case is in STARBEASTGenerator.writeStartingTreeForCalibration(XMLWriter writer)
writer.writeOpenTag(OldCoalescentSimulatorParser.CONSTRAINED_TAXA);
writer.writeTag(TaxaParser.TAXA, taxaAttribute, true);
for (Taxa taxa : options.taxonSets) {
if (options.taxonSetsTreeModel.get(taxa).equals(model)) {
Parameter statistic = options.getStatistic(taxa);
Attribute mono = new Attribute.Default<Boolean>(
OldCoalescentSimulatorParser.IS_MONOPHYLETIC, options.taxonSetsMono.get(taxa));
writer.writeOpenTag(OldCoalescentSimulatorParser.TMRCA_CONSTRAINT, mono);
writer.writeIDref(TaxaParser.TAXA, taxa.getId());
if (model.getPartitionTreePrior().getNodeHeightPrior() == TreePriorType.YULE_CALIBRATION
&& statistic.priorType == PriorType.UNIFORM_PRIOR) {
writeDistribution(statistic, false, writer);
}
writer.writeCloseTag(OldCoalescentSimulatorParser.TMRCA_CONSTRAINT);
}
}
writer.writeCloseTag(OldCoalescentSimulatorParser.CONSTRAINED_TAXA);
} else {
writer.writeTag(TaxaParser.TAXA, taxaAttribute, true);
}
}
private void writeSubTree(String treeId, String taxaId, Taxa taxa, PartitionTreeModel model, XMLWriter writer) {
Double height = options.taxonSetsHeights.get(taxa);
if (height == null) {
height = Double.NaN;
}
Attribute[] attributes = new Attribute[] {};
if (treeId != null) {
if (Double.isNaN(height)) {
attributes = new Attribute[] {
new Attribute.Default<String>(XMLParser.ID, treeId)
};
} else {
attributes = new Attribute[] {
new Attribute.Default<String>(XMLParser.ID, treeId),
new Attribute.Default<String>(CoalescentSimulatorParser.HEIGHT, "" + height)
};
}
} else {
if (!Double.isNaN(height)) {
attributes = new Attribute[] {
new Attribute.Default<String>(CoalescentSimulatorParser.HEIGHT, "" + height)
};
}
}
// construct a subtree
writer.writeOpenTag(
CoalescentSimulatorParser.COALESCENT_SIMULATOR,
attributes
);
List<Taxa> subsets = new ArrayList<Taxa>();
// Taxa remainingTaxa = new Taxa(taxa);
for (Taxa taxa2 : options.taxonSets) {
boolean sameTree = model.equals(options.taxonSetsTreeModel.get(taxa2));
boolean isMono = options.taxonSetsMono.get(taxa2);
boolean hasHeight = options.taxonSetsHeights.get(taxa2) != null;
boolean isSubset = taxa.containsAll(taxa2);
if (sameTree && (isMono || hasHeight) && taxa2 != taxa && isSubset) {
subsets.add(taxa2);
}
}
List<Taxa> toRemove = new ArrayList<Taxa>();
for (Taxa taxa3 : subsets) {
boolean isSubSubSet = false;
for (Taxa taxa4 : subsets) {
if (!taxa4.equals(taxa3) && taxa4.containsAll(taxa3)) {
isSubSubSet = true;
}
}
if (isSubSubSet) {
toRemove.add(taxa3);
}
}
subsets.removeAll(toRemove);
for (Taxa taxa5 : subsets) {
// remainingTaxa.removeTaxa(taxa5);
writeSubTree(null, null, taxa5, model, writer);
}
if (taxaId == null) {
writer.writeIDref(TaxaParser.TAXA, taxa.getId());
} else {
writer.writeIDref(TaxaParser.TAXA, taxaId);
}
writeInitialDemoModelRef(model, writer);
writer.writeCloseTag(CoalescentSimulatorParser.COALESCENT_SIMULATOR);
}
private void writeInitialDemoModelRef(PartitionTreeModel model, XMLWriter writer) {
PartitionTreePrior prior = model.getPartitionTreePrior();
if (prior.getNodeHeightPrior() == TreePriorType.CONSTANT || options.useStarBEAST) {
writer.writeIDref(ConstantPopulationModelParser.CONSTANT_POPULATION_MODEL, prior.getPrefix() + "constant");
} else if (prior.getNodeHeightPrior() == TreePriorType.EXPONENTIAL) {
writer.writeIDref(ExponentialGrowthModelParser.EXPONENTIAL_GROWTH_MODEL, prior.getPrefix() + "exponential");
} else {
writer.writeIDref(ConstantPopulationModelParser.CONSTANT_POPULATION_MODEL, prior.getPrefix() + "initialDemo");
}
}
private void writeNewickTree (Tree tree, XMLWriter writer) {
writer.writeComment("The user-specified starting tree in a newick tree format.");
writer.writeOpenTag(
NewickParser.NEWICK,
new Attribute[]{
// new Attribute.Default<String>(XMLParser.ID, modelPrefix + STARTING_TREE),
// new Attribute.Default<String>(DateParser.UNITS, options.datesUnits.getAttribute()),
new Attribute.Default<Boolean>(SimpleTreeParser.USING_DATES, options.clockModelOptions.isTipCalibrated())
}
);
writer.writeText(TreeUtils.newick(tree));
writer.writeCloseTag(NewickParser.NEWICK);
}
// private void writeNewickNode(Tree tree, NodeRef node, XMLWriter writer) {
// if (tree.getChildCount(node) > 0)
// writer.writeText("(");
// if (tree.getChildCount(node) == 0)
// writer.writeText(tree.getNodeTaxon(node).getId() + " : " + tree.getBranchLength(node));
//
// for (int i = 0; i < tree.getChildCount(node); i++) {
// if (i > 0) writer.writeText(", ");
// writeNewickNode(tree, tree.getChild(node, i), writer);
//
// }
// if (tree.getChildCount(node) > 0)
// writer.writeText(")");
// }
/**
* Generate XML for the user tree
*
* @param tree the user tree
* @param writer the writer
*/
private void writeSimpleTree(Tree tree, XMLWriter writer) {
writer.writeComment("The user-specified starting tree in a simple tree format.");
writer.writeOpenTag(
SimpleTreeParser.SIMPLE_TREE,
new Attribute[]{
// new Attribute.Default<String>(XMLParser.ID, modelPrefix + STARTING_TREE),
// new Attribute.Default<String>(DateParser.UNITS, options.datesUnits.getAttribute()),
new Attribute.Default<Object>(DateParser.UNITS, options.units.toString()),
new Attribute.Default<Boolean>(SimpleTreeParser.USING_DATES, options.clockModelOptions.isTipCalibrated())
}
);
writeSimpleNode(tree, tree.getRoot(), writer);
writer.writeCloseTag(SimpleTreeParser.SIMPLE_TREE);
}
/**
* Generate XML for the node of a user tree.
*
* @param tree the user tree
* @param node the current node
* @param writer the writer
*/
private void writeSimpleNode(Tree tree, NodeRef node, XMLWriter writer) {
writer.writeOpenTag(
SimpleNodeParser.NODE,
new Attribute[]{new Attribute.Default<Double>(SimpleNodeParser.HEIGHT, tree.getNodeHeight(node))}
);
if (tree.getChildCount(node) == 0) {
writer.writeIDref(TaxonParser.TAXON, tree.getNodeTaxon(node).getId());
}
for (int i = 0; i < tree.getChildCount(node); i++) {
writeSimpleNode(tree, tree.getChild(node, i), writer);
}
writer.writeCloseTag(SimpleNodeParser.NODE);
}
}