Skip to content

Commit

Permalink
add BEAUti template for SelfTuningMCMC
Browse files Browse the repository at this point in the history
  • Loading branch information
@rbouckaert
rbouckaert committed Oct 27, 2024
1 parent ee6f446 commit 7a60e26
Show file tree
Hide file tree
Showing 3 changed files with 104 additions and 11 deletions.
60 changes: 60 additions & 0 deletions fxtemplates/SelfTuningMCMC.xml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
<beast version='2.7'
namespace='beastfx.app.beauti:beastfx.app.inputeditor:beast.pkgmgmt:beast.base.core:beast.base.inference:beast.base.evolution.branchratemodel:beast.base.evolution.speciation:beast.base.evolution.tree.coalescent:beast.base.util:beast.base.math:beast.evolution.nuc:beast.base.evolution.operator:beast.base.inference.operator:beast.base.evolution.sitemodel:beast.base.evolution.substitutionmodel:beast.base.evolution.likelihood:beast.evolution:beast.base.inference.distribution'
templateinfo='template for standard phylogenetic analysis,&lt;br> supporting tip-date analysis and calibrations'>

<map name='connect' reserved='true'>beastfx.app.inputeditor.BeautiConnector</map>
<map name='subtemplate' reserved='true'>beastfx.app.inputeditor.BeautiSubTemplate</map>
<map name='Uniform'>beast.base.inference.distribution.Uniform</map>
<map name='Normal'>beast.base.inference.distribution.Normal</map>
<map name='OneOnX'>beast.base.inference.distribution.OneOnX</map>
<map name='LogNormal'>beast.base.inference.distribution.LogNormalDistributionModel</map>
<map name='Exponential'>beast.base.inference.distribution.Exponential</map>
<map name='Gamma'>beast.base.inference.distribution.Gamma</map>
<map name='Beta'>beast.base.inference.distribution.Beta</map>
<map name='LaplaceDistribution'>beast.base.inference.distribution.LaplaceDistribution</map>
<map name='InverseGamma'>beast.base.inference.distribution.InverseGamma</map>
<map name='prior'>beast.base.inference.distribution.Prior</map>


<mergewith point='runnables'>
<subtemplate id='Self Tuning MCMC' class='beastlabs.evolution.likelihood.SelfTuningMCMC' mainid='SelfTuningMCMC'
collapsedInputs="beastlabs.evolution.likelihood.SelfTuningMCMC.logger"
suppressInputs="beastlabs.evolution.likelihood.SelfTuningMCMC.operator,
beastlabs.evolution.likelihood.SelfTuningMCMC.operatorschedule,
beastlabs.evolution.likelihood.SelfTuningMCMC.state,
beastlabs.evolution.likelihood.SelfTuningMCMC.distribution,
beastlabs.evolution.likelihood.SelfTuningMCMC.init">
<![CDATA[
<run spec="beastlabs.evolution.likelihood.SelfTuningMCMC" id="SelfTuningMCMC" chainLength="10000000">
<state storeEvery='5000' id='state'>
</state>
<distribution spec="CompoundDistribution" id="posterior">
<distribution spec="CompoundDistribution" id="prior">
</distribution>
<distribution spec="CompoundDistribution" id="likelihood" useThreads="true">
</distribution>
</distribution>
<logger id='tracelog' spec="beast.base.inference.Logger" logEvery="1000" fileName="beast.log" sort="smart" sanitiseHeaders='true'>
<model idref='posterior'/>
<log idref="posterior"/>
<log idref="likelihood"/>
<log idref="prior"/>
</logger>
<logger id='screenlog' logEvery="1000">
<log idref="posterior"/>
<ESS spec='ESS' name='log' arg="@posterior"/>
<log idref="likelihood"/>
<log idref="prior"/>
</logger>
</run>
]]>
</subtemplate>

</mergewith>

</beast>

