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

Example 21 with FlexibleTree

use of dr.evolution.tree.FlexibleTree 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 22 with FlexibleTree

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

the class CaseToCaseTreeLikelihood method outputTreeToFile.

public void outputTreeToFile(BranchMapModel map, String fileName, boolean includeTransmissionNodes) {
    try {
        FlexibleTree treeCopy;
        if (!includeTransmissionNodes) {
            treeCopy = new FlexibleTree(treeModel);
            for (int j = 0; j < treeCopy.getNodeCount(); j++) {
                FlexibleNode node = (FlexibleNode) treeCopy.getNode(j);
                node.setAttribute("Number", node.getNumber());
                node.setAttribute("Time", heightToTime(node.getHeight()));
                node.setAttribute(PARTITIONS_KEY, map.get(node.getNumber()));
            }
        } else {
            treeCopy = addTransmissionNodes(treeModel);
        }
        NexusExporter testTreesOut = new NexusExporter(new PrintStream(fileName));
        testTreesOut.exportTree(treeCopy);
    } catch (IOException ignored) {
        System.out.println("IOException");
    }
}

Example 23 with FlexibleTree

use of dr.evolution.tree.FlexibleTree 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 24 with FlexibleTree

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

the class BranchJumpPlotter method rewireTree.

private FlexibleTree rewireTree(Tree tree, boolean verbose) {
    int totalJumps = 0;
    FlexibleTree outTree = new FlexibleTree(tree, true);
    for (int nodeNo = 0; nodeNo < outTree.getNodeCount(); nodeNo++) {
        FlexibleNode node = (FlexibleNode) outTree.getNode(nodeNo);
        String finalHost = (String) node.getAttribute(traitName);
        node.setAttribute(traitName, finalHost.replaceAll("\"", ""));
        Object[] jumps = readCJH(node);
        if (verbose) {
            System.out.print("Node " + nodeNo + ": ");
        }
        if (jumps != null) {
            FlexibleNode needsNewParent = node;
            Double height = tree.getNodeHeight(node);
            for (int i = jumps.length - 1; i >= 0; i--) {
                totalJumps++;
                Object[] jump = (Object[]) jumps[i];
                if (i < jumps.length - 1 && (Double) jump[1] <= height) {
                    throw new RuntimeException("Jumps do not appear to be in descending order of height");
                }
                height = (Double) jump[1];
                if (!needsNewParent.getAttribute(traitName).equals(jump[3])) {
                    throw new RuntimeException("Destination traits do not match");
                }
                FlexibleNode parent = (FlexibleNode) outTree.getParent(needsNewParent);
                outTree.beginTreeEdit();
                outTree.removeChild(parent, needsNewParent);
                needsNewParent.setLength(height - needsNewParent.getHeight());
                FlexibleNode jumpNode = new FlexibleNode();
                jumpNode.setHeight(height);
                jumpNode.setLength(parent.getHeight() - height);
                jumpNode.setAttribute(traitName, jump[2]);
                outTree.addChild(parent, jumpNode);
                outTree.addChild(jumpNode, needsNewParent);
                outTree.endTreeEdit();
                needsNewParent = jumpNode;
            }
        }
        if (verbose) {
            if (jumps == null) {
                System.out.println(0 + " (" + totalJumps + ")");
            } else {
                System.out.println(jumps.length + " (" + totalJumps + ")");
            }
        }
    }
    outTree = new FlexibleTree((FlexibleNode) outTree.getRoot());
    if (verbose) {
        System.out.println("Total jumps: " + totalJumps);
        int[] childCounts = new int[3];
        for (int i = 0; i < outTree.getNodeCount(); i++) {
            childCounts[outTree.getChildCount(outTree.getNode(i))]++;
        }
        for (int i = 0; i < 3; i++) {
            System.out.println(childCounts[i] + " nodes have " + i + " children");
        }
    }
    return outTree;
}

Example 25 with FlexibleTree

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

the class GetAncestralSequenceFromSplitTrait method analyze.

/**
     * Recursively analyzes log files.
     *
     * @param treeAnnotatorInputFile       the file to analyze (if this is a directory then the files within it are analyzed)
     * @throws dr.inference.trace.TraceException
     *          if the trace file is in the wrong format or corrupted
     */
private void analyze(File treeAnnotatorInputFile) throws TraceException {
    try {
        FileReader fileReader = new FileReader(treeAnnotatorInputFile);
        TreeImporter importer = new NexusImporter(fileReader);
        FlexibleTree tree = (FlexibleTree) importer.importNextTree();
        for (int i = 0; i < tree.getNodeCount(); i++) {
            Hashtable<Integer, State> states = new Hashtable<Integer, State>();
            for (Iterator<String> j = tree.getNodeAttributeNames(tree.getNode(i)); j.hasNext(); ) {
                String name = j.next();
                if (name.indexOf("states_") >= 0) {
                    Integer d = Integer.parseInt(name.replaceFirst("states_", "").replaceFirst("\\..+", ""));
                    //= new State(name.);
                    State s;
                    if (states.containsKey(d)) {
                        s = states.get(d);
                    } else {
                        s = new State(d);
                    }
                    if (name.matches("states_" + d + ".prob")) {
                        Object o = tree.getNodeAttribute(tree.getNode(i), name);
                        double probability = (Double) o;
                        s.setProbability(probability);
                    } else if (name.matches("states_" + d)) {
                        Object o = tree.getNodeAttribute(tree.getNode(i), name);
                        String value = (String) o;
                        s.setState(value.replaceAll("\"", ""));
                    } else if (name.matches("states_" + d + ".set.prob")) {
                        /* Not necessary but lets parse it anyways */
                        Object[] o = (Object[]) tree.getNodeAttribute(tree.getNode(i), name);
                        double[] probabilities = new double[o.length];
                        for (int k = 0; k < o.length; k++) {
                            probabilities[k] = (Double) o[k];
                        }
                        s.setProbabilities(probabilities);
                    } else if (name.matches("states_" + d + ".set")) {
                        /* Not necessary but lets parse it anyways */
                        Object[] o = (Object[]) tree.getNodeAttribute(tree.getNode(i), name);
                        String[] set = new String[o.length];
                        for (int k = 0; k < o.length; k++) {
                            set[k] = ((String) o[k]).replaceAll("\"", "");
                        }
                        s.setSet(set);
                    }
                    //                        }
                    states.put(d, s);
                }
            }
            State[] statesArray = states.values().toArray(new State[states.size()]);
            Arrays.sort(statesArray);
            /* Set the default length to the number of characters that it would need */
            StringBuffer sb = new StringBuffer(statesArray.length * statesArray[0].getState().length());
            for (State s : statesArray) {
                sb.append(s.getState());
            }
            tree.setNodeAttribute(tree.getNode(i), "seq", sb.toString());
        }
        /* Export the new tree with the new sequences */
        TreeExporter exporter = new NexusExporter(System.out);
        exporter.exportTree(tree);
        System.out.println("Begin trees;");
        System.out.println("\ttree max_tree = " + tree.toString());
        System.out.println("End;");
    } catch (IOException e) {
        System.err.println("Error Parsing Input log: " + e.getMessage());
    } catch (Importer.ImportException e) {
        System.err.println("Error Parsing Input Tree: " + e.getMessage());
    }
}