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

Example 1 with BranchMapModel

use of dr.evomodel.epidemiology.casetocase.BranchMapModel 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;
}

Example 2 with BranchMapModel

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

the class TransmissionExchangeOperatorB method getPossibleExchanges.

// get a set of candidates for an exchange of a given node. Does not include the node itself or its sibling, so
// the check for a failed move will be if this set is of size 0
public ArrayList<NodeRef> getPossibleExchanges(TreeModel tree, NodeRef node) {
    BranchMapModel map = c2cLikelihood.getBranchMap();
    ArrayList<NodeRef> out = new ArrayList<NodeRef>();
    NodeRef parent = tree.getParent(node);
    if (parent == null) {
        throw new RuntimeException("Can't exchange the root node");
    }
    if (map.get(parent.getNumber()) == map.get(node.getNumber())) {
        throw new RuntimeException("This node is not exchangeable by this operator");
    }
    for (NodeRef candidate : tree.getNodes()) {
        NodeRef newParent = tree.getParent(candidate);
        if (newParent != parent && newParent != null) {
            if (candidate != parent && node != newParent && tree.getNodeHeight(candidate) < tree.getNodeHeight(parent) && tree.getNodeHeight(node) < tree.getNodeHeight(newParent) && map.get(newParent.getNumber()) != map.get(candidate.getNumber())) {
                if (out.contains(candidate) || candidate == node) {
                    throw new RuntimeException("Adding a candidate that already exists in the list or" + " the node itself");
                }
                out.add(candidate);
            }
        }
    }
    return out;
}

Example 3 with BranchMapModel

use of dr.evomodel.epidemiology.casetocase.BranchMapModel 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()) {
        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());
        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();
    }
}

Example 4 with BranchMapModel

use of dr.evomodel.epidemiology.casetocase.BranchMapModel 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);
}

Example 5 with BranchMapModel

use of dr.evomodel.epidemiology.casetocase.BranchMapModel 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);
}