Add thorney beast support in Beauti

.gitignore

@@ -24,6 +24,8 @@ release/Mac
 release/Linux
 release/Windows
 src/revision.txt
+.vscode/launch.json
+.vscode/settings.json
 
 *.citations.txt
 *.ops

build.xml

@@ -238,6 +238,7 @@
                 <include name="dr/evomodel/speciation/**/*.class"/>
                 <include name="dr/evomodel/substmodel/**/*.class"/>
                 <include name="dr/evomodel/tree/**/*.class"/>
+                <include name="dr/evomodel/bigfasttree/**/*.class"/>
                 <include name="dr/evomodel/treelikelihood/**/*.class"/>
                 <include name="dr/evomodel/treedatalikelihood/**/*.class"/>
                 <include name="dr/evomodelxml/**/*.class"/>

src/dr/app/beast/development_parsers.properties

@@ -1,7 +1,7 @@
 #
 # development_parsers.properties
 #
-# Copyright © 2002-2024 the BEAST Development Team
+# Copyright © 2002-2024 the BEAST Development Team
 # http://beast.community/about
 #
 # This file is part of BEAST.
@@ -264,6 +264,8 @@ dr.evomodelxml.bigfasttree.thorney.ThorneyTreeGradientParser
 dr.evomodelxml.bigfasttree.thorney.UniformSubtreePruneRegraftParser
 dr.evomodel.bigfasttree.thorney.ZeroMutationMasker
 dr.evomodelxml.bigfasttree.thorney.SubtreeRootHeightStatisticParser
+dr.evomodel.bigfasttree.thorney.RootHeightProxyParameter
+
 # GLM
 dr.inferencexml.operators.MaskMoveOperatorParser
 dr.evomodelxml.continuous.hmc.GlmSubstitutionModelGradientParser

src/dr/app/beauti/generator/BeastGenerator.java

@@ -239,6 +239,23 @@ public void checkOptions() throws GeneratorException {
                 }
             }
 
