<commit_before>package starbeast2;

import org.apache.commons.math.MathException;

import beast.core.Input;
import beast.core.parameter.IntegerParameter;
import beast.core.parameter.RealParameter;
import beast.evolution.tree.Node;
import beast.evolution.tree.TreeInterface;
import beast.evolution.branchratemodel.BranchRateModel;
import beast.math.distributions.LogNormalDistributionModel;

public class DiscreteRates extends BranchRateModel.Base implements SpeciesTreeRates {
    final public Input<TreeInterface> treeInput = new Input<>("tree", "(Species) tree to apply per-branch rates to.", Input.Validate.REQUIRED);
    final public Input<Boolean> estimateRootInput = new Input<>("estimateRoot", "Estimate rate of the root branch.", false);
    final public Input<Boolean> noCacheInput = new Input<>("noCache", "Always recalculate branch rates.", false);
    final public Input<IntegerParameter> branchRatesInput = new Input<>("rates", "Discrete per-branch rates.", Input.Validate.REQUIRED);
    final public Input<LogNormalDistributionModel> rateDistributionInput = new Input<>("distr", "The distribution governing the rates among branches. Must have mean of 1.", Input.Validate.REQUIRED);

    private int nBins;
    private double currentLogNormalStdev;
    private double storedLogNormalStdev;
    private double[] binRates;
    private double[] storedBinRates;
    private double[] ratesArray;
    private double[] storedRatesArray;

    private int nEstimatedRates;
    private int rootNodeNumber;
    private boolean estimateRoot;
    private boolean noCache;
    private boolean needsUpdate;
    private boolean binRatesNeedsUpdate;

    @Override
    public boolean requiresRecalculation() {
        // This check is not reliable
        /* if (rateDistributionInput.get().isDirtyCalculation()) {
            binRatesNeedsUpdate = true;
        } else {
            binRatesNeedsUpdate = false;
        } */

        final double proposedLogNormalStdev = rateDistributionInput.get().SParameterInput.get().getValue();
        if (proposedLogNormalStdev != currentLogNormalStdev) {
            binRatesNeedsUpdate = true;
            currentLogNormalStdev = proposedLogNormalStdev;
        } else {
            binRatesNeedsUpdate = false;
        }

        needsUpdate = binRatesNeedsUpdate || branchRatesInput.isDirty() || meanRateInput.isDirty();
        return needsUpdate;
    }

    @Override
    public void store() {
        storedLogNormalStdev = currentLogNormalStdev;
        System.arraycopy(binRates, 0, storedBinRates, 0, binRates.length);
        System.arraycopy(ratesArray, 0, storedRatesArray, 0, ratesArray.length);
        super.store();
    }

    @Override
    public void restore() {
        double tmpLogNormalStdev = currentLogNormalStdev;
        double[] tmpBinRates = binRates;
        double[] tmpRatesArray = ratesArray;

        currentLogNormalStdev = storedLogNormalStdev;
        binRates = storedBinRates;
        ratesArray = storedRatesArray;
        
        storedLogNormalStdev = tmpLogNormalStdev;
        storedBinRates = tmpBinRates;
        storedRatesArray = tmpRatesArray;

        super.restore();
    }

