extending the epochal discrete CTMC model

src/dr/app/beagle/tools/Partition.java

@@ -150,7 +150,7 @@ public Partition(TreeModel treeModel, //
 
 	private void setSubstitutionModelDelegate() {
 		substitutionModelDelegate = new SubstitutionModelDelegate(treeModel,
-				branchModel);
+				branchModel, branchRateModel);
 	}// END: setSubstitutionModelDelegate
 
 	private void setBufferHelpers() {

src/dr/evomodel/arg/ARGRelaxedClock.java

@@ -93,6 +93,20 @@ public double getBranchRate(Tree tree, NodeRef nodeRef) {
 
         return globalRateParameter.getParameterValue(0) * argNode.getRate(partition);
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     public static XMLObjectParser PARSER = new AbstractXMLObjectParser() {
 

src/dr/evomodel/arg/branchratemodel/ARGDiscretizedBranchRates.java

@@ -224,6 +224,20 @@ public double getBranchRate(Tree tree, NodeRef node) {
 //	    System.err.println("rate = "+rates[rateCategory]+" : "+rateCategory);
         return rates[rateCategory];
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     /**
      * Calculates the actual rates corresponding to the category indices.

src/dr/evomodel/branchmodel/EpochBranchModel.java

@@ -95,7 +95,7 @@ public Mapping getBranchModelMapping(NodeRef node) {
         double currentHeight = nodeHeight;
 
         // find the epoch that the parent height is in...
-        while (epoch < epochCount && parentHeight >= transitionTimes[epoch]) {
+        while (epoch < epochCount && parentHeight > transitionTimes[epoch]) {
             weightList.add( transitionTimes[epoch] - currentHeight );
             orderList.add(epoch);
 

src/dr/evomodel/branchratemodel/AncestralTraitBranchRateModel.java

@@ -61,6 +61,20 @@ public double getBranchRate(Tree tree, NodeRef node) {
             return 1.0;
         }
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     @Override
     protected void handleModelChangedEvent(Model model, Object object, int index) {

src/dr/evomodel/branchratemodel/ArbitraryBranchRates.java

@@ -112,6 +112,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
 
         return transform.transform(getUntransformedBranchRate(tree, node), tree, node);
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     double getUntransformedBranchRate(final Tree tree, final NodeRef node) {
         return rates.getNodeValue(tree, node);

src/dr/evomodel/branchratemodel/AttributeBranchRateModel.java

@@ -84,6 +84,20 @@ public double getBranchRate(Tree tree, NodeRef node) {
         Object value = tree.getNodeAttribute(node, rateAttributeName);
         return Double.parseDouble((String)value);
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     @Override
     public String getTraitName() {

src/dr/evomodel/branchratemodel/BranchRateModel.java

@@ -41,4 +41,29 @@ public interface BranchRateModel extends Model, BranchRates, TreeTraitProvider,
 
     // This is inherited from BranchRates:
     // double getBranchRate(Tree tree, NodeRef node);
+
+    /**
+     * Returns a mapping of branch rate models to the given branch. The Mapping
+     * contains a list of branch rates in order from rootward to tipward
+     * and a set of relative weights for each (may be times or proportions).
+     *
+     * @param branch the branch
+     * @return a Mapping object
+     */
+    Mapping getBranchRateModelMapping(final Tree tree, final NodeRef branch);
+
+    interface Mapping {
+        double[] getRates();
+        double[] getWeights();
+    }
+
+    Mapping DEFAULT = new Mapping() {
+        public double[] getRates() {
+            return new double[] { 1.0 };
+        }
+
+        public double[] getWeights() {
+            return new double[] { 1.0 };
+        }
+    };
 }

src/dr/evomodel/branchratemodel/CompoundBranchRateModel.java

@@ -85,5 +85,19 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
         }
         return rate;
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
 }

src/dr/evomodel/branchratemodel/ContinuousBranchRates.java

@@ -143,6 +143,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
 
         return rates[node.getNumber()] * scaleFactor;
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     /**
      * Calculates the actual rates corresponding to the quantiles.

src/dr/evomodel/branchratemodel/ContinuousTraitBranchRateModel.java

@@ -139,5 +139,19 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
 
         return rate;
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
 }

src/dr/evomodel/branchratemodel/CountableMixtureBranchRates.java

@@ -406,6 +406,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
         }
         return effect;
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     private final Helper helper = new Helper();
 

src/dr/evomodel/branchratemodel/CountableModelMixtureBranchRates.java

@@ -198,6 +198,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
         }
         return effect;
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     private final CountableBranchCategoryProvider rateCategories;
     private final boolean modelInLogSpace;

src/dr/evomodel/branchratemodel/DecayingRateModel.java

@@ -129,6 +129,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
 
         return rates[node.getNumber()];
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     /**
      * Traverse the tree calculating partial likelihoods.

src/dr/evomodel/branchratemodel/DefaultBranchRateModel.java

@@ -38,6 +38,10 @@ public final class DefaultBranchRateModel implements BranchRateModel {
     public double getBranchRate(Tree tree, NodeRef node) {
         return 1.0;
     }
+
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+        return DEFAULT;
+    }
 
     public void addModelListener(ModelListener listener) {
         // nothing to do

src/dr/evomodel/branchratemodel/DiscreteTraitBranchRateModel.java

@@ -288,6 +288,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
             return getRawBranchRate(tree, node);
         }
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     // produce weighted mean of rate for a branch
     // rate = absRate * branchWeight[0] * relativeRates[0] + absRate * branchWeight[1] * relativeRates[1]

src/dr/evomodel/branchratemodel/DiscretizedBranchRates.java

@@ -307,6 +307,20 @@ public final double getBranchRate(final Tree tree, final NodeRef node) {
         //System.out.println(rates[rateCategory] + "\t"  + rateCategory);
         return rates[currentRateArrayIndex][rateCategory] * scaleFactor;
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     @SuppressWarnings("WeakerAccess")
     public final int getBranchRateCategory(final Tree tree, final NodeRef node) {

src/dr/evomodel/branchratemodel/FixedDriftModel.java

@@ -111,6 +111,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
             return otherDrift.getParameterValue(0);
         }
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     /*
     public double getBranchRate(final Tree tree, final NodeRef node) {

src/dr/evomodel/branchratemodel/LatentStateBranchRateModel.java

@@ -184,6 +184,20 @@ public double getBranchRate(Tree tree, NodeRef node) {
 
         return calculateBranchRate(nonLatentRate, latentProportion);
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     public double getLatentProportion(Tree tree, NodeRef node) {
 

src/dr/evomodel/branchratemodel/LocalBranchRates.java

@@ -157,6 +157,20 @@ public void handleModelChangedEvent(Model model, Object object, int index) {
     public double getBranchRate(final Tree tree, final NodeRef node) {
         return transform.transform(branchRates.getNodeValue(tree, node), tree, node);
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        return new Mapping() {
+            public double[] getRates() {
+                return new double[] { getBranchRate(tree, node) };
+            }
+
+            public double[] getWeights() {
+                return new double[] { 1.0 };
+            }
+        };
+    }
 
     @Override
     public void setBranchRate(Tree tree, NodeRef node, double value) {