Search in sources :

Example 1 with AbstractCase

use of dr.evomodel.epidemiology.casetocase.AbstractCase 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()) {
        throw new RuntimeException("Root changes not allowed!");
    }
    if (jP == iP || j == iP || jP == i)
        throw new RuntimeException("Move failed");
    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);
}
Also used : PartitionedTreeModel(dr.evomodel.epidemiology.casetocase.PartitionedTreeModel) NodeRef(dr.evolution.tree.NodeRef) BranchMapModel(dr.evomodel.epidemiology.casetocase.BranchMapModel) AbstractCase(dr.evomodel.epidemiology.casetocase.AbstractCase) HashSet(java.util.HashSet)

Example 2 with AbstractCase

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

the class TransmissionWilsonBaldingB method proposeTree.

public void proposeTree() {
    TreeModel tree = c2cLikelihood.getTreeModel();
    BranchMapModel branchMap = c2cLikelihood.getBranchMap();
    NodeRef i;
    double oldMinAge, newMinAge, newRange, oldRange, newAge, q;
    // choose a random eligible node
    final int nodeCount = tree.getNodeCount();
    do {
        i = tree.getNode(MathUtils.nextInt(nodeCount));
    } while (!eligibleForMove(i, tree, branchMap));
    final NodeRef iP = tree.getParent(i);
    //this one can go anywhere
    NodeRef j = tree.getNode(MathUtils.nextInt(tree.getNodeCount()));
    NodeRef jP = tree.getParent(j);
    while ((jP != null && tree.getNodeHeight(jP) <= tree.getNodeHeight(i)) || (i == j)) {
        j = tree.getNode(MathUtils.nextInt(tree.getNodeCount()));
        jP = tree.getParent(j);
    }
    if (iP == tree.getRoot() || j == tree.getRoot()) {
        throw new RuntimeException("Root changes not allowed!");
    }
    if (jP == iP || j == iP || jP == i)
        throw new RuntimeException("move failed");
    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());
        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);
    if (branchMap.get(PiP.getNumber()) != branchMap.get(CiP.getNumber())) {
        q *= 0.5;
    }
    if (branchMap.get(jP.getNumber()) != branchMap.get(j.getNumber())) {
        q *= 2;
    }
    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();
    //
    logq = Math.log(q);
    if (MathUtils.nextInt(2) == 0) {
        branchMap.set(iP.getNumber(), branchMap.get(jP.getNumber()), true);
    } else {
        branchMap.set(iP.getNumber(), branchMap.get(j.getNumber()), true);
    }
    if (DEBUG) {
        c2cLikelihood.getTreeModel().checkPartitions();
    }
}
Also used : TreeModel(dr.evomodel.tree.TreeModel) NodeRef(dr.evolution.tree.NodeRef) BranchMapModel(dr.evomodel.epidemiology.casetocase.BranchMapModel) AbstractCase(dr.evomodel.epidemiology.casetocase.AbstractCase)

Example 3 with AbstractCase

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

the class TransmissionSubtreeSlideA method doOperation.

/**
     * Do a probabilistic subtree slide move.
     *
     * @return the log-transformed hastings ratio
     */
