Examples with DefaultBranchColouring dr.evolution.colouring.DefaultBranchColouring used on opensource projects

Example 1 with DefaultBranchColouring

use of dr.evolution.colouring.DefaultBranchColouring in project beast-mcmc by beast-dev.

the class TransmissionHistoryModel method createTreeColouring.

private Taxon createTreeColouring(Tree tree, NodeRef node, DefaultTreeColouring treeColouring) {
    Taxon parentHost = null;
    Taxon childHost;
    if (tree.isExternal(node)) {
        childHost = (Taxon) tree.getNodeTaxon(node).getAttribute("host");
        if (childHost == null) {
            throw new RuntimeException("One or more of the viruses tree's taxa are missing the 'host' attribute");
        }
    } else {
        Taxon h1 = createTreeColouring(tree, tree.getChild(node, 0), treeColouring);
        Taxon h2 = createTreeColouring(tree, tree.getChild(node, 1), treeColouring);
        if (h1 != h2) {
            throw new RuntimeException("Two children have different hosts at coalescent event");
        }
        childHost = h1;
    }
    NodeRef parent = tree.getParent(node);
    if (parent != null) {
        double height0 = tree.getNodeHeight(node);
        double height1 = tree.getNodeHeight(parent);
        TransmissionEvent event = getTransmissionEventToHost(childHost);
        DefaultBranchColouring branchColouring;
        if (event != null && event.getTransmissionTime() < height1) {
            if (event.getTransmissionTime() < height0) {
                throw new RuntimeException("Transmission event is before the node");
            }
            List<Taxon> hosts = new ArrayList<Taxon>();
            List<Double> times = new ArrayList<Double>();
            parentHost = childHost;
            while (event != null && event.getTransmissionTime() < height1) {
                hosts.add(parentHost);
                times.add(event.getTransmissionTime());
                parentHost = event.donor;
                event = getTransmissionEventToHost(parentHost);
            }
            int host1 = getHostIndex(childHost);
            int host2 = getHostIndex(parentHost);
            branchColouring = new DefaultBranchColouring(host2, host1);
            for (int i = hosts.size() - 1; i >= 0; i--) {
                int host = getHostIndex(hosts.get(i));
                double time = times.get(i);
                branchColouring.addEvent(host, time);
            }
        } else {
            int host = getHostIndex(childHost);
            branchColouring = new DefaultBranchColouring(host, host);
            parentHost = childHost;
        }
        treeColouring.setBranchColouring(node, branchColouring);
    }
    return parentHost;
}