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;
}
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;
}
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);
}
Aggregations