public double doOperation() {
    if (DEBUG) {
        c2cLikelihood.outputTreeToFile("beforeTSSA.nex", false);
    }
    BranchMapModel branchMap = c2cLikelihood.getBranchMap();
    double logq = 0;
    NodeRef i;
    // 1. choose a random eligible node
    ArrayList<NodeRef> eligibleNodes = getEligibleNodes(tree, branchMap);
    i = eligibleNodes.get(MathUtils.nextInt(eligibleNodes.size()));
    double eligibleNodeCount = eligibleNodes.size();
    final NodeRef iP = tree.getParent(i);
    final NodeRef CiP = getOtherChild(tree, iP, i);
    final NodeRef PiP = tree.getParent(iP);
    // 2. choose a delta to move
    final double delta = getDelta();
    final double oldHeight = tree.getNodeHeight(iP);
    final double newHeight = oldHeight + delta;
    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 (resampleInfectionTimes) {
        if (iCase != iPCase) {
            iCase.setInfectionBranchPosition(MathUtils.nextDouble());
        }
        // what happens between PiP and CiP
        if (PiPCase == null || CiPCase != PiPCase) {
            CiPCase.setInfectionBranchPosition(MathUtils.nextDouble());
        }
    }
    // 3. if the move is down
    if (delta > 0) {
        // 3.1 if the topology will change
        if (PiP != null && tree.getNodeHeight(PiP) < newHeight) {
            // find new parent
            NodeRef newParent = PiP;
            NodeRef newChild = iP;
            while (tree.getNodeHeight(newParent) < newHeight) {
                newChild = newParent;
                newParent = tree.getParent(newParent);
                if (newParent == null)
                    break;
            }
            // if the parent has slid out of its partition
            if (branchMap.get(newChild.getNumber()) != branchMap.get(iP.getNumber())) {
                //                    throw new OperatorFailedException("invalid slide");
                return Double.NEGATIVE_INFINITY;
            }
            // if iP is now the earliest node in its subtree
            tree.beginTreeEdit();
            // 3.1.1 if creating a new root
            if (tree.isRoot(newChild)) {
                tree.removeChild(iP, CiP);
                tree.removeChild(PiP, iP);
                tree.addChild(iP, newChild);
                tree.addChild(PiP, CiP);
                tree.setRoot(iP);
                if (tree.hasNodeTraits()) {
                    // **********************************************
                    // swap traits and rates so that root keeps it trait and rate values
                    // **********************************************
                    tree.swapAllTraits(newChild, iP);
                }
                if (tree.hasRates()) {
                    final double rootNodeRate = tree.getNodeRate(newChild);
                    tree.setNodeRate(newChild, tree.getNodeRate(iP));
                    tree.setNodeRate(iP, rootNodeRate);
                }
            // **********************************************
            } else // 3.1.2 no new root
            {
                tree.removeChild(iP, CiP);
                tree.removeChild(PiP, iP);
                tree.removeChild(newParent, newChild);
                tree.addChild(iP, newChild);
                tree.addChild(PiP, CiP);
                tree.addChild(newParent, iP);
            //System.err.println("No new root!");
            }
            tree.setNodeHeight(iP, newHeight);
            tree.endTreeEdit();
            // 3.1.3 count the hypothetical sources of this destination.
            final int possibleSources = intersectingEdges(tree, newChild, oldHeight, branchMap, branchMap.get(iP.getNumber()), null);
            //System.out.println("possible sources = " + possibleSources);
            logq -= Math.log(possibleSources);
        } else {
            // just change the node height
            tree.setNodeHeight(iP, newHeight);
            logq += 0.0;
        }
    } else // 4 if we are sliding the subtree up.
    {
        // 4.0 is it a valid move?
        if (tree.getNodeHeight(i) > newHeight) {
            return Double.NEGATIVE_INFINITY;
        }
        // 4.1 will the move change the topology
        if (tree.getNodeHeight(CiP) > newHeight) {
            List<NodeRef> newChildren = new ArrayList<NodeRef>();
            final int possibleDestinations = intersectingEdges(tree, CiP, newHeight, branchMap, branchMap.get(iP.getNumber()), newChildren);
            // if no valid destinations then return a failure
            if (newChildren.size() == 0) {
                return Double.NEGATIVE_INFINITY;
            }
            // pick a random parent/child destination edge uniformly from options
            final int childIndex = MathUtils.nextInt(newChildren.size());
            NodeRef newChild = newChildren.get(childIndex);
            NodeRef newParent = tree.getParent(newChild);
            if (resampleInfectionTimes) {
                AbstractCase newChildCase = branchMap.get(newChild.getNumber());
                if (newChildCase != iPCase) {
                    newChildCase.setInfectionBranchPosition(MathUtils.nextDouble());
                }
            }
            tree.beginTreeEdit();
            // 4.1.1 if iP was root
            if (tree.isRoot(iP)) {
                // new root is CiP
                tree.removeChild(iP, CiP);
                tree.removeChild(newParent, newChild);
                tree.addChild(iP, newChild);
                tree.addChild(newParent, iP);
                tree.setRoot(CiP);
                if (tree.hasNodeTraits()) {
                    // **********************************************
                    // swap traits and rates, so that root keeps it trait and rate values
                    // **********************************************
                    tree.swapAllTraits(iP, CiP);
                }
                if (tree.hasRates()) {
                    final double rootNodeRate = tree.getNodeRate(iP);
                    tree.setNodeRate(iP, tree.getNodeRate(CiP));
                    tree.setNodeRate(CiP, rootNodeRate);
                }
            // **********************************************
            //System.err.println("DOWN: Creating new root!");
            } else {
                tree.removeChild(iP, CiP);
                tree.removeChild(PiP, iP);
                tree.removeChild(newParent, newChild);
                tree.addChild(iP, newChild);
                tree.addChild(PiP, CiP);
                tree.addChild(newParent, iP);
            //System.err.println("DOWN: no new root!");
            }
            tree.setNodeHeight(iP, newHeight);
            tree.endTreeEdit();
            logq += Math.log(possibleDestinations);
        } else {
            tree.setNodeHeight(iP, newHeight);
            logq += 0.0;
        }
    }
    if (swapInRandomRate) {
        final NodeRef j = tree.getNode(MathUtils.nextInt(tree.getNodeCount()));
        if (j != i) {
            final double tmp = tree.getNodeRate(i);
            tree.setNodeRate(i, tree.getNodeRate(j));
            tree.setNodeRate(j, tmp);
        }
    }
    if (swapInRandomTrait) {
        final NodeRef j = tree.getNode(MathUtils.nextInt(tree.getNodeCount()));
        if (j != i) {
            tree.swapAllTraits(i, j);
        //                final double tmp = tree.getNodeTrait(i, TRAIT);
        //                tree.setNodeTrait(i, TRAIT, tree.getNodeTrait(j, TRAIT));
        //                tree.setNodeTrait(j, TRAIT, tmp);
        }
    }
    //        if (logq == Double.NEGATIVE_INFINITY) throw new OperatorFailedException("invalid slide");
    if (!Double.isInfinite(logq)) {
        if (DEBUG) {
            c2cLikelihood.getTreeModel().checkPartitions();
            c2cLikelihood.outputTreeToFile("afterTSSA.nex", false);
        }
        double reverseEligibleNodeCount = getEligibleNodes(tree, branchMap).size();
        logq += Math.log(eligibleNodeCount / reverseEligibleNodeCount);
    }
    return logq;
}
Also used : NodeRef(dr.evolution.tree.NodeRef) ArrayList(java.util.ArrayList) BranchMapModel(dr.evomodel.epidemiology.casetocase.BranchMapModel) AbstractCase(dr.evomodel.epidemiology.casetocase.AbstractCase)

