Skip to content

Commit

Permalink
Merge branch 'Simulator-Piecewise-model' into hmc-clock
Browse files Browse the repository at this point in the history
  • Loading branch information
xji3 committed Dec 2, 2024
2 parents a8677f1 + f8be375 commit 08ab0db
Show file tree
Hide file tree
Showing 9 changed files with 283 additions and 8 deletions.
1 change: 1 addition & 0 deletions src/dr/app/beast/development_parsers.properties
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,7 @@ dr.evomodelxml.continuous.RestrictedPartialsParser
dr.evomodelxml.coalescent.BNPRSamplingLikelihoodParser
dr.inferencexml.distribution.RandomWalkGeneratorParser
dr.inference.operators.RandomWalkGammaPrecisionGibbsOperator
dr.evomodelxml.coalescent.CoalescentRescalerParser

# Clamps
dr.evomodelxml.tree.AncestralTraitTreeModelParser
Expand Down
2 changes: 1 addition & 1 deletion src/dr/evolution/coalescent/CoalescentSimulator.java
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ public SimpleNode[] simulateCoalescent(SimpleNode[] nodes, DemographicFunction d
currentHeight = getMinimumInactiveHeight();
setCurrentHeight(currentHeight);
}

// TODO should it enforce this?
if(!enforceMaxHeight){
// nextCoalescentHeight = currentHeight + DemographicFunction.Utils.getMedianInterval(demographic, getActiveNodeCount(), currentHeight);
nextCoalescentHeight = currentHeight + DemographicFunction.Utils.getSimulatedInterval(demographic,
Expand Down
41 changes: 40 additions & 1 deletion src/dr/evolution/coalescent/PiecewiseConstantPopulation.java
Original file line number Diff line number Diff line change
Expand Up @@ -148,9 +148,48 @@ public double getIntensity(double t) {
}

public double getInverseIntensity(double x) {
throw new RuntimeException("Not implemented!");
// TODO write down error cases
// if (cif != null) {
// if (!intensitiesKnown) {
// setIntervals(intervals, thetas);
// }
//
// int epoch = Collections.binarySearch(cif, x);
//
// if (epoch < 0) {
// epoch = -epoch - 1;
//
// if (epoch > 0) {
// return endTime.get(epoch - 1) + getInverseIntensity(epoch, x - cif.get(epoch - 1));
// } else {
// assert epoch == 0;
// return getInverseIntensity(0, x);
// }
// } else {
// return endTime.get(epoch);
// }
// } else {

double inverseIntensity = 0.0;
int epoch = 0;
double x1 = x;

while (x1 > getIntensity(epoch)) {
x1 -= getIntensity(epoch);
epoch += 1;
if (epoch == intervals.length) {
break;
}
}

inverseIntensity += getEpochDuration(epoch) + x1 * getEpochDemographic(epoch + 1);
// inverseIntensity += getInverseIntensity(epoch, x1);

return inverseIntensity;
// }
}


public double getUpperBound(int i) {
return 1e9;
}
Expand Down
81 changes: 81 additions & 0 deletions src/dr/evomodel/coalescent/CoalescentRescaler.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
/*
* CoalescentRescaler.java
*
* Copyright © 2002-2024 the BEAST Development Team
* http://beast.community/about
*
* This file is part of BEAST.
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership and licensing.
*
* BEAST is free software; you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as
* published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* BEAST is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with BEAST; if not, write to the
* Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
* Boston, MA 02110-1301 USA
*
*/

package dr.evomodel.coalescent;

import dr.evolution.coalescent.ConstantPopulation;
import dr.evolution.tree.Tree;
import dr.evomodel.bigfasttree.BigFastTreeIntervals;
import dr.evomodel.coalescent.demographicmodel.ConstantPopulationModel;
import dr.evomodel.coalescent.demographicmodel.RescaleAwareDemographic;
import dr.evomodel.tree.TreeModel;

/**
* Simulates a set of coalescent intervals given a demographic model.
*
* @author Xiang Ji
* @author Marc Suchard
* @author Filippo Monti
*/
public class CoalescentRescaler {

private TreeModel tree;

private ConstantPopulationModel constantPopulationModel;

private RescaleAwareDemographic rescaleAwareDemographic;

private BigFastTreeIntervals bigFastTreeIntervals;

public CoalescentRescaler(TreeModel tree, ConstantPopulationModel constantPopulationModel, RescaleAwareDemographic rescaleAwareDemographic) {
this.tree = tree;
this.constantPopulationModel = constantPopulationModel;
this.rescaleAwareDemographic = rescaleAwareDemographic;
this.bigFastTreeIntervals = new BigFastTreeIntervals("constant.population.size.tree", tree);
}

public TreeModel rescaleTree() {
double[] intensityTimeProducts = new double[bigFastTreeIntervals.getIntervalCount()];
double[] intervals = new double[bigFastTreeIntervals.getIntervalCount()];
for (int i = 0; i < bigFastTreeIntervals.getIntervalCount(); i++) {
intervals[i] = bigFastTreeIntervals.getInterval(i);
intensityTimeProducts[i] = intervals[i] / ((ConstantPopulation) constantPopulationModel.getDemographicFunction()).getN0();
}
double[] scaledIntervals = rescaleAwareDemographic.rescaleInterval(intervals, intensityTimeProducts);

double height = 0;
for (int i = 0; i < scaledIntervals.length; i++) {
height += scaledIntervals[i];
int[] intervalNodeNumbers = bigFastTreeIntervals.getNodeNumbersForInterval(i);
tree.setNodeHeightQuietly(tree.getNode(intervalNodeNumbers[1]), height);
}
tree.pushTreeChangedEvent();

return tree;
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -1418,4 +1418,4 @@ public List<Citation> getCitations() {
)
);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@
* @author Alexei Drummond
* @author Andrew Rambaut
*/
public class PiecewisePopulationModel extends DemographicModel {
public class PiecewisePopulationModel extends DemographicModel implements RescaleAwareDemographic {

//
// Public stuff
Expand Down Expand Up @@ -73,7 +73,7 @@ public PiecewisePopulationModel(String name, Parameter N0Parameter, double[] epo
if (isLinear) {
piecewiseFunction = new PiecewiseLinearPopulation(epochLengths, new double[N0Parameter.getDimension()], units);
} else {
piecewiseFunction = new PiecewiseConstantPopulation(epochLengths, new double[N0Parameter.getDimension()], units);
piecewiseFunction = new PiecewiseConstantPopulation(epochLengths, N0Parameter.getParameterValues(), units);
}
}

Expand Down Expand Up @@ -166,6 +166,38 @@ protected void restoreState() {
protected void acceptState() {
} // no additional state needs accepting

@Override
public double[] rescaleInterval(double[] intervals, double[] targetIntensityIntervalProducts) {

if (!(piecewiseFunction instanceof PiecewiseConstantPopulation)) {
throw new RuntimeException("Not yet implemented!");
}

PiecewiseConstantPopulation piecewiseConstantPopulation = (PiecewiseConstantPopulation) piecewiseFunction;

double[] scaledIntervals = new double[intervals.length];

int currentEpochIndex = 0;
for (int i = 0; i < intervals.length; i++) {
if (intervals[i] == 0) {
scaledIntervals[i] = 0;
} else {
double accumulatedProduct = 0;
double accumulatedIntervalLength = 0;
while(accumulatedProduct + piecewiseConstantPopulation.getEpochDuration(currentEpochIndex) / piecewiseConstantPopulation.getEpochDemographic(currentEpochIndex) < targetIntensityIntervalProducts[i]) {
accumulatedProduct += piecewiseConstantPopulation.getEpochDuration(currentEpochIndex) / piecewiseConstantPopulation.getEpochDemographic(currentEpochIndex);
accumulatedIntervalLength += piecewiseConstantPopulation.getEpochDuration(currentEpochIndex);
currentEpochIndex++;
}
final double remainingProduct = targetIntensityIntervalProducts[i] - accumulatedProduct;
accumulatedIntervalLength += remainingProduct * piecewiseConstantPopulation.getEpochDemographic(currentEpochIndex);
scaledIntervals[i] = accumulatedIntervalLength;
}
}

return scaledIntervals;
}

public class GrowthRateStatistic extends Statistic.Abstract {

public GrowthRateStatistic() {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
/*
* RescaleAwareDemographic.java
*
* Copyright © 2002-2024 the BEAST Development Team
* http://beast.community/about
*
* This file is part of BEAST.
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership and licensing.
*
* BEAST is free software; you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as
* published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* BEAST is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with BEAST; if not, write to the
* Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
* Boston, MA 02110-1301 USA
*
*/

package dr.evomodel.coalescent.demographicmodel;

public interface RescaleAwareDemographic {

double[] rescaleInterval(double[] intervals, double[] targetIntensityIntervalProducts);

}
83 changes: 83 additions & 0 deletions src/dr/evomodelxml/coalescent/CoalescentRescalerParser.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
/*
* CoalescentRescalerParser.java
*
* Copyright © 2002-2024 the BEAST Development Team
* http://beast.community/about
*
* This file is part of BEAST.
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership and licensing.
*
* BEAST is free software; you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as
* published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* BEAST is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with BEAST; if not, write to the
* Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
* Boston, MA 02110-1301 USA
*
*/

package dr.evomodelxml.coalescent;

import dr.evolution.coalescent.ConstantPopulation;
import dr.evolution.tree.Tree;
import dr.evolution.util.TaxonList;
import dr.evomodel.coalescent.CoalescentRescaler;
import dr.evomodel.coalescent.demographicmodel.ConstantPopulationModel;
import dr.evomodel.coalescent.demographicmodel.DemographicModel;
import dr.evomodel.coalescent.demographicmodel.RescaleAwareDemographic;
import dr.evomodel.tree.DefaultTreeModel;
import dr.evomodel.tree.TreeModel;
import dr.xml.*;

public class CoalescentRescalerParser extends AbstractXMLObjectParser {

public static final String COALESCENT_RESCALER = "coalescentRescaler";

@Override
public Object parseXMLObject(XMLObject xo) throws XMLParseException {
TreeModel tree = (TreeModel) xo.getChild(TreeModel.class);
if (tree == null) {
tree = new DefaultTreeModel((Tree) xo.getChild(Tree.class));
}
ConstantPopulationModel constantPopulation = (ConstantPopulationModel) xo.getChild(ConstantPopulationModel.class);
RescaleAwareDemographic rescaleAwareDemographic = (RescaleAwareDemographic) xo.getChild(RescaleAwareDemographic.class);
CoalescentRescaler coalescentRescaler = new CoalescentRescaler(tree, constantPopulation, rescaleAwareDemographic);
return coalescentRescaler.rescaleTree();
}

@Override
public XMLSyntaxRule[] getSyntaxRules() {
return rules;
}

private final XMLSyntaxRule[] rules = {
new XORRule(new ElementRule(TreeModel.class),
new ElementRule(Tree.class)),
new ElementRule(ConstantPopulationModel.class, 0, Integer.MAX_VALUE),
new ElementRule(RescaleAwareDemographic.class, 0, Integer.MAX_VALUE),
};

@Override
public String getParserDescription() {
return "";
}

@Override
public Class getReturnType() {
return TreeModel.class;
}

@Override
public String getParserName() {
return COALESCENT_RESCALER;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,13 @@ public Object parseXMLObject(XMLObject xo) throws XMLParseException {
return new PiecewisePopulationModel(PIECEWISE_POPULATION, epochSizes, epochWidths, isLinear, units);
} else {
Parameter populationSize = (Parameter) xo.getElementFirstChild(POPULATION_SIZE);
Parameter growthRates = (Parameter) xo.getElementFirstChild(GROWTH_RATES);
return new PiecewisePopulationModel(PIECEWISE_POPULATION, populationSize, growthRates, epochWidths, units);
if (xo.hasChildNamed(GROWTH_RATES)) {
Parameter growthRates = (Parameter) xo.getElementFirstChild(GROWTH_RATES);
return new PiecewisePopulationModel(PIECEWISE_POPULATION, populationSize, growthRates, epochWidths, units);
} else {
return new PiecewisePopulationModel(PIECEWISE_POPULATION, populationSize, epochWidths, false, units);

}
}
}

Expand Down Expand Up @@ -99,7 +104,7 @@ public XMLSyntaxRule[] getSyntaxRules() {
new ElementRule(POPULATION_SIZE,
new XMLSyntaxRule[]{new ElementRule(Parameter.class)}),
new ElementRule(GROWTH_RATES,
new XMLSyntaxRule[]{new ElementRule(Parameter.class)})
new XMLSyntaxRule[]{new ElementRule(Parameter.class)}, true)
)
),
new ElementRule(EPOCH_WIDTHS,
Expand Down

0 comments on commit 08ab0db

Please sign in to comment.