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

Example 6 with BranchMapModel

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

Example 7 with BranchMapModel

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

the class TransmissionTreeOperator method doOperation.

public double doOperation() {
    TreeModel tree = c2cLikelihood.getTreeModel();
    BranchMapModel branchMap = c2cLikelihood.getBranchMap();
    AbstractCase[] newBranchMap = branchMap.getArrayCopy();
    int[] oldParents = getParentsArray(tree);
    double[] oldHeights = getHeightsArray(tree);
    double hr = innerOperator.doOperation();
    int[] newParents = getParentsArray(tree);
    ArrayList<Integer> changedNodes = new ArrayList<Integer>();
    for (int i = 0; i < tree.getNodeCount(); i++) {
        if (oldParents[i] != newParents[i]) {
            changedNodes.add(i);
        }
    }
    if (changedNodes.size() != 0) {
        if (innerOperator instanceof ExchangeOperator) {
            //this is a node swap operator
            AbstractCase[] nodePaintings = new AbstractCase[2];
            AbstractCase[] parentPaintings = new AbstractCase[2];
            for (int i = 0; i < 2; i++) {
                nodePaintings[i] = branchMap.get(changedNodes.get(i));
                parentPaintings[i] = branchMap.get(oldParents[changedNodes.get(i)]);
            }
            if (nodePaintings[0] == parentPaintings[0] || nodePaintings[1] == parentPaintings[1]) {
                //If this is not true there is nothing to do - the result is already a valid painting
                for (int i = 0; i < 2; i++) {
                    paintUp(tree, nodePaintings[1 - i], nodePaintings[i], branchMap, newBranchMap, tree.getNode(changedNodes.get(i)).getNumber(), newParents);
                }
            }
        } else if (innerOperator instanceof SubtreeSlideOperator || innerOperator instanceof WilsonBalding) {
            //this is a node transplantation operator
            int movedNode = -1;
            int oldChild = -1;
            int newChild = -1;
            for (int i = 0; i < tree.getNodeCount(); i++) {
                // exception later, though.
                if (tree.getNodeHeight(tree.getNode(i)) != oldHeights[i]) {
                    movedNode = i;
                    break;
                }
            }
            for (int j : changedNodes) {
                if (j != movedNode) {
                    if (tree.getParent(tree.getNode(j)) == tree.getNode(movedNode)) {
                        newChild = j;
                    } else {
                        oldChild = j;
                    }
                }
            }
            if (movedNode == -1 || oldChild == -1 || newChild == -1) {
                // is this a bug or should the move be rejected (i.e., return -Inf HR)?
                throw new RuntimeException("Failed to establish relationship between relocated node and others");
            }
            NodeRef movedNodeObject = tree.getNode(movedNode);
            NodeRef oldChildObject = tree.getNode(oldChild);
            NodeRef newChildObject = tree.getNode(newChild);
            int otherChild = -1;
            NodeRef otherChildObject;
            //Find the other child of the moved node (the root of the transplanted subtree)
            for (int i = 0; i < tree.getChildCount(movedNodeObject); i++) {
                if (tree.getChild(movedNodeObject, i) != newChildObject) {
                    otherChildObject = tree.getChild(movedNodeObject, i);
                    otherChild = otherChildObject.getNumber();
                }
            }
            //If the child of the moved node is the earliest node with its painting:
            if (branchMap.get(otherChild) != branchMap.get(movedNode)) {
                newBranchMap[movedNode] = branchMap.get(newChild);
            } else {
                //Change all paintings up the tree from the old child that used to match the moved node to match
                //the old child
                paintUp(tree, branchMap.get(movedNode), branchMap.get(oldChild), branchMap, newBranchMap, oldChild, oldParents);
                //This may have resulted in the moved node being recoloured wrong.
                newBranchMap[movedNode] = branchMap.get(movedNode);
                branchMap.setAll(newBranchMap, true);
                //Change all paintings up the tree from the moved node that used to match the new child to match
                //the moved node
                paintUp(tree, branchMap.get(newChild), branchMap.get(movedNode), branchMap, newBranchMap, movedNode, newParents);
            }
        } else {
            //I don't know what this is
            throw new UnsupportedOperationException("Operator class " + innerOperator.getOperatorName() + " not yet " + "supported");
        }
    }
    branchMap.setAll(newBranchMap, true);
    c2cLikelihood.makeDirty();
    return hr;
}

Example 8 with BranchMapModel

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

the class TransmissionExchangeOperatorA method exchange.

public double exchange() {
    TreeModel tree = c2cLikelihood.getTreeModel();
    final int nodeCount = tree.getNodeCount();
    final NodeRef root = tree.getRoot();
    NodeRef i = root;
    while (root == i) {
        i = tree.getNode(MathUtils.nextInt(nodeCount));
    }
    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 iP = tree.getParent(i);
    NodeRef jP = tree.getParent(j);
    if (resampleInfectionTimes) {
        BranchMapModel branchMap = c2cLikelihood.getBranchMap();
        AbstractCase iCase = branchMap.get(i.getNumber());
        AbstractCase jCase = branchMap.get(j.getNumber());
        AbstractCase parentCase = branchMap.get(iP.getNumber());
        if (iCase != parentCase) {
            iCase.setInfectionBranchPosition(MathUtils.nextDouble());
        }
        if (jCase != parentCase) {
            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);
}