    @Override
    public void initAndValidate() {
        final IntegerParameter branchRates = branchRatesInput.get();
        final TreeInterface speciesTree = treeInput.get();
        final Node[] speciesNodes = speciesTree.getNodesAsArray();
        estimateRoot = estimateRootInput.get().booleanValue();
        noCache = noCacheInput.get().booleanValue();
        rootNodeNumber = speciesTree.getRoot().getNr();
        ratesArray = new double[speciesNodes.length];
        storedRatesArray = new double[speciesNodes.length];

        if (estimateRoot) {
            nEstimatedRates = speciesNodes.length;
        } else {
            nEstimatedRates = speciesNodes.length - 1;
        }

        nBins = nEstimatedRates;

        branchRates.setDimension(nEstimatedRates);
        branchRates.setLower(0);
        branchRates.setUpper(nBins - 1);

        currentLogNormalStdev = -1.0;
        storedLogNormalStdev = -1.0;

        binRates = new double[nBins];<commit_msg>Added support for user-supplied number of discrete rate categories<commit_after>package starbeast2;

import org.apache.commons.math.MathException;

import beast.core.Input;
import beast.core.parameter.IntegerParameter;
import beast.core.parameter.RealParameter;
import beast.evolution.tree.Node;
import beast.evolution.tree.TreeInterface;
import beast.evolution.branchratemodel.BranchRateModel;
import beast.math.distributions.LogNormalDistributionModel;

public class DiscreteRates extends BranchRateModel.Base implements SpeciesTreeRates {
    final public Input<TreeInterface> treeInput = new Input<>("tree", "(Species) tree to apply per-branch rates to.", Input.Validate.REQUIRED);
    final public Input<Integer> nBinsInput = new Input<>("nBins", "Number of discrete branch rate bins (default is equal to the number of estimated branch rates).", -1);
    final public Input<Boolean> estimateRootInput = new Input<>("estimateRoot", "Estimate rate of the root branch.", false);
    final public Input<Boolean> noCacheInput = new Input<>("noCache", "Always recalculate branch rates.", false);
    final public Input<IntegerParameter> branchRatesInput = new Input<>("rates", "Discrete per-branch rates.", Input.Validate.REQUIRED);
    final public Input<LogNormalDistributionModel> rateDistributionInput = new Input<>("distr", "The distribution governing the rates among branches. Must have mean of 1.", Input.Validate.REQUIRED);

    private int nBins;
    private double currentLogNormalStdev;
    private double storedLogNormalStdev;
    private double[] binRates;
    private double[] storedBinRates;
    private double[] ratesArray;
    private double[] storedRatesArray;

    private int nEstimatedRates;
    private int rootNodeNumber;
    private boolean estimateRoot;
    private boolean noCache;
    private boolean needsUpdate;
    private boolean binRatesNeedsUpdate;

    @Override
    public boolean requiresRecalculation() {
        // This check is not reliable
        /* if (rateDistributionInput.get().isDirtyCalculation()) {
            binRatesNeedsUpdate = true;
        } else {
            binRatesNeedsUpdate = false;
        } */

        final double proposedLogNormalStdev = rateDistributionInput.get().SParameterInput.get().getValue();
        if (proposedLogNormalStdev != currentLogNormalStdev) {
            binRatesNeedsUpdate = true;
            currentLogNormalStdev = proposedLogNormalStdev;
        } else {
            binRatesNeedsUpdate = false;
        }

        needsUpdate = binRatesNeedsUpdate || branchRatesInput.isDirty() || meanRateInput.isDirty();
        return needsUpdate;
    }

    @Override
    public void store() {
        storedLogNormalStdev = currentLogNormalStdev;
        System.arraycopy(binRates, 0, storedBinRates, 0, binRates.length);
        System.arraycopy(ratesArray, 0, storedRatesArray, 0, ratesArray.length);
        super.store();
    }

    @Override
    public void restore() {
        double tmpLogNormalStdev = currentLogNormalStdev;
        double[] tmpBinRates = binRates;
        double[] tmpRatesArray = ratesArray;

        currentLogNormalStdev = storedLogNormalStdev;
        binRates = storedBinRates;
        ratesArray = storedRatesArray;
        
        storedLogNormalStdev = tmpLogNormalStdev;
        storedBinRates = tmpBinRates;
        storedRatesArray = tmpRatesArray;

        super.restore();
    }

    @Override
    public void initAndValidate() {
        final IntegerParameter branchRates = branchRatesInput.get();
        final TreeInterface speciesTree = treeInput.get();
        final Node[] speciesNodes = speciesTree.getNodesAsArray();
        estimateRoot = estimateRootInput.get().booleanValue();
        noCache = noCacheInput.get().booleanValue();
        rootNodeNumber = speciesTree.getRoot().getNr();
        ratesArray = new double[speciesNodes.length];
        storedRatesArray = new double[speciesNodes.length];

        if (estimateRoot) {
            nEstimatedRates = speciesNodes.length;
        } else {
            nEstimatedRates = speciesNodes.length - 1;
        }

        final int nBinsSupplied = nBinsInput.get().intValue();
        nBins = (nBinsSupplied <= 0) ? nEstimatedRates : nBinsSupplied;

        branchRates.setDimension(nEstimatedRates);
        branchRates.setLower(0);
        branchRates.setUpper(nBins - 1);

        currentLogNormalStdev = -1.0;
        storedLogNormalStdev = -1.0;

        binRates = new double[nBins];
{'fraction_non_alphanumeric': 0.08881922675026123, 'fraction_numerical': 0.0032508998026239406, 'mean_word_length': 3.3308195072900957, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6939714', 'n_tokens_mistral': 2378, 'n_tokens_neox': 2327, 'n_words': 555}