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/branchmodel/EpochBranchModel.java

@@ -92,7 +92,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/AbstractBranchRateModel.java

@@ -84,4 +84,17 @@ public double getLogLikelihood() {
     public void makeDirty() {
         // Do nothing
     }
+
+    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/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/DefaultBranchRateModel.java

@@ -41,6 +41,10 @@ public final class DefaultBranchRateModel implements BranchRateModel, Differenti
     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/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/RateEpochBranchRateModel.java

@@ -32,6 +32,8 @@
 import dr.inference.model.Parameter;
 import dr.inference.model.Variable;
 
+import java.util.*;
+
 /**
  * Implements a model where time is broken into 'epochs' each with a different but
  * constant rate. Parameters can be used to sample transition times but it is up
@@ -120,6 +122,69 @@ public double getBranchRate(final Tree tree, final NodeRef node) {
         }
         throw new IllegalArgumentException("root node doesn't have a rate!");
     }
+
+    @Override
+    public Mapping getBranchRateModelMapping(final Tree tree, final NodeRef node) {
+
+        NodeRef parent = tree.getParent(node);
+
+        List<Double> weightList = new ArrayList<Double>();
+        List<Double> rateList = new ArrayList<Double>();
+
+        if (parent != null) {
+            double height0 = tree.getNodeHeight(node);
+            double height1 = tree.getNodeHeight(parent);
+            int i = 0;
+
+            double rate = 0.0;
+            double lastHeight = height0;
+
+            // First find the epoch which contains the node height
+            while (i < timeParameters.length && height0 >= timeParameters[i].getParameterValue(0)) {
+                i++;
+            }
+
+            // Now walk up the branch until we reach the last epoch or the height of the parent
+            while (i < timeParameters.length && height1 > timeParameters[i].getParameterValue(0)) {
+                // insert each epoch that this branch overlaps with to the list,
+                // ordered from parent (rootward) to child (tipward),
+                // as the transition probability matrix of a given branch is the product
+                // of matrices multiplying from parent to child
+                weightList.add( 0, timeParameters[i].getParameterValue(0) - lastHeight );
+                rateList.add( 0, rateParameters[i].getParameterValue(0) );
+                lastHeight = timeParameters[i].getParameterValue(0);
+                i++;
+            }
+
+            // Add that last rate segment
+            weightList.add( 0, height1 - lastHeight );
+            rateList.add( 0, rateParameters[i].getParameterValue(0) );
+
+        } else {
+            throw new IllegalArgumentException("root node doesn't have a rate!");
+        }
+
+        if (weightList.size() == 0) {
+            throw new RuntimeException("RateEpochBranchRateModel failed to give a valid mapping");
+        }
+
+        final double[] rates = new double[rateList.size()];
+        final double[] weights = new double[weightList.size()];
+        for (int i = 0; i < weightList.size(); i++) {
+            rates[i] = rateList.get(i);
+            weights[i] = weightList.get(i);
+        }
+
+        return new Mapping() {
+            public double[] getRates() {
+                return rates;
+            }
+
+            public double[] getWeights() {
+                return weights;
+            }
+        };
+    }
 
     protected double normalizeRate(double rate) {
         return rate;

src/dr/evomodel/branchratemodel/SericolaLatentStateBranchRateModel.java

@@ -183,6 +183,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/substmodel/UniformizedSubstitutionModel.java

@@ -128,16 +128,20 @@ public String getCompleteHistory() {
         return getCompleteHistory(null, null);
     }
 
-   public String getCompleteHistory(Double newStartTime, Double newEndTime) {
+    public String getCompleteHistory(Double newStartTime, Double newEndTime) {
         return getCompleteHistory(-1, newStartTime, newEndTime);
-   }
+    }
 
     public String getCompleteHistory(int site, Double newStartTime, Double newEndTime) {
+        return getCompleteHistory(site, newStartTime, newEndTime, true);
+    }
+
+    public String getCompleteHistory(int site, Double newStartTime, Double newEndTime, boolean wrap) {
         if (newStartTime != null && newEndTime != null) {
             // Rescale time of events
             completeHistory.rescaleTimesOfEvents(newStartTime, newEndTime);
         }
-        return completeHistory.toStringChanges(site, dataType); //, 0.0);
+        return completeHistory.toStringChanges(site, dataType, wrap); //, 0.0);
     }
 
     public int getNumberOfJumpsInCompleteHistory() {