8 changes: 5 additions & 3 deletions src/beastlabs/evolution/likelihood/SelfTuningCompoundDistribution.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import beast.base.evolution.branchratemodel.BranchRateModel;
import beast.base.evolution.likelihood.BeagleTreeLikelihood.PartialsRescalingScheme;
import beast.base.evolution.likelihood.GenericTreeLikelihood;
import beast.base.evolution.likelihood.ThreadedTreeLikelihood;
import beast.base.evolution.likelihood.TreeLikelihood;
import beast.base.evolution.sitemodel.SiteModel;
import beast.base.evolution.tree.TreeInterface;
Expand Down Expand Up @@ -55,7 +56,6 @@ public class SelfTuningCompoundDistribution extends Distribution {

final public Input<Boolean> includeMPTLInput = new Input<>("includeMPTL", "include multi-partition (BEAGLE 3) tree likelihood in configurations", true);
final public Input<Boolean> includeSPTLInput = new Input<>("includeSPTL", "include single-partition (BEAGLE 2) tree likelihood in configurations", true);


class Configuration {
long nrOfSamples;
Expand Down Expand Up @@ -210,6 +210,8 @@ public void initAndValidate() {
initialMeasurement = true;

switchTime = System.currentTimeMillis();


}


Expand All @@ -235,7 +237,7 @@ private MultiPartitionTreeLikelihood createMultiPartitionTreeLikelihood() {
tree = tl0.treeInput.get();
branchRateModel = tl0.branchRateModelInput.get();
useAmbiguities = (boolean) tl0.getInput("useAmbiguities").get();
useTipLikelihoods = (boolean) tl0.getInput("useTipLikelihoods").get();
useTipLikelihoods = tl0 instanceof TreeLikelihood ? (boolean) tl0.getInput("useTipLikelihoods").get() : false;
rescalingScheme = MultiPartitionTreeLikelihood.getRescalingScheme(tl0);
dataType = tl0.dataInput.get().getDataType().toString();
} else {
Expand Down Expand Up @@ -263,7 +265,7 @@ private MultiPartitionTreeLikelihood createMultiPartitionTreeLikelihood() {
tl.getID() + " and " + tl0.getID()+" -- MultiPartitionTreeLikelihood not considered");
return null;
}
if (useTipLikelihoods != (boolean)tl.getInput("useTipLikelihoods").get()) {
if (tl instanceof TreeLikelihood && useTipLikelihoods != (boolean)tl.getInput("useTipLikelihoods").get()) {
Log.warning("All partitions must use tip likelihoods, or ignore tip likelihoods, but found a difference between " +
tl.getID() + " and " + tl0.getID() + " -- MultiPartitionTreeLikelihood not considered");
return null;
Expand Down
47 changes: 39 additions & 8 deletions src/beastlabs/evolution/likelihood/SelfTuningMCMC.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,32 +5,63 @@
import beast.base.core.BEASTInterface;
import beast.base.core.Description;
import beast.base.core.Input;
import beast.base.core.ProgramStatus;
import beast.base.inference.CompoundDistribution;
import beast.base.inference.Logger;
import beast.base.inference.MCMC;

@Description("MCMC that can tune the treelikelihood for efficiency")
public class SelfTuningMCMC extends MCMC {

final public Input<Boolean> useThreadsInput = new Input<>("useThreads", "calculated the distributions in parallel using threads (default true)-- only used if no SelfTuningCompoundDistribution specified.", true);
final public Input<Integer> minNrOfThreadsInput = new Input<>("minThreads","minimum number of threads to use (default 1)-- only used if no SelfTuningCompoundDistribution specified.", 1);
final public Input<Integer> maxNrOfThreadsInput = new Input<>("maxThreads","maximum number of threads to use, if less than 1 the number of threads in BeastMCMC is used (default -1)-- only used if no SelfTuningCompoundDistribution specified.", -1);

final public Input<Long> swithcCountInput = new Input<>("switchCount", "number of milli seconds to calculate likelihood before switching configuration-- only used if no SelfTuningCompoundDistribution specified.", 500l);
final public Input<Long> reconfigCountInput = new Input<>("reconfigCount", "number of times to calculate likelihood before self tuning again-- only used if no SelfTuningCompoundDistribution specified.", 100000l);
final public Input<Integer> stopAfterSamerResultsInput = new Input<>("stopAfterSamerResults", "number of times the same configuration is optimal in a row before stopping to tune-- only used if no SelfTuningCompoundDistribution specified.", 3);

final public Input<Boolean> includeMPTLInput = new Input<>("includeMPTL", "include multi-partition (BEAGLE 3) tree likelihood in configurations-- only used if no SelfTuningCompoundDistribution specified.", true);
final public Input<Boolean> includeSPTLInput = new Input<>("includeSPTL", "include single-partition (BEAGLE 2) tree likelihood in configurations-- only used if no SelfTuningCompoundDistribution specified.", true);


SelfTuningCompoundDistribution stCompoundDistribution;

@Override
public void initAndValidate() {
super.initAndValidate();

stCompoundDistribution = findSTCompoundDistribution(posteriorInput.get());
if (stCompoundDistribution == null) {
// no SelfTuningCompoundDistribution found, so replace CompoundDistribution with id="likelihood"
// with one
stCompoundDistribution = createdSTCompoundDistribution(posteriorInput.get());
if (!ProgramStatus.name.equals("BEAUti")) {
stCompoundDistribution = findSTCompoundDistribution(posteriorInput.get());
if (stCompoundDistribution == null) {
// no SelfTuningCompoundDistribution found, so replace CompoundDistribution with id="likelihood"
// with one
stCompoundDistribution = createdSTCompoundDistribution(posteriorInput.get());
}

for (Logger logger: loggersInput.get()) {
logger.initAndValidate();
}
}

super.initAndValidate();
}

private SelfTuningCompoundDistribution createdSTCompoundDistribution(BEASTInterface o) {
if (o instanceof CompoundDistribution && o.getID().equals("likelihood")) {
CompoundDistribution cd = (CompoundDistribution) o;
// replace cd with SelfTuningCompoundDistribution
SelfTuningCompoundDistribution stcd = new SelfTuningCompoundDistribution();
stcd.initByName("distribution", cd.pDistributions.get(), "ignore", cd.ignoreInput.get());
stcd.initByName("distribution", cd.pDistributions.get(),
"ignore", cd.ignoreInput.get(),
"useThreads", useThreadsInput.get(),
"minThreads",minNrOfThreadsInput.get(),
"maxThreads",maxNrOfThreadsInput.get(),
"switchCount",swithcCountInput.get(),
"reconfigCount",reconfigCountInput.get(),
"stopAfterSamerResults",stopAfterSamerResultsInput.get(),
"includeMPTL",includeMPTLInput.get(),
"includeSPTL",includeSPTLInput.get()
);
for (BEASTInterface out : cd.getOutputs()) {
for (Input<?> in : out.listInputs()) {
if (in.get() instanceof List) {
Expand Down Expand Up @@ -80,5 +111,5 @@ protected void callUserFunction(final long sample) {
robustlyCalcPosterior(posterior);
}
}

}

0 comments on commit 7a60e26

Please sign in to comment.