Improve two operators, and provide a starting tree generator

src/beast/evolution/operators/LeafToSampledAncestorJump.java

@@ -1,9 +1,14 @@
 package beast.evolution.operators;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
 import beast.core.Description;
 import beast.core.Input;
 import beast.core.parameter.IntegerParameter;
 import beast.core.parameter.RealParameter;
+import beast.evolution.alignment.Taxon;
 import beast.evolution.tree.Node;
 import beast.evolution.tree.Tree;
 import beast.util.Randomizer;
@@ -24,15 +29,46 @@ public class LeafToSampledAncestorJump extends TreeOperator {
 
     public Input<RealParameter> rInput =
             new Input<RealParameter>("removalProbability", "The probability of an individual to be removed from the process immediately after the sampling");
-
+    public Input<List<Taxon>> sampledTaxa =
+    		new Input<List<Taxon>>(
+    				"sampledTaxa",
+    				"Taxa that this operator should be allowed to let jump between sampled ancestor and leaf. Default: All non-recent leaves.",
+    				new ArrayList<>());
+
+    protected List<Integer> validLeaves = new ArrayList<Integer>();
+
     @Override
     public void initAndValidate() {
+    	if (sampledTaxa.get().size() == 0) {
+    		validLeaves = new ArrayList<Integer>(treeInput.get().getLeafNodeCount());
+            for (Node leaf: treeInput.get().getExternalNodes()) {
+            	if (leaf.getHeight() > 1e-6) {
+            		validLeaves.add(leaf.getNr());
+            	}
+            }
+    	} else {
+    		List<Taxon> taxa = sampledTaxa.get();
+    		List<String> taxaNames = new ArrayList<String>(taxa.size());
+    		for (Taxon taxon: taxa) {
+    			taxaNames.add(taxon.getID());
+    		}
+        	validLeaves = new ArrayList<Integer>(taxa.size());
+        	Integer i = 0;
+            for (String leaf: treeInput.get().getTaxaNames()) {
+            	if (taxaNames.contains(leaf)) {
+            		validLeaves.add(i);
+            	}
+        		i += 1;
+            }
+    	}
+    	// System.out.println("Nodes to be jumped:");
+    	// System.out.println(Arrays.toString(validLeaves.toArray()));
     }
 
     @Override
     public double proposal() {
 
-        double newHeight, newRange, oldRange;
+        double newHeight, logNewRange, logOldRange;
         int categoryCount = 1;
         if (categoriesInput.get() != null) {
 
@@ -41,20 +77,21 @@ public double proposal() {
 
         Tree tree = treeInput.get();
 
-        int leafNodeCount = tree.getLeafNodeCount();
+        int leafNodeCount = validLeaves.size();
 
-        Node leaf = tree.getNode(Randomizer.nextInt(leafNodeCount));
+        Node leaf = tree.getNode(validLeaves.get(Randomizer.nextInt(leafNodeCount)));
         Node parent = leaf.getParent();
 
         if (leaf.isDirectAncestor()) {
-            oldRange = (double) 1;
+            logOldRange = (double) 0;
             if (parent.isRoot()) {
                 final double randomNumber = Randomizer.nextExponential(1);
                 newHeight = parent.getHeight() + randomNumber;
-                newRange = Math.exp(randomNumber);
+                logNewRange = randomNumber;
             } else {
-                newRange = parent.getParent().getHeight() - parent.getHeight();
+                double newRange = parent.getParent().getHeight() - parent.getHeight();
                 newHeight = parent.getHeight() + Randomizer.nextDouble() * newRange;
+                logNewRange = Math.log(newRange);
             }
 
             if (categoriesInput.get() != null) {
@@ -63,15 +100,15 @@ public double proposal() {
                 categoriesInput.get().setValue(index, newValue);
             }
         } else {
-            newRange = (double) 1;
+            logNewRange = (double) 0;
             //make sure that the branch where a new sampled node to appear is not above that sampled node
             if (getOtherChild(parent, leaf).getHeight() >= leaf.getHeight())  {
                 return Double.NEGATIVE_INFINITY;
             }
             if (parent.isRoot()) {
-                oldRange = Math.exp(parent.getHeight() - leaf.getHeight());
+                logOldRange = parent.getHeight() - leaf.getHeight();
             } else {
-                oldRange = parent.getParent().getHeight() - leaf.getHeight();
+                logOldRange = Math.log(parent.getParent().getHeight() - leaf.getHeight());
             }
             newHeight = leaf.getHeight();
             if  (categoriesInput.get() != null) {
@@ -86,6 +123,6 @@ public double proposal() {
             return Double.NEGATIVE_INFINITY;
         }
 
-        return Math.log(newRange/oldRange);
+        return logNewRange - logOldRange;
     }
 }