lphy.base.evolution.birthdeath.SimFBDAge Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of lphy-base Show documentation
The standard library of LPhy, which contains the required generative distributions and basic functions.
The newest version!
package lphy.base.evolution.birthdeath;

import lphy.base.evolution.tree.TimeTree;
import lphy.base.evolution.tree.TimeTreeNode;
import lphy.base.function.tree.PruneTree;
import lphy.core.model.GenerativeDistribution;
import lphy.core.model.RandomVariable;
import lphy.core.model.Value;
import lphy.core.model.annotation.GeneratorCategory;
import lphy.core.model.annotation.GeneratorInfo;
import lphy.core.model.annotation.ParameterInfo;
import lphy.core.simulator.RandomUtils;
import org.apache.commons.math3.random.RandomGenerator;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import static lphy.base.evolution.birthdeath.BirthDeathConstants.*;

/**
 * A Birth-death tree generative distribution
 */
public class SimFBDAge implements GenerativeDistribution {

    private Value birthRate;
    private Value deathRate;
    private Value psiVal;
    private Value fracVal;
    private Value originAge;

    RandomGenerator random;

    private static final int MAX_ATTEMPTS = 1000;

    public SimFBDAge(@ParameterInfo(name = lambdaParamName, description = "per-lineage birth rate.") Value birthRate,
                     @ParameterInfo(name = muParamName, description = "per-lineage death rate.") Value deathRate,
                     @ParameterInfo(name = fracParamName, description = "fraction of extant taxa sampled.") Value fracVal,
                     @ParameterInfo(name = psiParamName, description = "per-lineage sampling-through-time rate.") Value psiVal,
                     @ParameterInfo(name = originAgeParamName, description = "the age of the origin.") Value originAge) {


        this.birthRate = birthRate;
        this.deathRate = deathRate;
        this.fracVal = fracVal;
        this.psiVal = psiVal;
        this.originAge = originAge;

        random = RandomUtils.getRandom();
    }

    @GeneratorInfo(name = "SimFBDAge",
            category = GeneratorCategory.BD_TREE, examples = {"simFBDAge.lphy"},
            description = "A tree of extant species and those sampled through time, which is conceptually embedded in a full species tree produced by a speciation-extinction (birth-death) branching process.
" +
            "Conditioned on origin age.")
    public RandomVariable sample() {

        int nonNullLeafCount = 0;
        TimeTree sampleTree = null;

        int attempts = 0;

        while (nonNullLeafCount < 1 && attempts < MAX_ATTEMPTS) {
            FullBirthDeathTree birthDeathTree = new FullBirthDeathTree(birthRate, deathRate, null, originAge);
            RandomVariable fullTree = birthDeathTree.sample();

            SimFossilsPoisson simFossilsPoisson = new SimFossilsPoisson(fullTree, psiVal);

            Value fullTreeWithFossils = simFossilsPoisson.sample();

            sampleTree = new TimeTree(fullTreeWithFossils.value());

            List leafNodes = new ArrayList<>();

            for (TimeTreeNode node : sampleTree.getNodes()) {
                if (node.isLeaf() && node.getAge() == 0.0) {
                   leafNodes.add(node);
                }
            }

            int toNull = (int)Math.round(leafNodes.size()* (1.0-fracVal.value()));
            List nullList = new ArrayList<>();
            for (int i =0; i < toNull; i++) {
                nullList.add(leafNodes.remove(random.nextInt(leafNodes.size())));
            }
            for (TimeTreeNode node : nullList) {
                node.setId(null);
            }

            nonNullLeafCount = leafNodes.size();
            attempts += 1;
        }

        if (attempts == MAX_ATTEMPTS) throw new RuntimeException("Failed to simulate SimFBDAge after " + MAX_ATTEMPTS + " attempts.");

        PruneTree pruneTree = new PruneTree(new Value<>(null, sampleTree));

        TimeTree tree = pruneTree.apply().value();
        return new RandomVariable<>(null, tree, this);
    }

    @Override
    public double logDensity(TimeTree timeTree) {
        throw new UnsupportedOperationException("Not implemented!");
    }

    @Override
    public Map getParams() {
        return new TreeMap<>() {{
            put(lambdaParamName, birthRate);
            put(muParamName, deathRate);
            put(fracParamName, fracVal);
            put(psiParamName, psiVal);
            put(originAgeParamName, originAge);
        }};
    }

    @Override
    public void setParam(String paramName, Value value) {
        switch (paramName) {
            case lambdaParamName:
                birthRate = value;
                break;
            case muParamName:
                deathRate = value;
                break;
            case fracParamName:
                fracVal = value;
                break;
            case psiParamName:
                psiVal = value;
                break;
            case originAgeParamName:
                originAge = value;
                break;
            default:
                throw new RuntimeException("Unexpected parameter " + paramName);
        }
    }

    public Value getBirthRate() {
        return birthRate;
    }

    public Value getDeathRate() {
        return deathRate;
    }

    public Value getRho() {
        return fracVal;
    }

    public Value getPsi() {
        return psiVal;
    }
}