Example 4 with AbstractCase

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

the class TransmissionExchangeOperatorB method exchange.

public double exchange() {
    TreeModel tree = c2cLikelihood.getTreeModel();
    BranchMapModel branchMap = c2cLikelihood.getBranchMap();
    final int nodeCount = tree.getNodeCount();
    final NodeRef root = tree.getRoot();
    NodeRef i = root;
    NodeRef iP = tree.getParent(i);
    boolean partitionsMatch = true;
    while (root == i || partitionsMatch) {
        i = tree.getNode(MathUtils.nextInt(nodeCount));
        iP = tree.getParent(i);
        partitionsMatch = i == root || branchMap.get(i.getNumber()) == branchMap.get(iP.getNumber());
    }
    ArrayList<NodeRef> candidates = getPossibleExchanges(tree, i);
    int candidateCount = candidates.size();
    if (candidateCount == 0) {
        //            throw new OperatorFailedException("No valid exchanges for this node");
        return Double.NEGATIVE_INFINITY;
    }
    NodeRef j = candidates.get(MathUtils.nextInt(candidates.size()));
    int jFirstCandidateCount = getPossibleExchanges(tree, j).size();
    double HRDenom = (1 / ((double) candidateCount)) + (1 / ((double) jFirstCandidateCount));
    NodeRef jP = tree.getParent(j);
    if (resampleInfectionTimes) {
        AbstractCase iCase = branchMap.get(i.getNumber());
        AbstractCase jCase = branchMap.get(j.getNumber());
        iCase.setInfectionBranchPosition(MathUtils.nextDouble());
        jCase.setInfectionBranchPosition(MathUtils.nextDouble());
    }
    exchangeNodes(tree, i, j, iP, jP);
    ArrayList<NodeRef> reverseCandidatesIfirst = getPossibleExchanges(tree, i);
    ArrayList<NodeRef> reverseCandidatesJfirst = getPossibleExchanges(tree, j);
    double HRNum = (1 / ((double) reverseCandidatesIfirst.size())) + (1 / ((double) reverseCandidatesJfirst.size()));
    return Math.log(HRNum / HRDenom);
}
Also used : TreeModel(dr.evomodel.tree.TreeModel) NodeRef(dr.evolution.tree.NodeRef) BranchMapModel(dr.evomodel.epidemiology.casetocase.BranchMapModel) AbstractCase(dr.evomodel.epidemiology.casetocase.AbstractCase)

Example 5 with AbstractCase

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

the class TransmissionSubtreeSlideB method doOperation.

/**
     * Do a probablistic subtree slide move.
     *
     * @return the log-transformed hastings ratio
     */
