Examples with PartitionedTreeModel dr.evomodel.epidemiology.casetocase.PartitionedTreeModel used on opensource projects

Example 1 with PartitionedTreeModel

use of dr.evomodel.epidemiology.casetocase.PartitionedTreeModel in project beast-mcmc by beast-dev.

the class TransmissionWilsonBaldingA method proposeTree.

public void proposeTree() {
    PartitionedTreeModel tree = c2cLikelihood.getTreeModel();
    BranchMapModel branchMap = c2cLikelihood.getBranchMap();
    NodeRef i;
    double oldMinAge, newMinAge, newRange, oldRange, newAge, q;
    // choose a random node avoiding root, and nodes that are ineligible for this move because they have nowhere to
    // go
    ArrayList<NodeRef> eligibleNodes = getEligibleNodes(tree, branchMap);
    i = eligibleNodes.get(MathUtils.nextInt(eligibleNodes.size()));
    double eligibleNodeCount = eligibleNodes.size();
    final NodeRef iP = tree.getParent(i);
    Integer[] sameElements = tree.samePartitionElement(iP);
    HashSet<Integer> possibleDestinations = new HashSet<Integer>();
    // we can insert the node above OR BELOW any node in the same partition
    for (Integer sameElement : sameElements) {
        possibleDestinations.add(sameElement);
        if (!tree.isExternal(tree.getNode(sameElement))) {
            possibleDestinations.add(tree.getChild(tree.getNode(sameElement), 0).getNumber());
            possibleDestinations.add(tree.getChild(tree.getNode(sameElement), 1).getNumber());
        }
    }
    Integer[] pd = possibleDestinations.toArray(new Integer[possibleDestinations.size()]);
    NodeRef j = tree.getNode(pd[MathUtils.nextInt(pd.length)]);
    NodeRef jP = tree.getParent(j);
    while ((jP != null && (tree.getNodeHeight(jP) <= tree.getNodeHeight(i))) || (i == j)) {
        j = tree.getNode(pd[MathUtils.nextInt(pd.length)]);
        jP = tree.getParent(j);
    }
    if (iP == tree.getRoot() || j == tree.getRoot()) {
        logq = Double.NEGATIVE_INFINITY;
        return;
    }
    if (jP == iP || j == iP || jP == i) {
        logq = Double.NEGATIVE_INFINITY;
        return;
    }
    final NodeRef CiP = getOtherChild(tree, iP, i);
    NodeRef PiP = tree.getParent(iP);
    if (resampleInfectionTimes) {
        AbstractCase iCase = branchMap.get(i.getNumber());
        AbstractCase iPCase = branchMap.get(iP.getNumber());
        AbstractCase CiPCase = branchMap.get(CiP.getNumber());
        AbstractCase PiPCase = null;
        if (PiP != null) {
            PiPCase = branchMap.get(PiP.getNumber());
        }
        if (iCase != iPCase) {
            iCase.setInfectionBranchPosition(MathUtils.nextDouble());
        }
        // what happens between PiP and CiP
        if (PiPCase == null || CiPCase != PiPCase) {
            CiPCase.setInfectionBranchPosition(MathUtils.nextDouble());
        }
        // what happens between k and j
        AbstractCase jCase = branchMap.get(j.getNumber());
        AbstractCase kCase = branchMap.get(jP.getNumber());
        if (iPCase != jCase && iPCase != kCase) {
            throw new RuntimeException("TWBA misbehaving.");
        }
        jCase.setInfectionBranchPosition(MathUtils.nextDouble());
    }
    newMinAge = Math.max(tree.getNodeHeight(i), tree.getNodeHeight(j));
    newRange = tree.getNodeHeight(jP) - newMinAge;
    newAge = newMinAge + (MathUtils.nextDouble() * newRange);
    oldMinAge = Math.max(tree.getNodeHeight(i), tree.getNodeHeight(CiP));
    oldRange = tree.getNodeHeight(PiP) - oldMinAge;
    q = newRange / Math.abs(oldRange);
    tree.beginTreeEdit();
    if (j == tree.getRoot()) {
        // 1. remove edges <iP, CiP>
        tree.removeChild(iP, CiP);
        tree.removeChild(PiP, iP);
        // 2. add edges <k, iP>, <iP, j>, <PiP, CiP>
        tree.addChild(iP, j);
        tree.addChild(PiP, CiP);
        // iP is the new root
        tree.setRoot(iP);
    } else if (iP == tree.getRoot()) {
        // 1. remove edges <k, j>, <iP, CiP>, <PiP, iP>
        tree.removeChild(jP, j);
        tree.removeChild(iP, CiP);
        // 2. add edges <k, iP>, <iP, j>, <PiP, CiP>
        tree.addChild(iP, j);
        tree.addChild(jP, iP);
        // CiP is the new root
        tree.setRoot(CiP);
    } else {
        // 1. remove edges <k, j>, <iP, CiP>, <PiP, iP>
        tree.removeChild(jP, j);
        tree.removeChild(iP, CiP);
        tree.removeChild(PiP, iP);
        // 2. add edges <k, iP>, <iP, j>, <PiP, CiP>
        tree.addChild(iP, j);
        tree.addChild(jP, iP);
        tree.addChild(PiP, CiP);
    }
    tree.setNodeHeight(iP, newAge);
    tree.endTreeEdit();
    if (DEBUG) {
        c2cLikelihood.getTreeModel().checkPartitions();
    }
    logq = Math.log(q);
    double reverseEligibleNodeCount = getEligibleNodes(tree, branchMap).size();
    logq += Math.log(eligibleNodeCount / reverseEligibleNodeCount);
}