Examples with FlexibleNode dr.evolution.tree.FlexibleNode used on opensource projects

Example 31 with FlexibleNode

use of dr.evolution.tree.FlexibleNode in project beast-mcmc by beast-dev.

the class CaseToCaseTreeLikelihood method explodeTree.

protected void explodeTree() {
    for (int i = 0; i < outbreak.size(); i++) {
        AbstractCase aCase = outbreak.getCase(i);
        if (aCase.wasEverInfected() && elementsAsTrees.get(aCase) == null) {
            NodeRef partitionRoot = ((PartitionedTreeModel) treeModel).getEarliestNodeInElement(aCase);
            double extraHeight;
            if (treeModel.isRoot(partitionRoot)) {
                extraHeight = maxFirstInfToRoot.getParameterValue(0) * aCase.getInfectionBranchPosition().getParameterValue(0);
            } else {
                extraHeight = treeModel.getBranchLength(partitionRoot) * aCase.getInfectionBranchPosition().getParameterValue(0);
            }
            FlexibleNode newRoot = new FlexibleNode();
            FlexibleTree littleTree = new FlexibleTree(newRoot);
            littleTree.beginTreeEdit();
            if (!treeModel.isExternal(partitionRoot)) {
                for (int j = 0; j < treeModel.getChildCount(partitionRoot); j++) {
                    copyElementToTreelet(littleTree, treeModel.getChild(partitionRoot, j), newRoot, aCase);
                }
            }
            littleTree.endTreeEdit();
            littleTree.resolveTree();
            Treelet treelet = new Treelet(littleTree, littleTree.getRootHeight() + extraHeight);
            elementsAsTrees.put(aCase, treelet);
        }
    }
}

Example 32 with FlexibleNode

use of dr.evolution.tree.FlexibleNode in project beast-mcmc by beast-dev.

the class CaseToCaseTreeLikelihood method copyElementToTreelet.

private void copyElementToTreelet(FlexibleTree littleTree, NodeRef oldNode, NodeRef newParent, AbstractCase element) {
    if (element.wasEverInfected()) {
        if (getBranchMap().get(oldNode.getNumber()) == element) {
            if (treeModel.isExternal(oldNode)) {
                NodeRef newTip = new FlexibleNode(new Taxon(treeModel.getNodeTaxon(oldNode).getId()));
                littleTree.addChild(newParent, newTip);
                littleTree.setBranchLength(newTip, treeModel.getBranchLength(oldNode));
            } else {
                NodeRef newChild = new FlexibleNode();
                littleTree.addChild(newParent, newChild);
                littleTree.setBranchLength(newChild, treeModel.getBranchLength(oldNode));
                for (int i = 0; i < treeModel.getChildCount(oldNode); i++) {
                    copyElementToTreelet(littleTree, treeModel.getChild(oldNode, i), newChild, element);
                }
            }
        } else {
            // we need a new tip
            NodeRef transmissionTip = new FlexibleNode(new Taxon("Transmission_" + getBranchMap().get(oldNode.getNumber()).getName()));
            double parentTime = getNodeTime(treeModel.getParent(oldNode));
            double childTime = getInfectionTime(getBranchMap().get(oldNode.getNumber()));
            littleTree.addChild(newParent, transmissionTip);
            littleTree.setBranchLength(transmissionTip, childTime - parentTime);
        }
    }
}

Example 33 with FlexibleNode

use of dr.evolution.tree.FlexibleNode in project beast-mcmc by beast-dev.

the class CaseToCaseTreeLikelihood method addTransmissionNodes.

