<commit_before>package starbeast2;

import beast.core.Description;
import beast.core.Input;
import beast.core.Input.Validate;
import beast.core.Operator;
import beast.evolution.alignment.TaxonSet;
import beast.evolution.tree.Node;
import beast.util.Randomizer;

import java.util.ArrayList;
import java.util.List;

@Description("Tree operator which randomly changes the height of a node, " +
        "then reconstructs the tree from node heights.")
public class NodeReheight2 extends Operator {
    public final Input<SpeciesTree> treeInput = new Input<>("tree", "the species tree", Validate.REQUIRED);
    public final Input<TaxonSet> taxonSetInput = new Input<>("taxonset", "taxon set describing species tree taxa and their gene trees", Validate.REQUIRED);
    public final Input<List<GeneTree>> geneTreesInput = new Input<>("geneTree", "list of gene trees that constrain species tree movement", new ArrayList<>());

    final private double window = 1.0;

    private int lastIndex;
    private int trueBifurcationCount;
    private Node[] canonicalOrder;
    private int[] trueBifurcations;
    private double[] nodeHeights;

    @Override
    public void initAndValidate() {
        final int nodeCount = treeInput.get().getNodeCount();
        canonicalOrder = new Node[nodeCount];
        trueBifurcations = new int[nodeCount];
        nodeHeights = new double[nodeCount];
    }

    @Override
    public double proposal() {
        final Node rootNode = treeInput.get().getRoot();

        // chooseCanonicalOrder also fills in nodeHeights and trueBifurcations
        // the lastIndex will be the last and right-most node index
        lastIndex = -1;
        trueBifurcationCount = 0;
        chooseCanonicalOrder(rootNode);

        // no nodes can be changed by this operator
        if (trueBifurcationCount == 0) {
            return Double.NEGATIVE_INFINITY;
        }

        // pick a bifurcation at random and change the height
        final int chosenNode = trueBifurcations[Randomizer.nextInt(trueBifurcationCount)];
        final double minHeight = Double.min(nodeHeights[chosenNode - 1], nodeHeights[chosenNode + 1]);

        double newHeight = window * (Randomizer.nextDouble() - 0.5);
        if (newHeight < minHeight) { // reflection
            newHeight = minHeight + minHeight - newHeight;
        }

        nodeHeights[chosenNode] = newHeight;
        canonicalOrder[chosenNode].setHeight(newHeight);

        rebuildTree(0, lastIndex);

        int[] visitedCounts = new int[lastIndex + 1];
        checkConsistency(rootNode, visitedCounts);
        for (int i = 0; i <= lastIndex; i++) {
            assert visitedCounts[i] == 1;
        }

        return 0.0;
    }