public double doOperation() {
    if (DEBUG) {
        c2cLikelihood.outputTreeToFile("beforeTSSB.nex", false);
    }
    BranchMapModel branchMap = c2cLikelihood.getBranchMap();
    double logq;
    NodeRef i;
    // 1. choose a random eligible node avoiding root
    do {
        i = tree.getNode(MathUtils.nextInt(tree.getNodeCount()));
    } while (!eligibleForMove(i, tree, branchMap));
    final NodeRef iP = tree.getParent(i);
    final NodeRef CiP = getOtherChild(tree, iP, i);
    final NodeRef PiP = tree.getParent(iP);
    // 2. choose a delta to move
    final double delta = getDelta();
    final double oldHeight = tree.getNodeHeight(iP);
    final double newHeight = oldHeight + delta;
    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 (resampleInfectionTimes) {
        // what happens on i's branch (there has always been a change)
        iCase.setInfectionBranchPosition(MathUtils.nextDouble());
        if (PiPCase == null || CiPCase != PiPCase) {
            CiPCase.setInfectionBranchPosition(MathUtils.nextDouble());
        }
    }
    // 3. if the move is down
    if (delta > 0) {
        // 3.1 if the topology will change
        if (PiP != null && tree.getNodeHeight(PiP) < newHeight) {
            // find new parent
            NodeRef newParent = PiP;
            NodeRef newChild = iP;
            while (tree.getNodeHeight(newParent) < newHeight) {
                newChild = newParent;
                newParent = tree.getParent(newParent);
                if (newParent == null)
                    break;
            }
            tree.beginTreeEdit();
            // 3.1.1 if creating a new root
            if (tree.isRoot(newChild)) {
                tree.removeChild(iP, CiP);
                tree.removeChild(PiP, iP);
                tree.addChild(iP, newChild);
                tree.addChild(PiP, CiP);
                tree.setRoot(iP);
                if (tree.hasNodeTraits()) {
                    // **********************************************
                    // swap traits and rates so that root keeps it trait and rate values
                    // **********************************************
                    tree.swapAllTraits(newChild, iP);
                }
                if (tree.hasRates()) {
                    final double rootNodeRate = tree.getNodeRate(newChild);
                    tree.setNodeRate(newChild, tree.getNodeRate(iP));
                    tree.setNodeRate(iP, rootNodeRate);
                }
            // **********************************************
            } else // 3.1.2 no new root
            {
                tree.removeChild(iP, CiP);
                tree.removeChild(PiP, iP);
                tree.removeChild(newParent, newChild);
                tree.addChild(iP, newChild);
                tree.addChild(PiP, CiP);
                tree.addChild(newParent, iP);
            //System.err.println("No new root!");
            }
            tree.setNodeHeight(iP, newHeight);
            tree.endTreeEdit();
            // 3.1.3 count the hypothetical sources of this destination.
            final int possibleSources = intersectingEdges(tree, newChild, oldHeight, null);
            logq = -Math.log(possibleSources);
            if (PiPCase != CiPCase) {
                logq += Math.log(0.5);
            }
            AbstractCase newiPCase;
            AbstractCase newChildCase = branchMap.get(newChild.getNumber());
            if (newParent != null && branchMap.get(newParent.getNumber()) != branchMap.get(newChild.getNumber())) {
                if (MathUtils.nextInt(2) == 0) {
                    newiPCase = branchMap.get(newParent.getNumber());
                } else {
                    newiPCase = newChildCase;
                }
                if (resampleInfectionTimes) {
                    //whichever we picked for iP, it's the new child's case whose infection branch is modified
                    // (even if this infection branch is iP's branch)
                    newChildCase.setInfectionBranchPosition(MathUtils.nextDouble());
                }
                logq += Math.log(2);
            } else {
                newiPCase = newChildCase;
                if (resampleInfectionTimes) {
                    if (newParent == null) {
                        newChildCase.setInfectionBranchPosition(MathUtils.nextDouble());
                    }
                }
            }
            branchMap.set(iP.getNumber(), newiPCase, true);
        } else {
            // just change the node height
            // todo you could actually randomise whether the subtree containing iP is changed here
            tree.setNodeHeight(iP, newHeight);
            logq = 0.0;
        }
    } else // 4 if we are sliding the subtree up.
    {
        // 4.0 is it a valid move?
        if (tree.getNodeHeight(i) > newHeight) {
            return Double.NEGATIVE_INFINITY;
        }
        // 4.1 will the move change the topology?
        if (tree.getNodeHeight(CiP) > newHeight) {
            List<NodeRef> newChildren = new ArrayList<NodeRef>();
            final int possibleDestinations = intersectingEdges(tree, CiP, newHeight, newChildren);
            // if no valid destinations then return a failure
            if (newChildren.size() == 0) {
                return Double.NEGATIVE_INFINITY;
            }
            // pick a random parent/child destination edge uniformly from options
            final int childIndex = MathUtils.nextInt(newChildren.size());
            NodeRef newChild = newChildren.get(childIndex);
            NodeRef newParent = tree.getParent(newChild);
            tree.beginTreeEdit();
            // 4.1.1 if iP was root
            if (tree.isRoot(iP)) {
                // new root is CiP
                tree.removeChild(iP, CiP);
                tree.removeChild(newParent, newChild);
                tree.addChild(iP, newChild);
                tree.addChild(newParent, iP);
                tree.setRoot(CiP);
                if (tree.hasNodeTraits()) {
                    // **********************************************
                    // swap traits and rates, so that root keeps it trait and rate values
                    // **********************************************
                    tree.swapAllTraits(iP, CiP);
                }
                if (tree.hasRates()) {
                    final double rootNodeRate = tree.getNodeRate(iP);
                    tree.setNodeRate(iP, tree.getNodeRate(CiP));
                    tree.setNodeRate(CiP, rootNodeRate);
                }
            // **********************************************
            //System.err.println("DOWN: Creating new root!");
            } else {
                tree.removeChild(iP, CiP);
                tree.removeChild(PiP, iP);
                tree.removeChild(newParent, newChild);
                tree.addChild(iP, newChild);
                tree.addChild(PiP, CiP);
                tree.addChild(newParent, iP);
            //System.err.println("DOWN: no new root!");
            }
            tree.setNodeHeight(iP, newHeight);
            tree.endTreeEdit();
            logq = Math.log(possibleDestinations);
            if (PiP != null && PiPCase != CiPCase) {
                logq += Math.log(0.5);
            }
            AbstractCase newiPCase;
            AbstractCase newChildCase = branchMap.get(newChild.getNumber());
            if (branchMap.get(newParent.getNumber()) != branchMap.get(newChild.getNumber())) {
                if (MathUtils.nextInt(2) == 0) {
                    newiPCase = branchMap.get(newParent.getNumber());
                } else {
                    newiPCase = newChildCase;
                }
                if (resampleInfectionTimes) {
                    //whichever we picked for iP, it's the new child's case whose infection branch is modified
                    // (even if this infection branch is iP's branch)
                    newChildCase.setInfectionBranchPosition(MathUtils.nextDouble());
                }
                logq += Math.log(2);
            } else {
                //upward, so don't have to worry about newParent being the root if the topology changed
                newiPCase = newChildCase;
            }
            branchMap.set(iP.getNumber(), newiPCase, true);
        } else {
            tree.setNodeHeight(iP, newHeight);
            logq = 0.0;
        }
    }
    if (swapInRandomRate) {
        final NodeRef j = tree.getNode(MathUtils.nextInt(tree.getNodeCount()));
        if (j != i) {
            final double tmp = tree.getNodeRate(i);
            tree.setNodeRate(i, tree.getNodeRate(j));
            tree.setNodeRate(j, tmp);
        }
    }
    if (swapInRandomTrait) {
        final NodeRef j = tree.getNode(MathUtils.nextInt(tree.getNodeCount()));
        if (j != i) {
            tree.swapAllTraits(i, j);
        //                final double tmp = tree.getNodeTrait(i, TRAIT);
        //                tree.setNodeTrait(i, TRAIT, tree.getNodeTrait(j, TRAIT));
        //                tree.setNodeTrait(j, TRAIT, tmp);
        }
    }
    if (DEBUG) {
        c2cLikelihood.getTreeModel().checkPartitions();
        c2cLikelihood.outputTreeToFile("afterTSSB.nex", false);
    }
    return logq;
}
Also used : NodeRef(dr.evolution.tree.NodeRef) ArrayList(java.util.ArrayList) BranchMapModel(dr.evomodel.epidemiology.casetocase.BranchMapModel) AbstractCase(dr.evomodel.epidemiology.casetocase.AbstractCase)

Aggregations

NodeRef (dr.evolution.tree.NodeRef)7 AbstractCase (dr.evomodel.epidemiology.casetocase.AbstractCase)7 BranchMapModel (dr.evomodel.epidemiology.casetocase.BranchMapModel)7 TreeModel (dr.evomodel.tree.TreeModel)4 ArrayList (java.util.ArrayList)3 PartitionedTreeModel (dr.evomodel.epidemiology.casetocase.PartitionedTreeModel)1 ExchangeOperator (dr.evomodel.operators.ExchangeOperator)1 SubtreeSlideOperator (dr.evomodel.operators.SubtreeSlideOperator)1 WilsonBalding (dr.evomodel.operators.WilsonBalding)1 HashSet (java.util.HashSet)1