public FlexibleTree addTransmissionNodes(Tree tree) {
    prepareTimings();
    FlexibleTree outTree = new FlexibleTree(tree, true);
    for (int j = 0; j < outTree.getNodeCount(); j++) {
        FlexibleNode node = (FlexibleNode) outTree.getNode(j);
        node.setAttribute("Number", node.getNumber());
        node.setAttribute("Time", heightToTime(node.getHeight()));
        node.setAttribute(PARTITIONS_KEY, getBranchMap().get(node.getNumber()));
    }
    for (AbstractCase aCase : outbreak.getCases()) {
        if (aCase.wasEverInfected()) {
            NodeRef originalNode = ((PartitionedTreeModel) treeModel).getEarliestNodeInElement(aCase);
            int infectionNodeNo = originalNode.getNumber();
            if (!treeModel.isRoot(originalNode)) {
                NodeRef originalParent = treeModel.getParent(originalNode);
                double nodeTime = getNodeTime(originalNode);
                double infectionTime = getInfectionTime(aCase);
                double heightToBreakBranch = getHeight(originalNode) + (nodeTime - infectionTime);
                FlexibleNode newNode = (FlexibleNode) outTree.getNode(infectionNodeNo);
                FlexibleNode oldParent = (FlexibleNode) outTree.getParent(newNode);
                outTree.beginTreeEdit();
                outTree.removeChild(oldParent, newNode);
                FlexibleNode infectionNode = new FlexibleNode();
                infectionNode.setHeight(heightToBreakBranch);
                infectionNode.setLength(oldParent.getHeight() - heightToBreakBranch);
                infectionNode.setAttribute(PARTITIONS_KEY, getNodePartition(treeModel, originalParent));
                infectionNode.setAttribute("Time", heightToTime(heightToBreakBranch));
                newNode.setLength(nodeTime - infectionTime);
                outTree.addChild(oldParent, infectionNode);
                outTree.addChild(infectionNode, newNode);
                outTree.endTreeEdit();
            } else {
                double nodeTime = getNodeTime(originalNode);
                double infectionTime = getInfectionTime(aCase);
                double heightToInstallRoot = getHeight(originalNode) + (nodeTime - infectionTime);
                FlexibleNode newNode = (FlexibleNode) outTree.getNode(infectionNodeNo);
                outTree.beginTreeEdit();
                FlexibleNode infectionNode = new FlexibleNode();
                infectionNode.setHeight(heightToInstallRoot);
                infectionNode.setAttribute("Time", heightToTime(heightToInstallRoot));
                infectionNode.setAttribute(PARTITIONS_KEY, "Origin");
                outTree.addChild(infectionNode, newNode);
                newNode.setLength(heightToInstallRoot - getHeight(originalNode));
                outTree.setRoot(infectionNode);
                outTree.endTreeEdit();
            }
        }
    }
    outTree = new FlexibleTree((FlexibleNode) outTree.getRoot());
    for (int i = 0; i < outTree.getNodeCount(); i++) {
        NodeRef node = outTree.getNode(i);
        NodeRef parent = outTree.getParent(node);
        if (parent != null && outTree.getNodeHeight(node) > outTree.getNodeHeight(parent)) {
            try {
                NexusExporter exporter = new NexusExporter(new PrintStream("fancyProblem.nex"));
                exporter.exportTree(outTree);
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                ((PartitionedTreeModel) treeModel).checkPartitions();
            } catch (BadPartitionException e) {
                System.out.print("Rewiring messed up because of partition problem.");
            }
            throw new RuntimeException("Rewiring messed up; investigate");
        }
    }
    return outTree;
}

Example 34 with FlexibleNode

use of dr.evolution.tree.FlexibleNode in project beast-mcmc by beast-dev.

the class WithinCaseCoalescent method debugTreelet.

public void debugTreelet(Tree treelet, String fileName) {
    try {
        FlexibleTree treeCopy = new FlexibleTree(treelet);
        for (int j = 0; j < treeCopy.getNodeCount(); j++) {
            FlexibleNode node = (FlexibleNode) treeCopy.getNode(j);
            node.setAttribute("Number", node.getNumber());
        }
        NexusExporter testTreesOut = new NexusExporter(new PrintStream(fileName));
        testTreesOut.exportTree(treeCopy);
    } catch (IOException ignored) {
        System.out.println("IOException");
    }
}