+            //+++++++++++++++++ ThorneyBEAST / clock model Verification ++++++++++++++
+
+            for (PartitionClockModel model : options.getPartitionClockModels()) {
+                for (AbstractPartitionData pd : options.getDataPartitions(model)) {
+                    if(pd.getDataType().getType() ==DataType.TREE){
+                        PartitionTreeModel tm = pd.getPartitionTreeModel();
+                        if(tm.isUsingThorneyBEAST() && !PartitionTreeModel.thorneyCompatibleClocks.contains(model.getClockType())){
+                            throw new GeneratorException("Gradients for "+model.getClockType()+ 
+                            " are not implemented for Thorney BEAST." +
+                            "Please use a clock model that does not rely on gradients.", BeautiFrame.CLOCK_MODELS);
+
+                        }
+                    }
+                }
+            }
+
+
             //++++++++++++++++ Prior Bounds ++++++++++++++++++
             for (Parameter param : options.selectParameters()) {
                 if (!Double.isNaN(param.getInitial()) ) {

src/dr/app/beauti/generator/InitialTreeGenerator.java

@@ -107,8 +107,39 @@ public void writeStartingTree(PartitionTreeModel model, XMLWriter writer) {
                 break;
 
             case RANDOM:
-                // generate a coalescent tree
                 String simulatorId = prefix + STARTING_TREE;
+                String constraintID = prefix + "constraintsTree";
+                if(model.isUsingThorneyBEAST()){
+                    // write tree with dates 
+                    writer.writeComment("The user-specified tree with dates used to simualte a compatible bifurcating starting tree.");
+                    writer.writeOpenTag(
+                            NewickParser.NEWICK,
+                            new Attribute[]{
+                                new Attribute.Default<>(XMLParser.ID, constraintID),
+                                new Attribute.Default<Boolean>(SimpleTreeParser.USING_DATES, options.useTipDates)
+                            }
+                    );
+                    writer.writeText(TreeUtils.newick(model.getTreePartitionData().getTrees().getTrees().get(0)));
+                    writer.writeCloseTag(NewickParser.NEWICK);
+
+                    writer.writeComment("Generate a random starting tree under the coalescent process");
+                    writer.writeOpenTag(
+                        CoalescentSimulatorParser.COALESCENT_SIMULATOR,
+                        new Attribute[]{
+                                new Attribute.Default<String>(XMLParser.ID, simulatorId)
+                        });
+                        writeInitialDemoModelRef(model, writer);
+                        writer.writeTag(
+                            CoalescentSimulatorParser.CONSTRAINTS_TREE,
+                            false
+                        );
+                        writer.writeIDref(NewickParser.NEWICK, constraintID);
+                        writer.writeCloseTag(CoalescentSimulatorParser.CONSTRAINTS_TREE);
+                    writer.writeCloseTag(CoalescentSimulatorParser.COALESCENT_SIMULATOR);
+
+                }else{
+
+                // generate a coalescent tree
 
                 String taxaId = TaxaParser.TAXA;
                 AbstractPartitionData partition = options.getDataPartitions(model).get(0);
@@ -132,6 +163,7 @@ public void writeStartingTree(PartitionTreeModel model, XMLWriter writer) {
                     writeInitialDemoModelRef(model, writer);
                     writer.writeCloseTag(CoalescentSimulatorParser.COALESCENT_SIMULATOR);
                 }
+            }
                 break;
 
             default:

src/dr/app/beauti/generator/OperatorsGenerator.java

@@ -34,10 +34,12 @@
 import dr.app.beauti.util.XMLWriter;
 import dr.evolution.datatype.DataType;
 import dr.evomodel.branchratemodel.BranchSpecificFixedEffects;
+import dr.evomodel.bigfasttree.thorney.ConstrainedTreeModel;
 import dr.evomodel.operators.BitFlipInSubstitutionModelOperator;
 import dr.evomodel.operators.EmpiricalTreeDistributionOperator;
 import dr.evomodel.tree.DefaultTreeModel;
 import dr.evomodel.tree.EmpiricalTreeDistributionModel;
+import dr.evomodelxml.bigfasttree.thorney.UniformSubtreePruneRegraftParser;
 import dr.evomodelxml.branchratemodel.AutoCorrelatedBranchRatesDistributionParser;
 import dr.evomodelxml.branchratemodel.AutoCorrelatedGradientWrtIncrementsParser;
 import dr.evomodelxml.branchratemodel.BranchRateGradientWrtIncrementsParser;
@@ -255,6 +257,15 @@ private void writeOperator(Operator operator, XMLWriter writer) {
             case SHRINKAGE_CLOCK_GIBBS_OPERATOR:
                 writeShrinkageClockGibbsOperator(operator, prefix, writer);
                 break;
+            case NODE_HEIGHT_OPERATOR_UNIFORM :
+                writeBFTUniformNodeHeightOperator(operator, prefix, writer);
+                break;
+            case NODE_HEIGHT_OPERATOR_ROOT:
+                writeBFTRootScaleOperator(operator, prefix, writer);
+                break;
+            case UNIFORM_SUBTREE_PRUNE_REGRAFT:
+                writeUniformSPROperator(operator, prefix, writer);
+                break;
             default:
                 throw new IllegalArgumentException("Unknown operator type");
         }
@@ -747,6 +758,40 @@ private void writeEmpiricalTreeSwapOperator(Operator operator, String treeModelP
         writer.writeCloseTag(EmpiricalTreeDistributionOperator.EMPIRICAL_TREE_DISTRIBUTION_OPERATOR);
     }
 
+    private void writeBFTUniformNodeHeightOperator(Operator operator, String treeModelPrefix, XMLWriter writer) {
+        writer.writeOpenTag(NodeHeightOperatorParser.NODE_HEIGHT_OPERATOR, // should this be in the class and not the parser?
+                new Attribute[]{
+                    new Attribute.Default<String>("type",NodeHeightOperatorParser.OperatorType.UNIFORM.toString()),
+                        getWeightAttribute(operator.getWeight())
+                }
+        );
+        writer.writeIDref(ConstrainedTreeModel.CONSTRAINED_TREE_MODEL, treeModelPrefix + DefaultTreeModel.TREE_MODEL);
+        writer.writeCloseTag(NodeHeightOperatorParser.NODE_HEIGHT_OPERATOR);
+    }
+
+    private void writeBFTRootScaleOperator(Operator operator, String treeModelPrefix, XMLWriter writer) {
+        writer.writeOpenTag(NodeHeightOperatorParser.NODE_HEIGHT_OPERATOR, // should this be in the class and not the parser?
+                new Attribute[]{
+                    new Attribute.Default<String>("type",NodeHeightOperatorParser.OperatorType.SCALEROOT.toString()),
+                    // todo scale factor here
+                    new Attribute.Default<Double>(NodeHeightOperatorParser.SCALE_FACTOR, operator.getTuning()),
+                        getWeightAttribute(operator.getWeight())
+                }
+        );
+        writer.writeIDref(ConstrainedTreeModel.CONSTRAINED_TREE_MODEL, treeModelPrefix + DefaultTreeModel.TREE_MODEL);
+        writer.writeCloseTag(NodeHeightOperatorParser.NODE_HEIGHT_OPERATOR);
+    }
+
+    private void writeUniformSPROperator(Operator operator, String treeModelPrefix, XMLWriter writer) {
+        writer.writeOpenTag(UniformSubtreePruneRegraftParser.UNIFORM_SUBTREE_PRUNE_REGRAFT, // should this be in the class and not the parser?
+            new Attribute[]{
+                getWeightAttribute(operator.getWeight())
+            });
+        writer.writeIDref(ConstrainedTreeModel.CONSTRAINED_TREE_MODEL, treeModelPrefix + DefaultTreeModel.TREE_MODEL);
+        writer.writeCloseTag(UniformSubtreePruneRegraftParser.UNIFORM_SUBTREE_PRUNE_REGRAFT);
+    }
+
+
     // tuneable version of FHSPR but not currently being used
     private void writeSubtreeJumpOperator(Operator operator, String treeModelPrefix, XMLWriter writer) {
         writer.writeOpenTag(SubtreeJumpOperatorParser.SUBTREE_JUMP,

src/dr/app/beauti/generator/TreeLikelihoodGenerator.java

@@ -29,19 +29,29 @@
 
 import dr.app.beauti.types.ClockType;
 import dr.evomodel.tree.DefaultTreeModel;
+import dr.evomodelxml.bigfasttree.thorney.ConstrainedBranchLengthProviderParser;
+import dr.evomodelxml.bigfasttree.thorney.PoissonBranchLengthLikelihoodParser;
+import dr.evomodelxml.bigfasttree.thorney.ThorneyTreeLikelihoodParser;
 import dr.evomodelxml.treedatalikelihood.TreeDataLikelihoodParser;
 import dr.evomodelxml.treelikelihood.MarkovJumpsTreeLikelihoodParser;
 import dr.app.beauti.components.ComponentFactory;
 import dr.app.beauti.components.ancestralstates.AncestralStatesComponentOptions;
 import dr.app.beauti.options.*;
 import dr.app.beauti.util.XMLWriter;
 import dr.evolution.datatype.DataType;
+import dr.evolution.tree.TreeUtils;
+import dr.evomodel.bigfasttree.thorney.ConstrainedTreeBranchLengthProvider;
+import dr.evomodel.bigfasttree.thorney.ConstrainedTreeModel;
+import dr.evomodel.bigfasttree.thorney.MutationBranchMap;
+import dr.evomodel.bigfasttree.thorney.PoissonBranchLengthLikelihoodDelegate;
 import dr.evomodel.branchratemodel.BranchRateModel;
 import dr.oldevomodel.sitemodel.GammaSiteModel;
 import dr.oldevomodel.sitemodel.SiteModel;
 import dr.oldevomodelxml.treelikelihood.AncestralStateTreeLikelihoodParser;
 import dr.oldevomodelxml.treelikelihood.TreeLikelihoodParser;
 import dr.evoxml.AlignmentParser;
+import dr.evoxml.NewickParser;
+import dr.evoxml.SimpleTreeParser;
 import dr.evoxml.SitePatternsParser;
 import dr.util.Attribute;
 import dr.xml.XMLParser;
@@ -74,13 +84,81 @@ public void writeAllTreeLikelihoods(XMLWriter writer) throws GeneratorException
 
                     // if the partition isn't an instanceof PartitionData then it doesn't
                     // need a TreeLikelihood (it is probably a Tree Partition)
-//                } else {
-//                    throw new GeneratorException("Unrecognized partition:\n" + partition.getName());
+               } else {
+                   if(partition instanceof TreePartitionData) {
+                        PartitionTreeModel model = partition.getPartitionTreeModel();
+                        if(model.isUsingThorneyBEAST()){
+                            writeThorneyTreeLikelihood((TreePartitionData) partition, writer);
+                            writer.writeText("");
+                        }
+
+                   }
                 }
             }
         }
     }
 
+    public void writeThorneyTreeLikelihood(TreePartitionData partition,XMLWriter writer){
+       // write mutation branch map
+
+/**
+ *     <simpleMutationBranchMap id="branchMutationCounts" scale="29903.0">
+        <constrainedTreeModel idref="treeModel"/>
+        <dataTree>
+            <tree idref="dataTree"/>
+        </dataTree>
+    </simpleMutationBranchMap>
+ */
+
+        PartitionTreeModel treeModel = partition.getPartitionTreeModel();
+        PartitionClockModel clockModel = partition.getPartitionClockModel();
+        String prefix = treeModel.getPrefix() + clockModel.getPrefix(); // use the treemodel prefix
+
+        String idString = prefix + "treeLikelihood";
+        String branchMapIdD = prefix + "branchMutationCounts";
+
+
+        writer.writeComment("A map between the branches in the tree model and the data providing the number of mutations along each branch.");
+        writer.writeTag(ConstrainedBranchLengthProviderParser.MUTATION_BRANCH_MAP, new Attribute[]{
+            new Attribute.Default<>(XMLParser.ID,branchMapIdD),
+            new Attribute.Default<Double>("scale",treeModel.getThorneyScaler()),
+        },false);
+            writer.writeIDref(ConstrainedTreeModel.CONSTRAINED_TREE_MODEL,prefix+ DefaultTreeModel.TREE_MODEL);
+            writer.writeTag(ConstrainedBranchLengthProviderParser.DATA_TREE, new Attribute[]{}, false);
+                writer.writeOpenTag(
+                            NewickParser.NEWICK,
+                            new Attribute[]{
+                                new Attribute.Default<>(XMLParser.ID, prefix + "dataTree"),
+                                new Attribute.Default<Boolean>(NewickParser.USING_HEIGHTS, true),
+                                new Attribute.Default<Boolean>(NewickParser.USING_DATES, false)
+                            }
+                    );
+                        writer.writeText(TreeUtils.newick(treeModel.getTreePartitionData().getTrees().getTrees().get(0)));
+                    writer.writeCloseTag(NewickParser.NEWICK);
+            writer.writeCloseTag(ConstrainedBranchLengthProviderParser.DATA_TREE);
+        writer.writeCloseTag(ConstrainedBranchLengthProviderParser.MUTATION_BRANCH_MAP);
+
+
+
+        writer.writeOpenTag(ThorneyTreeLikelihoodParser.THORNEY_DATA_LIKELIHOOD_DELEGATE,  new Attribute[]{
+            new Attribute.Default<String>(XMLParser.ID, idString)
+        });
+
+        writer.writeIDref(ConstrainedTreeModel.CONSTRAINED_TREE_MODEL, prefix + DefaultTreeModel.TREE_MODEL);
+        writer.writeIDref(ConstrainedBranchLengthProviderParser.MUTATION_BRANCH_MAP, branchMapIdD );
+
+
+        writer.writeTag( PoissonBranchLengthLikelihoodParser.POISSON_BRANCH_LENGTH_LIKELIHOOD,
+        new Attribute[] {
+            new Attribute.Default<>(XMLParser.ID, prefix + "branchMutationLikelihood"),
+            new Attribute.Default<Double>("scale", treeModel.getThorneyScaler()),
+    }, true);
+
+        ClockModelGenerator.writeBranchRatesModelRef(clockModel, writer);
+
+        writer.writeCloseTag(ThorneyTreeLikelihoodParser.THORNEY_DATA_LIKELIHOOD_DELEGATE);
+    }
+
     public void writeTreeDataLikelihood(List<PartitionData> partitions, XMLWriter writer) {
 
         PartitionSubstitutionModel substModel = partitions.get(0).getPartitionSubstitutionModel();
@@ -287,6 +365,16 @@ public void writeTreeLikelihoodReferences(XMLWriter writer) {
             AncestralStatesComponentOptions ancestralStatesOptions = (AncestralStatesComponentOptions) options
                     .getComponentOptions(AncestralStatesComponentOptions.class);
 
+            // TODO moving poisson branch length to a substitution model would make this logic simplier
+
+            PartitionTreeModel treeModel = partition.getPartitionTreeModel();
+            if(treeModel.isUsingThorneyBEAST()){
+                PartitionClockModel clockModel = partition.getPartitionClockModel();
+                String prefix = treeModel.getPrefix() + clockModel.getPrefix(); // use the treemodel prefix
+                String idString = prefix + "treeLikelihood"; // Line 117
+                writer.writeIDref(ThorneyTreeLikelihoodParser.THORNEY_DATA_LIKELIHOOD_DELEGATE, idString);
+
+            }
             PartitionSubstitutionModel model = partition.getPartitionSubstitutionModel();
 
             if (model == null) {