Fix for constant drift model following issue #1083

beast-dev · May 7, 2020 · a1510cc · a1510cc
1 parent f42f4bb
commit a1510cc
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/src/dr/app/beauti/components/continuous/ContinuousComponentGenerator.java b/src/dr/app/beauti/components/continuous/ContinuousComponentGenerator.java
@@ -548,7 +548,7 @@ private void writeMultivariateTreeLikelihood(XMLWriter writer,
                                                  String diffusionModelId,
                                                  String treeModelId) {
 
-        int traitDimension = 1; // todo - set this to trait dimension
+        int traitDimension = partitionData.getTraits().size();
         writer.writeOpenTag(continuousDataLikelihoodParser.getParserTag(),
                 new Attribute[]{
                         new Attribute.Default<String>("id", continuousDataLikelihoodParser.getId(partitionData.getName())),
@@ -602,11 +602,13 @@ private void writeMultivariateTreeLikelihood(XMLWriter writer,
             writer.writeOpenTag("driftModels");
             for (int i = 0; i < traitDimension; i++) {
                 writer.writeOpenTag("strictClockBranchRates");
+                writer.writeOpenTag("rate");
                 writer.writeTag("parameter", new Attribute[]{
                         new Attribute.Default<String>("id", partitionData.getName() + "." + ContinuousComponentOptions.DRIFT_RATE +
                                 (traitDimension > 1 ? "." + i : "")),
                         new Attribute.Default<String>("value", "0.0"),
                 }, true);
+                writer.writeCloseTag("rate");
                 writer.writeCloseTag("strictClockBranchRates");
             }
             writer.writeCloseTag("driftModels");
@@ -645,7 +647,8 @@ private void writeMultivariateTreeLikelihood(XMLWriter writer,
 
         writer.writeCloseTag("conjugateRootPrior");
 
-        if (partitionData.getPartitionSubstitutionModel().getContinuousSubstModelType() != ContinuousSubstModelType.HOMOGENOUS) {
+        if (partitionData.getPartitionSubstitutionModel().getContinuousSubstModelType() != ContinuousSubstModelType.HOMOGENOUS &&
+                partitionData.getPartitionSubstitutionModel().getContinuousSubstModelType() != ContinuousSubstModelType.DRIFT) {
             if (ContinuousComponentOptions.USE_ARBITRARY_BRANCH_RATE_MODEL) {
                 writer.writeIDref("arbitraryBranchRates", partitionData.getName() + "." + "diffusion.branchRates");
             } else {
@@ -661,7 +664,7 @@ private void writeMultivariateTreeLikelihood(XMLWriter writer,
             for (int i = 0; i < traitDimension; i++) { // todo iterate over dimension of trait
                 writer.writeTag("parameter", new Attribute.Default<String>("idref", partitionData.getName() + "." + ContinuousComponentOptions.DRIFT_RATE + "." + i), true);
             }
-            writer.writeCloseTag("priorSampleSize");
+            writer.writeCloseTag("compoundParameter");
         }
 
     }
@@ -885,7 +888,8 @@ private void writeTreeLogEntries(PartitionTreeModel treeModel, XMLWriter writer)
                 PartitionSubstitutionModel model = partitionData.getPartitionSubstitutionModel();
                 writer.writeIDref("multivariateDiffusionModel", model.getName() + ".diffusionModel");
                 continuousDataLikelihoodParser.writeIDrefFromName(writer, partitionData.getName());
-                if (model.getContinuousSubstModelType() != ContinuousSubstModelType.HOMOGENOUS) {
+                if (model.getContinuousSubstModelType() != ContinuousSubstModelType.HOMOGENOUS &&
+                        model.getContinuousSubstModelType() != ContinuousSubstModelType.DRIFT) {
                     writer.writeOpenTag(TreeLoggerParser.TREE_TRAIT,
                             new Attribute[]{
                                     new Attribute.Default<String>(TreeLoggerParser.NAME, "rate"),