    /* Performs an in-order traversal of the species tree, randomly shuffling left and right nodes, to produce
       a canonical order in the sense of Mau et al 1999. Also identify which nodes are true bifurcations
       (not fake nodes used for sampled ancestors) */
    private double chooseCanonicalOrder(final Node node) {
        final double height = node.getHeight();

        if (node.isLeaf()) {
            lastIndex++;

            canonicalOrder[lastIndex] = node;
            nodeHeights[lastIndex] = height;

            return height;
        }

        Node canonicalLeft;
        Node canonicalRight;

        if (Randomizer.nextBoolean()) {
            canonicalLeft = node.getLeft();
            canonicalRight = node.getRight();
        } else {
            canonicalLeft = node.getRight();
            canonicalRight = node.getLeft();
        }

        final double leftChildHeight = chooseCanonicalOrder(canonicalLeft);

        lastIndex++;
        final int thisIndex = lastIndex;

        canonicalOrder[thisIndex] = node;

        final double rightChildHeight = chooseCanonicalOrder(canonicalRight);<commit_msg>Fixed some bugs, seems to work now<commit_after>package starbeast2;

import beast.core.Description;
import beast.core.Input;
import beast.core.Input.Validate;
import beast.core.Operator;
import beast.evolution.alignment.TaxonSet;
import beast.evolution.tree.Node;
import beast.util.Randomizer;

import java.util.ArrayList;
import java.util.List;

@Description("Tree operator which randomly changes the height of a node, " +
        "then reconstructs the tree from node heights.")
public class NodeReheight2 extends Operator {
    public final Input<SpeciesTree> treeInput = new Input<>("tree", "the species tree", Validate.REQUIRED);
    public final Input<TaxonSet> taxonSetInput = new Input<>("taxonset", "taxon set describing species tree taxa and their gene trees", Validate.REQUIRED);
    public final Input<List<GeneTree>> geneTreesInput = new Input<>("geneTree", "list of gene trees that constrain species tree movement", new ArrayList<>());

    private int lastIndex;
    private int trueBifurcationCount;
    private Node[] canonicalOrder;
    private int[] trueBifurcations;
    private double[] nodeHeights;

    @Override
    public void initAndValidate() {
        final int nodeCount = treeInput.get().getNodeCount();
        canonicalOrder = new Node[nodeCount];
        trueBifurcations = new int[nodeCount];
        nodeHeights = new double[nodeCount];
    }

    @Override
    public double proposal() {
        final SpeciesTree tree = treeInput.get();
        final Node originalRoot = tree.getRoot();

        // chooseCanonicalOrder also fills in nodeHeights and trueBifurcations
        // the lastIndex will be the last and right-most node index
        lastIndex = -1;
        trueBifurcationCount = 0;
        chooseCanonicalOrder(originalRoot);

        // no nodes can be changed by this operator
        if (trueBifurcationCount == 0) {
            return Double.NEGATIVE_INFINITY;
        }

        // pick a bifurcation at random and change the height
        final int chosenNode = trueBifurcations[Randomizer.nextInt(trueBifurcationCount)];
        final double minHeight = Double.max(nodeHeights[chosenNode - 1], nodeHeights[chosenNode + 1]);

        double newHeight = nodeHeights[chosenNode] + Randomizer.nextDouble() - 0.5;
        if (newHeight < minHeight) { // reflection
            newHeight = minHeight + minHeight - newHeight;
        }

        nodeHeights[chosenNode] = newHeight;
        canonicalOrder[chosenNode].setHeight(newHeight);

        final int rootIndex = rebuildTree(0, lastIndex);
        if (canonicalOrder[rootIndex] != originalRoot) {
            canonicalOrder[rootIndex].setParent(null);
            tree.setRoot(canonicalOrder[rootIndex]);
        }

        /* int[] visitedCounts = new int[lastIndex + 1];
        checkConsistency(canonicalOrder[rootIndex], visitedCounts);
        for (int i = 0; i <= lastIndex; i++) {
            assert visitedCounts[i] == 1;
        } */

        return 0.0;
    }

    /* Performs an in-order traversal of the species tree, randomly shuffling left and right nodes, to produce
       a canonical order in the sense of Mau et al 1999. Also identify which nodes are true bifurcations
       (not fake nodes used for sampled ancestors) */
    private double chooseCanonicalOrder(final Node node) {
        final double height = node.getHeight();

        if (node.isLeaf()) {
            lastIndex++;

            canonicalOrder[lastIndex] = node;
            nodeHeights[lastIndex] = height;

            return height;
        }

        Node canonicalLeft;
        Node canonicalRight;

        if (Randomizer.nextBoolean()) {
            canonicalLeft = node.getLeft();
            canonicalRight = node.getRight();
        } else {
            canonicalLeft = node.getRight();
            canonicalRight = node.getLeft();
        }

        final double leftChildHeight = chooseCanonicalOrder(canonicalLeft);

        lastIndex++;
        final int thisIndex = lastIndex;

        canonicalOrder[thisIndex] = node;
        nodeHeights[thisIndex] = height;

        final double rightChildHeight = chooseCanonicalOrder(canonicalRight);
{'fraction_non_alphanumeric': 0.08898623279098873, 'fraction_numerical': 0.0050062578222778474, 'mean_word_length': 2.8216164514586324, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 19, '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': '2867038', 'n_tokens_mistral': 2202, 'n_tokens_neox': 2086, 'n_words': 638}