From 94e91fe72ba4b7adb3de74e4a4023c5850dc020f Mon Sep 17 00:00:00 2001 From: Jiansi Gao Date: Fri, 4 Mar 2022 22:05:35 -0800 Subject: [PATCH 1/3] adding a getBranchRateModelMapping function to the epochal rate model --- src/dr/evomodel/arg/ARGRelaxedClock.java | 14 ++++ .../ARGDiscretizedBranchRates.java | 14 ++++ .../AncestralTraitBranchRateModel.java | 14 ++++ .../branchratemodel/ArbitraryBranchRates.java | 14 ++++ .../AttributeBranchRateModel.java | 14 ++++ .../branchratemodel/BranchRateModel.java | 25 +++++++ .../CompoundBranchRateModel.java | 14 ++++ .../ContinuousBranchRates.java | 14 ++++ .../ContinuousTraitBranchRateModel.java | 14 ++++ .../CountableMixtureBranchRates.java | 14 ++++ .../CountableModelMixtureBranchRates.java | 14 ++++ .../branchratemodel/DecayingRateModel.java | 14 ++++ .../DefaultBranchRateModel.java | 4 ++ .../DiscreteTraitBranchRateModel.java | 14 ++++ .../DiscretizedBranchRates.java | 14 ++++ .../branchratemodel/FixedDriftModel.java | 14 ++++ .../LatentStateBranchRateModel.java | 14 ++++ .../branchratemodel/LocalBranchRates.java | 14 ++++ .../branchratemodel/LocalClockModel.java | 14 ++++ .../MixtureModelBranchRates.java | 14 ++++ .../PassageBranchRateModel.java | 14 ++++ .../RandomLocalClockModel.java | 14 ++++ .../RateEpochBranchRateModel.java | 65 +++++++++++++++++++ .../branchratemodel/RelaxedDriftModel.java | 14 ++++ .../ScaledTreeLengthRateModel.java | 14 ++++ .../SericolaLatentStateBranchRateModel.java | 14 ++++ .../StrictClockBranchRates.java | 14 ++++ .../branchratemodel/TipBranchRateModel.java | 14 ++++ .../clock/RateEvolutionLikelihood.java | 14 ++++ src/dr/oldevomodel/clock/UniversalClock.java | 14 ++++ 30 files changed, 472 insertions(+) diff --git a/src/dr/evomodel/arg/ARGRelaxedClock.java b/src/dr/evomodel/arg/ARGRelaxedClock.java index c964010925..ee2891c086 100644 --- a/src/dr/evomodel/arg/ARGRelaxedClock.java +++ b/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() { diff --git a/src/dr/evomodel/arg/branchratemodel/ARGDiscretizedBranchRates.java b/src/dr/evomodel/arg/branchratemodel/ARGDiscretizedBranchRates.java index 02907d6fb7..cab1acabf0 100644 --- a/src/dr/evomodel/arg/branchratemodel/ARGDiscretizedBranchRates.java +++ b/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. diff --git a/src/dr/evomodel/branchratemodel/AncestralTraitBranchRateModel.java b/src/dr/evomodel/branchratemodel/AncestralTraitBranchRateModel.java index 232ce50320..3e67bd5284 100644 --- a/src/dr/evomodel/branchratemodel/AncestralTraitBranchRateModel.java +++ b/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) { diff --git a/src/dr/evomodel/branchratemodel/ArbitraryBranchRates.java b/src/dr/evomodel/branchratemodel/ArbitraryBranchRates.java index 283b5e6857..6f118504d6 100644 --- a/src/dr/evomodel/branchratemodel/ArbitraryBranchRates.java +++ b/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); diff --git a/src/dr/evomodel/branchratemodel/AttributeBranchRateModel.java b/src/dr/evomodel/branchratemodel/AttributeBranchRateModel.java index fd8c5f2ff2..7cc1244b1d 100644 --- a/src/dr/evomodel/branchratemodel/AttributeBranchRateModel.java +++ b/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() { diff --git a/src/dr/evomodel/branchratemodel/BranchRateModel.java b/src/dr/evomodel/branchratemodel/BranchRateModel.java index 498335cf3b..1dcb00545d 100644 --- a/src/dr/evomodel/branchratemodel/BranchRateModel.java +++ b/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 }; + } + }; } diff --git a/src/dr/evomodel/branchratemodel/CompoundBranchRateModel.java b/src/dr/evomodel/branchratemodel/CompoundBranchRateModel.java index af1c0fe7ee..7b112eb72c 100644 --- a/src/dr/evomodel/branchratemodel/CompoundBranchRateModel.java +++ b/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 }; + } + }; + } } \ No newline at end of file diff --git a/src/dr/evomodel/branchratemodel/ContinuousBranchRates.java b/src/dr/evomodel/branchratemodel/ContinuousBranchRates.java index 7265a39b91..08d5d5411b 100644 --- a/src/dr/evomodel/branchratemodel/ContinuousBranchRates.java +++ b/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. diff --git a/src/dr/evomodel/branchratemodel/ContinuousTraitBranchRateModel.java b/src/dr/evomodel/branchratemodel/ContinuousTraitBranchRateModel.java index 88a995743e..aa1587f8a5 100644 --- a/src/dr/evomodel/branchratemodel/ContinuousTraitBranchRateModel.java +++ b/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 }; + } + }; + } } \ No newline at end of file diff --git a/src/dr/evomodel/branchratemodel/CountableMixtureBranchRates.java b/src/dr/evomodel/branchratemodel/CountableMixtureBranchRates.java index 3663b6a4b9..4adc4a2a4e 100644 --- a/src/dr/evomodel/branchratemodel/CountableMixtureBranchRates.java +++ b/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(); diff --git a/src/dr/evomodel/branchratemodel/CountableModelMixtureBranchRates.java b/src/dr/evomodel/branchratemodel/CountableModelMixtureBranchRates.java index b2f6c3b1d2..eb4a1c9f02 100644 --- a/src/dr/evomodel/branchratemodel/CountableModelMixtureBranchRates.java +++ b/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; diff --git a/src/dr/evomodel/branchratemodel/DecayingRateModel.java b/src/dr/evomodel/branchratemodel/DecayingRateModel.java index 0af5f9678e..0aa2072c1f 100644 --- a/src/dr/evomodel/branchratemodel/DecayingRateModel.java +++ b/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. diff --git a/src/dr/evomodel/branchratemodel/DefaultBranchRateModel.java b/src/dr/evomodel/branchratemodel/DefaultBranchRateModel.java index 0a2ee31531..37dee99e16 100644 --- a/src/dr/evomodel/branchratemodel/DefaultBranchRateModel.java +++ b/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 diff --git a/src/dr/evomodel/branchratemodel/DiscreteTraitBranchRateModel.java b/src/dr/evomodel/branchratemodel/DiscreteTraitBranchRateModel.java index 419f1e2c95..98a0b1d42f 100644 --- a/src/dr/evomodel/branchratemodel/DiscreteTraitBranchRateModel.java +++ b/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] diff --git a/src/dr/evomodel/branchratemodel/DiscretizedBranchRates.java b/src/dr/evomodel/branchratemodel/DiscretizedBranchRates.java index f47cc7dcd5..81605112dd 100644 --- a/src/dr/evomodel/branchratemodel/DiscretizedBranchRates.java +++ b/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) { diff --git a/src/dr/evomodel/branchratemodel/FixedDriftModel.java b/src/dr/evomodel/branchratemodel/FixedDriftModel.java index e49f9aed8a..79ba847c64 100644 --- a/src/dr/evomodel/branchratemodel/FixedDriftModel.java +++ b/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) { diff --git a/src/dr/evomodel/branchratemodel/LatentStateBranchRateModel.java b/src/dr/evomodel/branchratemodel/LatentStateBranchRateModel.java index 72c3ba2bb0..254afe0156 100644 --- a/src/dr/evomodel/branchratemodel/LatentStateBranchRateModel.java +++ b/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) { diff --git a/src/dr/evomodel/branchratemodel/LocalBranchRates.java b/src/dr/evomodel/branchratemodel/LocalBranchRates.java index 59bc0a8379..160a672996 100644 --- a/src/dr/evomodel/branchratemodel/LocalBranchRates.java +++ b/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) { diff --git a/src/dr/evomodel/branchratemodel/LocalClockModel.java b/src/dr/evomodel/branchratemodel/LocalClockModel.java index ecd2dddf97..f986a4e373 100644 --- a/src/dr/evomodel/branchratemodel/LocalClockModel.java +++ b/src/dr/evomodel/branchratemodel/LocalClockModel.java @@ -273,6 +273,20 @@ 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 }; + } + }; + } /** * Set up the map from node to clock. diff --git a/src/dr/evomodel/branchratemodel/MixtureModelBranchRates.java b/src/dr/evomodel/branchratemodel/MixtureModelBranchRates.java index 23ddff0425..6550855511 100644 --- a/src/dr/evomodel/branchratemodel/MixtureModelBranchRates.java +++ b/src/dr/evomodel/branchratemodel/MixtureModelBranchRates.java @@ -233,6 +233,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) { //int rateCategory = (int) Math.round(rateCategories.getNodeValue(tree, 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 category indices. diff --git a/src/dr/evomodel/branchratemodel/PassageBranchRateModel.java b/src/dr/evomodel/branchratemodel/PassageBranchRateModel.java index 77a126aeb2..f04293c1c9 100644 --- a/src/dr/evomodel/branchratemodel/PassageBranchRateModel.java +++ b/src/dr/evomodel/branchratemodel/PassageBranchRateModel.java @@ -126,5 +126,19 @@ public double getBranchRate(final Tree tree, final NodeRef 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 }; + } + }; + } } \ No newline at end of file diff --git a/src/dr/evomodel/branchratemodel/RandomLocalClockModel.java b/src/dr/evomodel/branchratemodel/RandomLocalClockModel.java index 18e8be15b9..8adfc1a608 100644 --- a/src/dr/evomodel/branchratemodel/RandomLocalClockModel.java +++ b/src/dr/evomodel/branchratemodel/RandomLocalClockModel.java @@ -171,6 +171,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) { } return unscaledBranchRates[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 }; + } + }; + } private void calculateUnscaledBranchRates(TreeModel tree) { recursivelyCompute(tree, tree.getRoot(), 1.0); diff --git a/src/dr/evomodel/branchratemodel/RateEpochBranchRateModel.java b/src/dr/evomodel/branchratemodel/RateEpochBranchRateModel.java index 493eae401a..25412f8110 100644 --- a/src/dr/evomodel/branchratemodel/RateEpochBranchRateModel.java +++ b/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 weightList = new ArrayList(); + List rateList = new ArrayList(); + + 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; diff --git a/src/dr/evomodel/branchratemodel/RelaxedDriftModel.java b/src/dr/evomodel/branchratemodel/RelaxedDriftModel.java index 505fee7a82..4efcedac9b 100644 --- a/src/dr/evomodel/branchratemodel/RelaxedDriftModel.java +++ b/src/dr/evomodel/branchratemodel/RelaxedDriftModel.java @@ -137,6 +137,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) { } return branchRates[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 }; + } + }; + } private void calculateBranchRates(TreeModel tree) { branchRates[tree.getRoot().getNumber()] = getVariable(tree, tree.getRoot()); diff --git a/src/dr/evomodel/branchratemodel/ScaledTreeLengthRateModel.java b/src/dr/evomodel/branchratemodel/ScaledTreeLengthRateModel.java index 54903f1031..860f67d983 100644 --- a/src/dr/evomodel/branchratemodel/ScaledTreeLengthRateModel.java +++ b/src/dr/evomodel/branchratemodel/ScaledTreeLengthRateModel.java @@ -68,6 +68,20 @@ public double getBranchRate(final Tree tree, final NodeRef node) { } return rateFactor; } + + @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 getTotalLength() { return totalLength.getParameterValue(0); diff --git a/src/dr/evomodel/branchratemodel/SericolaLatentStateBranchRateModel.java b/src/dr/evomodel/branchratemodel/SericolaLatentStateBranchRateModel.java index b8c42ef19a..4852b9acc0 100644 --- a/src/dr/evomodel/branchratemodel/SericolaLatentStateBranchRateModel.java +++ b/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) { diff --git a/src/dr/evomodel/branchratemodel/StrictClockBranchRates.java b/src/dr/evomodel/branchratemodel/StrictClockBranchRates.java index 38e15572ef..1acb91aba7 100644 --- a/src/dr/evomodel/branchratemodel/StrictClockBranchRates.java +++ b/src/dr/evomodel/branchratemodel/StrictClockBranchRates.java @@ -73,5 +73,19 @@ protected void acceptState() { public double getBranchRate(final Tree tree, final NodeRef node) { return rateParameter.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 }; + } + }; + } } \ No newline at end of file diff --git a/src/dr/evomodel/branchratemodel/TipBranchRateModel.java b/src/dr/evomodel/branchratemodel/TipBranchRateModel.java index ec8328c9ec..facb4d55d5 100644 --- a/src/dr/evomodel/branchratemodel/TipBranchRateModel.java +++ b/src/dr/evomodel/branchratemodel/TipBranchRateModel.java @@ -122,5 +122,19 @@ public double getBranchRate(final Tree tree, final NodeRef 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 }; + } + }; + } } \ No newline at end of file diff --git a/src/dr/oldevomodel/clock/RateEvolutionLikelihood.java b/src/dr/oldevomodel/clock/RateEvolutionLikelihood.java index 9d5886ea8b..435c22a619 100644 --- a/src/dr/oldevomodel/clock/RateEvolutionLikelihood.java +++ b/src/dr/oldevomodel/clock/RateEvolutionLikelihood.java @@ -194,6 +194,20 @@ public double getBranchRate(Tree tree, NodeRef node) { if (tree.isRoot(node)) return rootRateParameter.getParameterValue(0); return ratesParameter.getNodeValue(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 }; + } + }; + } public boolean isEpisodic() { return isEpisodic; diff --git a/src/dr/oldevomodel/clock/UniversalClock.java b/src/dr/oldevomodel/clock/UniversalClock.java index 34c3af2013..143702e843 100644 --- a/src/dr/oldevomodel/clock/UniversalClock.java +++ b/src/dr/oldevomodel/clock/UniversalClock.java @@ -130,6 +130,20 @@ private void calculateRate(int index) { public double getBranchRate(Tree tree, NodeRef node) { throw new RuntimeException("Look at code before running this class!"); } + + @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 }; + } + }; + } Parameter rateParameter = null; Parameter massParameter = null; From e86792cb1b443abb33f28be4caede6c05b33eb79 Mon Sep 17 00:00:00 2001 From: Jiansi Gao Date: Fri, 4 Mar 2022 22:09:13 -0800 Subject: [PATCH 2/3] allowing the transition-probability matrix to be computed correctly along a branch when that branch overlaps with both multiple epochal matrices and multiple epochal rates --- src/dr/app/beagle/tools/Partition.java | 2 +- .../treelikelihood/BeagleTreeLikelihood.java | 2 +- .../MultiPartitionTreeLikelihood.java | 2 +- .../SubstitutionModelDelegate.java | 78 +++++++++++++++++-- 4 files changed, 74 insertions(+), 10 deletions(-) diff --git a/src/dr/app/beagle/tools/Partition.java b/src/dr/app/beagle/tools/Partition.java index 51ebc5e439..7ad22d92c7 100644 --- a/src/dr/app/beagle/tools/Partition.java +++ b/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() { diff --git a/src/dr/evomodel/treelikelihood/BeagleTreeLikelihood.java b/src/dr/evomodel/treelikelihood/BeagleTreeLikelihood.java index 2d4ebcfea9..4dec77022f 100644 --- a/src/dr/evomodel/treelikelihood/BeagleTreeLikelihood.java +++ b/src/dr/evomodel/treelikelihood/BeagleTreeLikelihood.java @@ -206,7 +206,7 @@ public BeagleTreeLikelihood(PatternList patternList, if (extraBufferOrder.size() > 0) { extraBufferCount = extraBufferOrder.get(instanceCount % extraBufferOrder.size()); } - substitutionModelDelegate = new SubstitutionModelDelegate(treeModel, branchModel, extraBufferCount); + substitutionModelDelegate = new SubstitutionModelDelegate(treeModel, branchModel, branchRateModel, extraBufferCount); // first set the rescaling scheme to use from the parser this.rescalingScheme = rescalingScheme; diff --git a/src/dr/evomodel/treelikelihood/MultiPartitionTreeLikelihood.java b/src/dr/evomodel/treelikelihood/MultiPartitionTreeLikelihood.java index 4a7f98d499..58cddea43c 100644 --- a/src/dr/evomodel/treelikelihood/MultiPartitionTreeLikelihood.java +++ b/src/dr/evomodel/treelikelihood/MultiPartitionTreeLikelihood.java @@ -197,7 +197,7 @@ public MultiPartitionTreeLikelihood(List patternLists, substitutionModelDelegates = new SubstitutionModelDelegate[partitionCount]; for (int i = 0; i < partitionCount; i++) { - substitutionModelDelegates[i] = new SubstitutionModelDelegate(treeModel, branchModel, extraBufferCount); + substitutionModelDelegates[i] = new SubstitutionModelDelegate(treeModel, branchModel, branchRateModel, extraBufferCount); } // first set the rescaling scheme to use from the parser diff --git a/src/dr/evomodel/treelikelihood/SubstitutionModelDelegate.java b/src/dr/evomodel/treelikelihood/SubstitutionModelDelegate.java index eb37ac1035..79fc8df230 100644 --- a/src/dr/evomodel/treelikelihood/SubstitutionModelDelegate.java +++ b/src/dr/evomodel/treelikelihood/SubstitutionModelDelegate.java @@ -27,6 +27,7 @@ import beagle.Beagle; import dr.evomodel.branchmodel.BranchModel; +import dr.evomodel.branchratemodel.BranchRateModel; import dr.evomodel.substmodel.EigenDecomposition; import dr.evomodel.substmodel.SubstitutionModel; import dr.evolution.tree.Tree; @@ -58,6 +59,7 @@ public final class SubstitutionModelDelegate implements EvolutionaryProcessDeleg private final Tree tree; private final List substitutionModelList; private final BranchModel branchModel; + private final BranchRateModel branchRateModel; private final int eigenCount; private final int nodeCount; @@ -71,10 +73,18 @@ public final class SubstitutionModelDelegate implements EvolutionaryProcessDeleg private Deque availableBuffers = new ArrayDeque(); public SubstitutionModelDelegate(Tree tree, BranchModel branchModel) { - this(tree, branchModel, BUFFER_POOL_SIZE_DEFAULT); + this(tree, branchModel, null, BUFFER_POOL_SIZE_DEFAULT); } - + + public SubstitutionModelDelegate(Tree tree, BranchModel branchModel, BranchRateModel branchRateModel) { + this(tree, branchModel, branchRateModel, BUFFER_POOL_SIZE_DEFAULT); + } + public SubstitutionModelDelegate(Tree tree, BranchModel branchModel, int bufferPoolSize) { + this(tree, branchModel, null, bufferPoolSize); + } + + public SubstitutionModelDelegate(Tree tree, BranchModel branchModel, BranchRateModel branchRateModel, int bufferPoolSize) { if (MEASURE_RUN_TIME) { updateTime = 0; @@ -86,6 +96,7 @@ public SubstitutionModelDelegate(Tree tree, BranchModel branchModel, int bufferP this.substitutionModelList = branchModel.getSubstitutionModels(); this.branchModel = branchModel; + this.branchRateModel = branchRateModel; eigenCount = substitutionModelList.size(); nodeCount = tree.getNodeCount(); @@ -179,10 +190,6 @@ public void updateTransitionMatrices(Beagle beagle, int[] branchIndices, double[ edgeLengths[k][counts[k]] = edgeLength[i]; counts[k]++; } else { - double sum = 0.0; - for (double w : weights) { - sum += w; - } if (getAvailableBufferCount() < order.length) { // too few buffers available, process what we have and continue... @@ -194,6 +201,63 @@ public void updateTransitionMatrices(Beagle beagle, int[] branchIndices, double[ counts[k] = 0; } } + + double sum = 0.0; + for (double w : weights) { + sum += w; + } + + double[] matrixWeights = new double[weights.length]; + double[] matrixWeightsCumsum = new double[weights.length + 1]; + double[] matrixEdgeLengths = new double[weights.length]; + for (int j = 0; j < weights.length; j++) { + matrixWeights[j] = weights[j] / sum; + matrixWeightsCumsum[j + 1] = matrixWeightsCumsum[j] + matrixWeights[j]; + matrixEdgeLengths[j] = matrixWeights[j] * edgeLength[i]; + } + + if (branchRateModel != null) { + + BranchRateModel.Mapping rateMapping = branchRateModel.getBranchRateModelMapping(tree, tree.getNode(branchIndices[i])); + double[] rates = rateMapping.getRates(); + double[] rateWeights = rateMapping.getWeights(); + + if (rates.length > 1) { + + sum = 0.0; + for (double w : rateWeights) { + sum += w; + } + double[] rateWeightsCumsum = new double[rateWeights.length + 1]; + for (int j = 0; j < rates.length; j++) { + rateWeights[j] = rateWeights[j] / sum; + rateWeightsCumsum[j + 1] = rateWeightsCumsum[j] + rateWeights[j]; + } + + double rate = branchRateModel.getBranchRate(tree, tree.getNode(branchIndices[i])); + double edgeTime = edgeLength[i] / rate; + + int k = 0; + double lastCumsum; + for (int j = 0; j < weights.length; j++) { + + matrixEdgeLengths[j] = 0.0; + lastCumsum = matrixWeightsCumsum[j]; + + while (k < rates.length && rateWeightsCumsum[k + 1] <= matrixWeightsCumsum[j + 1]) { + matrixEdgeLengths[j] += edgeTime * (rateWeightsCumsum[k + 1] - lastCumsum) * rates[k]; + lastCumsum = rateWeightsCumsum[k + 1]; + k++; + } + + if (matrixWeightsCumsum[j + 1] > lastCumsum && k < rates.length) { + matrixEdgeLengths[j] += edgeTime * (matrixWeightsCumsum[j + 1] - lastCumsum) * rates[k]; + } + + } + } + } + Deque bufferIndices = new ArrayDeque(); for (int j = 0; j < order.length; j++) { @@ -208,7 +272,7 @@ public void updateTransitionMatrices(Beagle beagle, int[] branchIndices, double[ int k = order[j]; probabilityIndices[k][counts[k]] = buffer; - edgeLengths[k][counts[k]] = weights[j] * edgeLength[i] / sum; + edgeLengths[k][counts[k]] = matrixEdgeLengths[j]; // edgeLengths[k][counts[k]] = weights[j] ; counts[k]++; From 9e81125b0af9c9dbe151ee5d4a49eca8941375d8 Mon Sep 17 00:00:00 2001 From: Jiansi Gao Date: Fri, 4 Mar 2022 22:13:09 -0800 Subject: [PATCH 3/3] implementing a stochastic mapping function that correctly simulates histories under either an epochal matrix model or an epochal rate model (or both) --- .../branchmodel/EpochBranchModel.java | 2 +- .../UniformizedSubstitutionModel.java | 12 +- .../AncestralStateBeagleTreeLikelihood.java | 190 +++++++++++++++++- .../MarkovJumpsBeagleTreeLikelihood.java | 157 +++++++++++---- 4 files changed, 314 insertions(+), 47 deletions(-) diff --git a/src/dr/evomodel/branchmodel/EpochBranchModel.java b/src/dr/evomodel/branchmodel/EpochBranchModel.java index 38ce8b2380..f449a5921a 100644 --- a/src/dr/evomodel/branchmodel/EpochBranchModel.java +++ b/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); diff --git a/src/dr/evomodel/substmodel/UniformizedSubstitutionModel.java b/src/dr/evomodel/substmodel/UniformizedSubstitutionModel.java index 773e3a7168..498f5a0fa3 100644 --- a/src/dr/evomodel/substmodel/UniformizedSubstitutionModel.java +++ b/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() { diff --git a/src/dr/evomodel/treelikelihood/AncestralStateBeagleTreeLikelihood.java b/src/dr/evomodel/treelikelihood/AncestralStateBeagleTreeLikelihood.java index 636502af72..2f460e3571 100644 --- a/src/dr/evomodel/treelikelihood/AncestralStateBeagleTreeLikelihood.java +++ b/src/dr/evomodel/treelikelihood/AncestralStateBeagleTreeLikelihood.java @@ -28,19 +28,24 @@ import dr.evolution.datatype.*; import dr.evolution.tree.*; import dr.evomodel.branchmodel.BranchModel; +import dr.evomodel.branchmodel.EpochBranchModel; import dr.evomodel.siteratemodel.SiteRateModel; import dr.evolution.alignment.PatternList; import dr.evolution.alignment.UncertainSiteList; import dr.evomodel.branchratemodel.BranchRateModel; +import dr.evomodel.branchratemodel.RateEpochBranchRateModel; import dr.evomodel.tipstatesmodel.TipStatesModel; import dr.inference.model.Model; import dr.inference.model.Parameter; +import dr.inference.markovjumps.MarkovJumpsCore; import dr.math.MathUtils; import java.util.Map; import java.util.Set; import java.util.function.Function; +import java.util.*; + /** * @author Marc Suchard * @author Andrew Rambaut @@ -82,6 +87,13 @@ public AncestralStateBeagleTreeLikelihood(PatternList patternList, MutableTreeMo probabilities = new double[stateCount * stateCount * categoryCount]; partials = new double[stateCount * patternCount * categoryCount]; + + probabilitiesAlongBranch = new ArrayList(); + probabilitiesConvolved = new ArrayList(); + combinedWeights = new ArrayList(); + combinedMatrixOrder = new ArrayList(); + combinedRates = new ArrayList(); + // rootPartials = new double[stateCount*patternCount]; // cumulativeScaleBuffers = new int[nodeCount][]; // scaleBufferIndex = getScaleBufferCount() - 1; @@ -238,7 +250,7 @@ protected int getScaleBufferCount() { return internalNodeCount + 2; } - private int drawChoice(double[] measure) { + protected int drawChoice(double[] measure) { if (useMAP) { double max = measure[0]; int choice = 0; @@ -326,6 +338,167 @@ protected void getMatrix(int branchIndex, double[] probabilities) { beagle.getTransitionMatrix(substitutionModelDelegate.getMatrixIndex(branchIndex), probabilities); // NB: It may be faster to compute matrices in BEAST via substitutionModel } + + protected void getTransitionProbabilityMatrices(Tree tree, NodeRef childNode) { + + NodeRef parentNode = tree.getParent(childNode); + + final double parentTime = tree.getNodeHeight(parentNode); + final double childTime = tree.getNodeHeight(childNode); + final double branchTime = parentTime - childTime; + final double branchRate = branchRateModel.getBranchRate(tree, childNode); + + // fetch the matrix and/or rate epochal layout along the branch + // combinedWeights contains the (absolute instead of relative) duration of each piece + combinedWeights.clear(); + combinedMatrixOrder.clear(); + combinedRates.clear(); + + BranchModel.Mapping matrixMapping = branchModel.getBranchModelMapping(childNode); + int[] matrixOrder = matrixMapping.getOrder(); + double[] matrixWeights = matrixMapping.getWeights(); + int nmatrices = matrixOrder.length; + + if (!(branchModel instanceof EpochBranchModel)) { + for (int j = 0; j < nmatrices; j++) { + matrixWeights[j] *= branchTime; + } + } + + BranchRateModel.Mapping rateMapping = branchRateModel.getBranchRateModelMapping(tree, childNode); + double[] rates = rateMapping.getRates(); + double[] rateWeights = rateMapping.getWeights(); + int nrates = rates.length; + + if (!(branchRateModel instanceof RateEpochBranchRateModel)) { + for (int j = 0; j < nrates; j++) { + rateWeights[j] *= branchTime; + } + } + + // generate cumulative sum vector for matrices + double[] matrixWeightsCumsum = new double[nmatrices + 1]; + for (int j = 0; j < nmatrices - 1; j++) { + matrixWeightsCumsum[j + 1] = matrixWeightsCumsum[j] + matrixWeights[j]; + } + + // generate cumulative sum vector for rates + double[] rateWeightsCumsum = new double[nrates + 1]; + for (int j = 0; j < nrates - 1; j++) { + rateWeightsCumsum[j + 1] = rateWeightsCumsum[j] + rateWeights[j]; + } + + // make sure that the last time in the matrix and rate weights cumsum vectors match (which should equal the branch duration) + matrixWeightsCumsum[nmatrices] = branchTime; + rateWeightsCumsum[nrates] = branchTime; + + if (nmatrices == 1 && nrates == 1) { + + combinedWeights.add(matrixWeights[0]); + combinedMatrixOrder.add(matrixOrder[0]); + combinedRates.add(branchRate); + + } else if (nmatrices > 1 && nrates == 1) { + + for (double w : matrixWeights) { + combinedWeights.add(w); + } + for (int i : matrixOrder) { + combinedMatrixOrder.add(i); + } + for (int j = 0; j < nmatrices; j++) { + combinedRates.add(branchRate); + } + + } else if (nmatrices == 1 && nrates > 1) { + + for (double w : rateWeights) { + combinedWeights.add(w); + } + for (int j = 0; j < nrates; j++) { + combinedMatrixOrder.add(matrixOrder[0]); + } + for (double r : rates) { + combinedRates.add(r); + } + + } else { + + int matrixId = 0; + int rateId = 0; + double lastCumsum = 0.0; + + while (matrixId < nmatrices || rateId < nrates) { + combinedMatrixOrder.add(matrixOrder[matrixId]); + combinedRates.add(rates[rateId]); + + if (matrixWeightsCumsum[matrixId + 1] < rateWeightsCumsum[rateId + 1]) { + combinedWeights.add(matrixWeightsCumsum[matrixId + 1] - lastCumsum); + lastCumsum = matrixWeightsCumsum[matrixId + 1]; + matrixId++; + } else if (matrixWeightsCumsum[matrixId + 1] > rateWeightsCumsum[rateId + 1]) { + combinedWeights.add(rateWeightsCumsum[rateId + 1] - lastCumsum); + lastCumsum = rateWeightsCumsum[rateId + 1]; + rateId++; + } else { + combinedWeights.add(matrixWeightsCumsum[matrixId + 1] - lastCumsum); + lastCumsum = matrixWeightsCumsum[matrixId + 1]; + matrixId++; + rateId++; + } + } + } + + int npieces = combinedWeights.size(); + + // compute transition probability matrix for each piece + probabilitiesAlongBranch.clear(); + for (int j = 0; j < npieces; j++) { + probabilitiesAlongBranch.add(new double[stateCount * stateCount * categoryCount]); + + for (int k = 0; k < categoryCount; k++) { + final double edgeLength = combinedWeights.get(j) * combinedRates.get(j) * siteRateModel.getRateForCategory(k); + double[] probabilitiesTmp = new double[stateCount * stateCount]; + substitutionModelDelegate.getSubstitutionModel(combinedMatrixOrder.get(j)).getTransitionProbabilities(edgeLength, probabilitiesTmp); + System.arraycopy(probabilitiesTmp, 0, probabilitiesAlongBranch.get(j), k * stateCount * stateCount, stateCount * stateCount); + } + } + + // convolve transition probability matrix + probabilitiesConvolved.clear(); + for (int j = npieces - 2; j >= 0; j--) { + probabilitiesConvolved.add(0, new double[stateCount * stateCount * categoryCount]); + + for (int k = 0; k < categoryCount; k++) { + double[] probabilitiesTmp0 = new double[stateCount * stateCount]; + double[] probabilitiesTmp1 = new double[stateCount * stateCount]; + double[] probabilitiesTmp2 = new double[stateCount * stateCount]; + + System.arraycopy(probabilitiesAlongBranch.get(j), k * stateCount * stateCount, probabilitiesTmp1, 0, stateCount * stateCount); + if (j == npieces - 2) { + System.arraycopy(probabilitiesAlongBranch.get(j + 1), k * stateCount * stateCount, probabilitiesTmp2, 0, stateCount * stateCount); + } else { + System.arraycopy(probabilitiesConvolved.get(1), k * stateCount * stateCount, probabilitiesTmp2, 0, stateCount * stateCount); + } + + MarkovJumpsCore.matrixMultiply(probabilitiesTmp1, probabilitiesTmp2, stateCount, probabilitiesTmp0); + System.arraycopy(probabilitiesTmp0, 0, probabilitiesConvolved.get(0), k * stateCount * stateCount, stateCount * stateCount); + } + } + } + + protected void getTransitionProbabilityMatrix(Tree tree, NodeRef childNode, double[] probabilities, boolean update) { + + if (update) { + getTransitionProbabilityMatrices(tree, childNode); + } + + if (probabilitiesConvolved.size() > 0) { + System.arraycopy(probabilitiesConvolved.get(0), 0, probabilities, 0, categoryCount * stateCount * stateCount); + } else { + System.arraycopy(probabilitiesAlongBranch.get(0), 0, probabilities, 0, categoryCount * stateCount * stateCount); + } + } public void setTipStates(int tipNum, int[] states) { System.arraycopy(states, 0, tipStates[tipNum], 0, states.length); @@ -490,7 +663,8 @@ public void traverseSample(Tree tree, NodeRef node, int[] parentState, int[] rat // if (categoryCount > 1) // throw new RuntimeException("Reconstruction not implemented for multiple categories yet."); - getMatrix(nodeNum, probabilities); +// getMatrix(nodeNum, probabilities); + getTransitionProbabilityMatrix(tree, node, probabilities, true); for (int j = 0; j < patternCount; j++) { @@ -533,10 +707,11 @@ public void traverseSample(Tree tree, NodeRef node, int[] parentState, int[] rat } else { // This is an external leaf + getTransitionProbabilityMatrix(tree, node, probabilities, true); if (useAmbiguities()) { - getMatrix(nodeNum, probabilities); +// getMatrix(nodeNum, probabilities); double[] partials = tipPartials[nodeNum]; for (int j = 0; j < patternCount; j++) { @@ -576,7 +751,7 @@ public void traverseSample(Tree tree, NodeRef node, int[] parentState, int[] rat int category = rateCategory == null ? 0 : rateCategory[j]; int matrixIndex = category * stateCount * stateCount; - getMatrix(nodeNum, probabilities); +// getMatrix(nodeNum, probabilities); System.arraycopy(probabilities, parentIndex + matrixIndex, conditionalProbabilities, 0, stateCount); if (useAmbiguities && !dataType.isUnknownState(thisState)) { // Not completely unknown @@ -599,7 +774,7 @@ public void traverseSample(Tree tree, NodeRef node, int[] parentState, int[] rat if (!returnMarginalLogLikelihood) { final int parentIndex = parentState[j] * stateCount; - getMatrix(nodeNum, probabilities); +// getMatrix(nodeNum, probabilities); if (!returnMarginalLogLikelihood) { double contrib = probabilities[parentIndex + reconstructedStates[nodeNum][j]]; jointLogLikelihood += Math.log(contrib); @@ -635,6 +810,11 @@ protected void hookCalculation(Tree tree, NodeRef parentNode, NodeRef childNode, private double[][] tipPartials; private double[] probabilities; + protected List probabilitiesAlongBranch; + protected List probabilitiesConvolved; + protected List combinedWeights = new ArrayList(); + protected List combinedMatrixOrder = new ArrayList(); + protected List combinedRates = new ArrayList(); private double[] partials; protected int[] rateCategory = null; diff --git a/src/dr/evomodel/treelikelihood/MarkovJumpsBeagleTreeLikelihood.java b/src/dr/evomodel/treelikelihood/MarkovJumpsBeagleTreeLikelihood.java index 80fbd65fde..2a15aa5669 100644 --- a/src/dr/evomodel/treelikelihood/MarkovJumpsBeagleTreeLikelihood.java +++ b/src/dr/evomodel/treelikelihood/MarkovJumpsBeagleTreeLikelihood.java @@ -42,6 +42,7 @@ import dr.evomodel.tipstatesmodel.TipStatesModel; import dr.inference.loggers.LogColumn; import dr.inference.loggers.NumberColumn; +import dr.inference.markovjumps.MarkovJumpsCore; import dr.inference.markovjumps.MarkovJumpsRegisterAcceptor; import dr.inference.markovjumps.MarkovJumpsType; import dr.inference.model.Parameter; @@ -133,11 +134,11 @@ public void addRegister(Parameter addRegisterParameter, } addVariable(addRegisterParameter); final String tag = addRegisterParameter.getId(); + + boolean isEpochModel = branchModel instanceof EpochBranchModel; for (int i = 0; i < substitutionModelDelegate.getSubstitutionModelCount(); ++i) { - boolean isEpochModel = branchModel instanceof EpochBranchModel; - registerParameter.add(addRegisterParameter); MarkovJumpsSubstitutionModel mjModel; SubstitutionModel substitutionModel = substitutionModelDelegate.getSubstitutionModel(i); @@ -206,7 +207,6 @@ public double[] getTrait(Tree tree, NodeRef node) { if (i == 0 || !isEpochModel) { - if (histories == null) { histories = new String[treeModel.getNodeCount()][patternCount]; } else { @@ -307,12 +307,8 @@ public boolean getLoggable() { } }); } - } - - if (isEpochModel) { - for (int j = 0; j < markovjumps.size(); ++j) { - ((UniformizedSubstitutionModel)markovjumps.get(j)).setSaveCompleteHistory(true); - } + } else if (isEpochModel) { + ((UniformizedSubstitutionModel) mjModel).setSaveCompleteHistory(true); } } @@ -413,7 +409,8 @@ protected void hookCalculation(Tree tree, NodeRef parentNode, NodeRef childNode, double[] probabilities = inProbabilities; if (probabilities == null) { // Leaf will call this hook with a null - getMatrix(childNum, tmpProbabilities); +// getMatrix(childNum, tmpProbabilities); + getTransitionProbabilityMatrix(tree, childNode, tmpProbabilities, false); probabilities = tmpProbabilities; } @@ -421,33 +418,116 @@ protected void hookCalculation(Tree tree, NodeRef parentNode, NodeRef childNode, final double parentTime = tree.getNodeHeight(parentNode); final double childTime = tree.getNodeHeight(childNode); final double substTime = parentTime - childTime; - + + if (histories != null) { + Arrays.fill(histories[childNum], "{"); + } + for (int r = 0; r < markovjumps.size(); r++) { - MarkovJumpsSubstitutionModel thisMarkovJumps = markovjumps.get(r); - - final int modelNumberFromrRegistry = branchModelNumber.get(r); -// int dummy = 0; -// final int modelNumberFromTree = branchSubstitutionModel.getBranchIndex(tree, childNode, dummy); - // @todo AR - not sure about this - if this is an epoch this is just going to get the most - // @todo tipward model for the branch. I think this was what was happening before (in comment, - // @todo above). - BranchModel.Mapping mapping = branchModel.getBranchModelMapping(childNode); - - if (modelNumberFromrRegistry == mapping.getOrder()[0]) { - if (useUniformization) { - computeSampledMarkovJumpsForBranch(((UniformizedSubstitutionModel) thisMarkovJumps), substTime, - branchRate, childNum, parentStates, childStates, parentTime, childTime, probabilities, scaleByTime[r], - expectedJumps.get(r), rateCategory, - (branchModel instanceof EpochBranchModel) || r == historyRegisterNumber - ); - } else { - computeIntegratedMarkovJumpsForBranch(thisMarkovJumps, substTime, branchRate, childNum, parentStates, - childStates, probabilities, condJumps, scaleByTime[r], expectedJumps.get(r), rateCategory); + Arrays.fill(expectedJumps.get(r)[childNum], 0.0); + } + + BranchModel.Mapping mapping = branchModel.getBranchModelMapping(childNode); + int[] order = mapping.getOrder(); + double[] weights = mapping.getWeights(); + + if (combinedWeights.size() == 1 || (!useUniformization)) { + for (int r = 0; r < markovjumps.size(); r++) { + MarkovJumpsSubstitutionModel thisMarkovJumps = markovjumps.get(r); + + final int modelNumberFromrRegistry = branchModelNumber.get(r); + // int dummy = 0; + // final int modelNumberFromTree = branchSubstitutionModel.getBranchIndex(tree, childNode, dummy); + // @todo AR - not sure about this - if this is an epoch this is just going to get the most + // @todo tipward model for the branch. I think this was what was happening before (in comment, + // @todo above). + + if (modelNumberFromrRegistry == mapping.getOrder()[0]) { + if (useUniformization) { + computeSampledMarkovJumpsForBranch(((UniformizedSubstitutionModel) thisMarkovJumps), substTime, + branchRate, childNum, parentStates, childStates, parentTime, childTime, probabilities, scaleByTime[r], + expectedJumps.get(r), rateCategory, + (branchModel instanceof EpochBranchModel) || r == historyRegisterNumber); + } else { + computeIntegratedMarkovJumpsForBranch(thisMarkovJumps, substTime, branchRate, childNum, parentStates, + childStates, probabilities, condJumps, scaleByTime[r], expectedJumps.get(r), rateCategory); + } } - } else { - // Fill with zeros - double[] result = expectedJumps.get(r)[childNum]; - Arrays.fill(result, 0.0); + } + } else if (markovjumps.size() > 0) { + + // compute time for each piece + int npieces = combinedWeights.size(); + double[] parentTimes = new double[npieces]; + double[] childTimes = new double[npieces]; + double[] substTimes = new double[npieces]; + parentTimes[0] = parentTime; + + for (int j = 0; j < npieces; j++) { + substTimes[j] = combinedWeights.get(j); + childTimes[j] = parentTimes[j] - substTimes[j]; + + if (j < npieces - 1) { + parentTimes[j + 1] = childTimes[j]; + } + } + + // sample intermediate states + int[][] parentStatesAll = new int[npieces][patternCount]; + int[][] childStatesAll = new int[npieces][patternCount]; + parentStatesAll[0] = parentStates; + childStatesAll[npieces - 1] = childStates; + + for (int j = 0; j < npieces - 1; j++) { + for (int k = 0; k < patternCount; k++) { + final int category = rateCategory == null ? 0 : rateCategory[k]; + final int matrixIndex = category * stateCount * stateCount; + final int startState = parentStatesAll[j][k]; + final int endState = childStates[k]; + double probSE = probabilitiesConvolved.get(j)[matrixIndex + startState * stateCount + endState]; + + double[] interStateProb = new double[stateCount]; + double probSum = 0.0; + + for (int interState = 0; interState < stateCount; interState++) { + double probSI = probabilitiesAlongBranch.get(j)[matrixIndex + startState * stateCount + interState]; + double probIE; + if (j == npieces - 2) { + probIE = probabilitiesAlongBranch.get(j + 1)[matrixIndex + interState * stateCount + endState]; + } else { + probIE = probabilitiesConvolved.get(j + 1)[matrixIndex + interState * stateCount + endState]; + } + + interStateProb[interState] = probSI * probIE / probSE; + probSum += interStateProb[interState]; + } + + for (int interState = 0; interState < stateCount; interState++) { + interStateProb[interState] /= probSum; + } + + int interState = drawChoice(interStateProb); + parentStatesAll[j + 1][k] = interState; + childStatesAll[j][k] = interState; + } + } + + for (int j = 0; j < npieces; j++) { + int r = branchModelNumber.indexOf(combinedMatrixOrder.get(j)); + MarkovJumpsSubstitutionModel thisMarkovJumps = markovjumps.get(r); + + computeSampledMarkovJumpsForBranch(((UniformizedSubstitutionModel) thisMarkovJumps), substTimes[j], + combinedRates.get(j), childNum, parentStatesAll[j], childStatesAll[j], parentTimes[j], childTimes[j], probabilitiesAlongBranch.get(j), scaleByTime[r], + expectedJumps.get(r), rateCategory, true); + } + } + + if (histories != null) { + for (int j = 0; j < patternCount; j++) { + if (histories[childNum][j].charAt(histories[childNum][j].length() - 1) == ',') { + histories[childNum][j] = histories[childNum][j].substring(0, histories[childNum][j].length() - 1); + } + histories[childNum][j] += '}'; } } } @@ -481,10 +561,13 @@ private void computeSampledMarkovJumpsForBranch(UniformizedSubstitutionModel thi if (scaleByTime) { value /= branchRate * categoryRate; } - thisExpectedJumps[childNum][j] = value; + thisExpectedJumps[childNum][j] += value; if (saveHistory) { int site = (useCompactHistory) ? j + 1 : -1; - histories[childNum][j] = thisMarkovJumps.getCompleteHistory(site, parentTime, childTime); + String historyTmp = thisMarkovJumps.getCompleteHistory(site, parentTime, childTime, false); + if (historyTmp != null && historyTmp.length() > 0) { + histories[childNum][j] += historyTmp + ','; + } } } }