adds checks to ConstrainedTreeOperator + smarter subtree selection

beast-dev · Jan 18, 2021 · 21be489 · 21be489
1 parent 6d6a1d2
commit 21be489
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git a/src/dr/evomodel/treelikelihood/thorneytreelikelihood/ConstrainedTreeModel.java b/src/dr/evomodel/treelikelihood/thorneytreelikelihood/ConstrainedTreeModel.java
@@ -158,8 +158,11 @@ private BitSet addBits(Tree referenceTree, Tree tree, NodeRef node, HashMap<BitS
         return bits;
     }
 
+    public int getSubtreeCount(){
+        return subtrees.size();
+    }
 
-    public WrappedSubtree getSubtree(NodeRef node,SubtreeContext context) {
+    public TreeModel getSubtree(NodeRef node,SubtreeContext context) {
         return subtrees.get(((Node) node).getSubtreeNumber(context));
     }
 
@@ -171,11 +174,11 @@ public WrappedSubtree getSubtree(NodeRef node,SubtreeContext context) {
      * @param node
      * @return
      */
-    public WrappedSubtree getSubtree(NodeRef node){
+    public TreeModel getSubtree(NodeRef node){
         return subtrees.get(((Node) node).getSubtreeNumber(SubtreeContext.IncludeRoot));
     }
 
-    public WrappedSubtree getSubtree(int i) {
+    public TreeModel getSubtree(int i) {
         return subtrees.get(i);
     }
 

diff --git a/src/dr/evomodel/treelikelihood/thorneytreelikelihood/ConstrainedTreeOperator.java b/src/dr/evomodel/treelikelihood/thorneytreelikelihood/ConstrainedTreeOperator.java
@@ -1,20 +1,28 @@
 package dr.evomodel.treelikelihood.thorneytreelikelihood;
 
-import dr.evolution.tree.NodeRef;
 import dr.evomodel.operators.AbstractTreeOperator;
 import dr.evomodel.tree.TreeModel;
-import dr.inference.operators.AdaptationMode;
 import dr.math.MathUtils;
 
 public class ConstrainedTreeOperator extends AbstractTreeOperator {
-    public ConstrainedTreeOperator(ConstrainedTreeModel tree,double weight,ConstrainableTreeOperator operator){
+    public ConstrainedTreeOperator(ConstrainedTreeModel tree, double weight, ConstrainableTreeOperator operator) {
         setWeight(weight);
         constrainedTreeModel = tree;
         this.operator = operator;
+        if (tree.getInternalNodeCount() == tree.getSubtreeCount()) {
+            throw new IllegalArgumentException(getOperatorName() + " is designed to resolve polytomies; however, the "+
+                    "constrained tree is fully resolved. Please remove this operator or provide an unresolved "+
+                    "constraints tree.");
+        }
+        subtreeSizes = new double[tree.getSubtreeCount()];
+        for (int i = 0; i < tree.getSubtreeCount(); i++) {
+            subtreeSizes[i] = (double) tree.getSubtree(i).getInternalNodeCount()-1; // don't choose subtrees with only 1 internal node. There's no topology to sample.
+        }
     }
+
     @Override
     public String getOperatorName() {
-        return "Constrained"+ operator.getOperatorName();
+        return "Constrained  " + operator.getOperatorName();
     }
 
     /**
@@ -24,16 +32,14 @@ public String getOperatorName() {
      */
     @Override
     public double doOperation() {
-            NodeRef node = constrainedTreeModel.getNode(MathUtils.nextInt(constrainedTreeModel.getNodeCount()));
-            TreeModel subtree = constrainedTreeModel.getSubtree(node);
 
-            while(subtree.getExternalNodeCount()<3){
-                node = constrainedTreeModel.getNode(MathUtils.nextInt(constrainedTreeModel.getNodeCount()));
-                subtree = constrainedTreeModel.getSubtree(node);
-            }
+        int subtreeIndex = MathUtils.randomChoicePDF(subtreeSizes);
+        TreeModel subtree = constrainedTreeModel.getSubtree(subtreeIndex);
         return operator.doOperation(subtree);
     }
+
     private final ConstrainedTreeModel constrainedTreeModel;
     private final ConstrainableTreeOperator operator;
+    private final double[] subtreeSizes;
 
 }