<commit_before>/*
 * DataSimulationDelegate.java
 *
 * Copyright (c) 2002-2016 <NAME>, <NAME> and <NAME>
 *
 * 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.treedatalikelihood.preorder;

import beagle.Beagle;
import dr.evolution.alignment.PatternList;
import dr.evolution.tree.*;
import dr.evomodel.siteratemodel.SiteRateModel;
import dr.evomodel.substmodel.SubstitutionModel;
import dr.evomodel.treedatalikelihood.*;
import dr.inference.model.Model;

import java.util.Arrays;
import java.util.List;

import static dr.evolution.tree.TreeTrait.DA.factory;
import static dr.evomodel.treedatalikelihood.preorder.AbstractDiscreteTraitDelegate.MatrixChoice.GRADIENT;
import static dr.evomodel.treedatalikelihood.preorder.AbstractDiscreteTraitDelegate.MatrixChoice.HESSIAN;

/**
 * AbstractDiscreteTraitDelegate - interface for a plugin delegate for data simulation on a tree.
 *
 * @author <NAME>
 * @author <NAME>
 */
public class AbstractDiscreteTraitDelegate extends ProcessSimulationDelegate.AbstractDelegate
        implements TreeTrait.TraitInfo<double[]> {

    public AbstractDiscreteTraitDelegate(String name,
                                         Tree tree,
                                         BeagleDataLikelihoodDelegate likelihoodDelegate) {
        super(name, tree);
        this.likelihoodDelegate = likelihoodDelegate;
        this.beagle = likelihoodDelegate.getBeagleInstance();

        assert (this.likelihoodDelegate.isUsePreOrder()); /// TODO: reinitialize beagle instance if usePreOrder = false

        this.evolutionaryProcessDelegate = likelihoodDelegate.getEvolutionaryProcessDelegate();
        this.siteRateModel = likelihoodDelegate.getSiteRateModel();

        this.patternCount = likelihoodDelegate.getPatternList().getPatternCount();
        this.stateCount = likelihoodDelegate.getPatternList().getDataType().getStateCount();
        this.categoryCount = siteRateModel.getCategoryCount();

        // put preOrder partials right after postOrder partials
        this.preOrderPartialOffset = likelihoodDelegate.getPartialBufferCount();

        // put scaleBuffers for preOrder partials right after those for postOrder partials
        this.preOrderScaleBufferOffset = likelihoodDelegate.getScaleBufferCount();

        this.patternList = likelihoodDelegate.getPatternList();

        this.postOrderPartial = new double[stateCount * patternCount * categoryCount];
        this.preOrderPartial = new double[stateCount * patternCount * categoryCount];
        this.rootPostOrderPartials = new double[stateCount * patternCount * categoryCount];
        this.Q = new double[stateCount * stateCount];
        this.cLikelihood = new double[categoryCount * patternCount];

        this.grandDenominator = new double[patternCount];
        this.grandNumerator = new double[patternCount];
        this.grandNumeratorIncrementLowerBound = new double[patternCount];
        this.grandNumeratorIncrementUpperBound = new double[patternCount];

    }

    @Override
    public void simulate(final int[] operations, final int operationCount,
                         final int rootNodeNumber) {
        //This function updates preOrder Partials for all nodes
        this.simulateRoot(rootNodeNumber); //set up pre-order partials at root node first
        beagle.updatePrePartials(operations, operationCount, Beagle.NONE);  // Update all nodes with no rescaling
        //TODO:error control

        if (COUNT_TOTAL_OPERATIONS) {
            ++simulateCount;
        }
    }

    @Override
    public void setupStatistics() {
        throw new RuntimeException("Not used (?) with BEAGLE");
    }

    @Override
    protected void simulateRoot(int rootNumber) {
        //This function sets preOrderPartials at Root for now.
        final double[] frequencies = evolutionaryProcessDelegate.getRootStateFrequencies();
        double[] rootPreOrderPartial = new double[stateCount * patternCount * categoryCount];
        for (int i = 0; i < patternCount * categoryCount; ++i) {
            System.arraycopy(frequencies, 0, rootPreOrderPartial, i * stateCount, stateCount);
        }
//        InstanceDetails instanceDetails = beagle.getDetails();

        beagle.setPartials(getPreOrderPartialIndex(rootNumber), rootPreOrderPartial); // Call beagle.setPartials()
        //TODO: find the right error message for control
    }

    @Override
    protected void simulateNode(int v0, int v1, int v2, int v3, int v4) {
        throw new RuntimeException("Not used with BEAGLE");
    }

    @Override
    protected void constructTraits(Helper treeTraitHelper) {
        treeTraitHelper.addTrait(factory(this));
        treeTraitHelper.addTrait(new TreeTrait.DA() {
            @Override
            public String getTraitName() {
                return "Hessian";
            }

            @Override
            public Intent getIntent() {
                return null;
            }

            @Override
            public double[] getTrait(Tree tree, NodeRef node) {
                return getHessian(tree, node);
            }
        });
    }

    public static String getName(String name) {
        return "derivative." + name;
    }

    @Override
    public String getTraitName() {
        return getName(name);
    }

    @Override
    public TreeTrait.Intent getTraitIntent() {
        return TreeTrait.Intent.NODE;
    }

    @Override
    public Class getTraitClass() {
        return double[].class;
    }

    public enum MatrixChoice {
        GRADIENT {
            @Override
            public void getMatrix(SubstitutionModel model, double[] matrix) {
                model.getInfinitesimalMatrix(matrix);
            }

            @Override
            public double getRateDifferential(double rate) {
                return rate;
            }
        },
        HESSIAN {
            @Override
            public void getMatrix(SubstitutionModel model, double[] matrix) {
                double[] tmp = new double[matrix.length];
                model.getInfinitesimalMatrix(tmp);
                Arrays.fill(matrix, 0.0);

                final int stateCount = model.getDataType().getStateCount();
                for (int i = 0; i < stateCount; ++i){
                    for ( int j = 0; j < stateCount; ++j){
                        for ( int k = 0; k < stateCount; ++k){
                            matrix[i * stateCount + j] += tmp[i * stateCount + k] * tmp[k * stateCount + j];
                        }
                    }
                }
            }

            @Override
            public double getRateDifferential(double rate) {
                return rate * rate;
            }
        };

        public abstract void getMatrix(SubstitutionModel model, double[] matrix);

        public abstract double getRateDifferential(double rate);

    }

    private double[] getHessian(Tree tree, NodeRef node) {
        final double[] patternGradient = getTrait(tree, node, GRADIENT); // this is 1/L * firstDerivative(L)
        final double[] patternDiagonalHessian  = getTrait(tree, node, HESSIAN); // this is 1/L * secondDerivative(L)
        final double[] patternWeights = patternList.getPatternWeights();
        double[] patternDiagonalLogHessian = new double[patternGradient.length];
        double[] diagonalLogHessian = new double[tree.getNodeCount() - 1];

        for (int i = 0; i < patternGradient.length; ++i){
            patternDiagonalLogHessian[i] = patternDiagonalHessian[i] - patternGradient[i] * patternGradient[i];
        }

        sumOverPatterns(tree, patternWeights, patternDiagonalLogHessian, diagonalLogHessian);
        return diagonalLogHessian;
    }

    private void sumOverPatterns(Tree tree, double[] patternWeights, double[] patternArray, double[] summedArray) {
        int v = 0;
        Arrays.fill(summedArray, 0.0);
        for (int nodeNum = 0; nodeNum < tree.getNodeCount(); ++nodeNum){
            if (!tree.isRoot(tree.getNode(nodeNum))) {
                for (int pattern = 0; pattern < patternCount; pattern++) {
                    summedArray[v] += patternArray[v * patternCount + pattern] * patternWeights[pattern];
                }
                v++;
            }
        }
    }

    @Override
    public double[] getTrait(Tree tree, NodeRef node) {

        if (COUNT_TOTAL_OPERATIONS) {
            ++getTraitCount;
        }

        final double[] patternGradient = getTrait(tree, node, GRADIENT);
        final double[] patternWeights = patternList.getPatternWeights();
        double[] gradient = new double[tree.getNodeCount() - 1];
        sumOverPatterns(tree, patternWeights, patternGradient, gradient);
        return gradient;<commit_msg>move gradient and hessian calculation into beagle<commit_after>/*
 * DataSimulationDelegate.java
 *
 * Copyright (c) 2002-2016 <NAME>, <NAME> and <NAME>
 *
 * 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.treedatalikelihood.preorder;

import beagle.Beagle;
import dr.evolution.alignment.PatternList;
import dr.evolution.tree.*;
import dr.evomodel.siteratemodel.SiteRateModel;
import dr.evomodel.substmodel.SubstitutionModel;
import dr.evomodel.treedatalikelihood.*;
import dr.inference.model.Model;

import java.util.Arrays;
import java.util.List;

import static dr.evolution.tree.TreeTrait.DA.factory;
import static dr.evomodel.treedatalikelihood.preorder.AbstractDiscreteTraitDelegate.MatrixChoice.GRADIENT;
import static dr.evomodel.treedatalikelihood.preorder.AbstractDiscreteTraitDelegate.MatrixChoice.HESSIAN;

/**
 * AbstractDiscreteTraitDelegate - interface for a plugin delegate for data simulation on a tree.
 *
 * @author <NAME>
 * @author <NAME>
 */
public class AbstractDiscreteTraitDelegate extends ProcessSimulationDelegate.AbstractDelegate
        implements TreeTrait.TraitInfo<double[]> {

    public AbstractDiscreteTraitDelegate(String name,
                                         Tree tree,
                                         BeagleDataLikelihoodDelegate likelihoodDelegate) {
        super(name, tree);
        this.likelihoodDelegate = likelihoodDelegate;
        this.beagle = likelihoodDelegate.getBeagleInstance();

        assert (this.likelihoodDelegate.isUsePreOrder()); /// TODO: reinitialize beagle instance if usePreOrder = false

        this.evolutionaryProcessDelegate = likelihoodDelegate.getEvolutionaryProcessDelegate();
        this.siteRateModel = likelihoodDelegate.getSiteRateModel();

        this.patternCount = likelihoodDelegate.getPatternList().getPatternCount();
        this.stateCount = likelihoodDelegate.getPatternList().getDataType().getStateCount();
        this.categoryCount = siteRateModel.getCategoryCount();

        // put preOrder partials right after postOrder partials
        this.preOrderPartialOffset = likelihoodDelegate.getPartialBufferCount();

        // put scaleBuffers for preOrder partials right after those for postOrder partials
        this.preOrderScaleBufferOffset = likelihoodDelegate.getScaleBufferCount();

        this.patternList = likelihoodDelegate.getPatternList();

        this.postOrderPartial = new double[stateCount * patternCount * categoryCount];
        this.preOrderPartial = new double[stateCount * patternCount * categoryCount];
        this.rootPostOrderPartials = new double[stateCount * patternCount * categoryCount];
        this.Q = new double[stateCount * stateCount];
        this.cLikelihood = new double[categoryCount * patternCount];

        this.grandDenominator = new double[patternCount];
        this.grandNumerator = new double[patternCount];
        this.grandNumeratorIncrementLowerBound = new double[patternCount];
        this.grandNumeratorIncrementUpperBound = new double[patternCount];

    }

    @Override
    public void simulate(final int[] operations, final int operationCount,
                         final int rootNodeNumber) {
        //This function updates preOrder Partials for all nodes
        this.simulateRoot(rootNodeNumber); //set up pre-order partials at root node first
        beagle.updatePrePartials(operations, operationCount, Beagle.NONE);  // Update all nodes with no rescaling
        //TODO:error control

        if (COUNT_TOTAL_OPERATIONS) {
            ++simulateCount;
        }
    }

    @Override
    public void setupStatistics() {
        throw new RuntimeException("Not used (?) with BEAGLE");
    }

    @Override
    protected void simulateRoot(int rootNumber) {
        //This function sets preOrderPartials at Root for now.
        final double[] frequencies = evolutionaryProcessDelegate.getRootStateFrequencies();
        double[] rootPreOrderPartial = new double[stateCount * patternCount * categoryCount];
        for (int i = 0; i < patternCount * categoryCount; ++i) {
            System.arraycopy(frequencies, 0, rootPreOrderPartial, i * stateCount, stateCount);
        }
//        InstanceDetails instanceDetails = beagle.getDetails();

        beagle.setPartials(getPreOrderPartialIndex(rootNumber), rootPreOrderPartial); // Call beagle.setPartials()
        //TODO: find the right error message for control
    }

    @Override
    protected void simulateNode(int v0, int v1, int v2, int v3, int v4) {
        throw new RuntimeException("Not used with BEAGLE");
    }

    @Override
    protected void constructTraits(Helper treeTraitHelper) {
        treeTraitHelper.addTrait(factory(this));
        treeTraitHelper.addTrait(new TreeTrait.DA() {
            @Override
            public String getTraitName() {
                return "Hessian";
            }

            @Override
            public Intent getIntent() {
                return null;
            }

            @Override
            public double[] getTrait(Tree tree, NodeRef node) {
                return getHessian(tree, node);
            }
        });
    }

    public static String getName(String name) {
        return "derivative." + name;
    }

    @Override
    public String getTraitName() {
        return getName(name);
    }

    @Override
    public TreeTrait.Intent getTraitIntent() {
        return TreeTrait.Intent.NODE;
    }

    @Override
    public Class getTraitClass() {
        return double[].class;
    }

    public enum MatrixChoice {
        GRADIENT {
            @Override
            public void getMatrix(SubstitutionModel model, double[] matrix) {
                model.getInfinitesimalMatrix(matrix);
            }

            @Override
            public double getRateDifferential(double rate) {
                return rate;
            }
        },
        HESSIAN {
            @Override
            public void getMatrix(SubstitutionModel model, double[] matrix) {
                double[] tmp = new double[matrix.length];
                model.getInfinitesimalMatrix(tmp);
                Arrays.fill(matrix, 0.0);

                final int stateCount = model.getDataType().getStateCount();
                for (int i = 0; i < stateCount; ++i){
                    for ( int j = 0; j < stateCount; ++j){
                        for ( int k = 0; k < stateCount; ++k){
                            matrix[i * stateCount + j] += tmp[i * stateCount + k] * tmp[k * stateCount + j];
                        }
                    }
                }
            }

            @Override
            public double getRateDifferential(double rate) {
                return rate * rate;
            }
        };

        public abstract void getMatrix(SubstitutionModel model, double[] matrix);

        public abstract double getRateDifferential(double rate);

    }

    private double[] getHessian(Tree tree, NodeRef node) {
        double[] patternGradient = new double[patternCount * (tree.getNodeCount() - 1)];
        double[] patternDiagonalHessian = new double[patternCount * (tree.getNodeCount() - 1)];
        getPatternGradientHessian(tree, patternGradient, patternDiagonalHessian);
        final double[] patternWeights = patternList.getPatternWeights();
        double[] patternDiagonalLogHessian = new double[patternGradient.length];
        double[] diagonalLogHessian = new double[tree.getNodeCount() - 1];

        for (int i = 0; i < patternGradient.length; ++i){
            patternDiagonalLogHessian[i] = patternDiagonalHessian[i] - patternGradient[i] * patternGradient[i];
        }

        sumOverPatterns(tree, patternWeights, patternDiagonalLogHessian, diagonalLogHessian);
        return diagonalLogHessian;
    }

    private void sumOverPatterns(Tree tree, double[] patternWeights, double[] patternArray, double[] summedArray) {
        int v = 0;
        Arrays.fill(summedArray, 0.0);
        for (int nodeNum = 0; nodeNum < tree.getNodeCount(); ++nodeNum){
            if (!tree.isRoot(tree.getNode(nodeNum))) {
                for (int pattern = 0; pattern < patternCount; pattern++) {
                    summedArray[v] += patternArray[v * patternCount + pattern] * patternWeights[pattern];
                }
                v++;
            }
        }
    }

    @Override
    public double[] getTrait(Tree tree, NodeRef node) {

        if (COUNT_TOTAL_OPERATIONS) {
            ++getTraitCount;
        }

//        final double[] patternGradientOld = getTrait(tree, node, GRADIENT);
        double[] patternGradient = new double[patternCount * (tree.getNodeCount() - 1)];
        getPatternGradientHessian(tree, patternGradient, null);
        final double[] patternWeights = patternList.getPatternWeights();
        double[] gradient = new double[tree.getNodeCount() - 1];
        sumOverPatterns(tree, patternWeights, patternGradient, gradient);
        return gradient;
    }

    private void getPatternGradientHessian(Tree tree, double[] patternGradient, double[] patternDiagonalHessian) {
        final int[] postBufferIndices = new int[tree.getNodeCount() - 1];
        final int[] preBufferIndices = new int[tree.getNodeCount() - 1];
        final int   rootNumber = tree.getRoot().getNumber();
        final int[] firstDervIndices = new int[tree.getNodeCount() - 1];
        final int[] secondDeriveIndices = new int[tree.getNodeCount() - 1];

        int u = 0;
        for (int nodeNum = 0; nodeNum < tree.getNodeCount(); nodeNum++) {
            if (!tree.isRoot(tree.getNode(nodeNum))) {
                postBufferIndices[u] = getPostOrderPartialIndex(nodeNum);
                preBufferIndices[u]  = getPreOrderPartialIndex(nodeNum);
                firstDervIndices[u]  = evolutionaryProcessDelegate.getInfinitesimalMatrixBufferIndex(nodeNum);
                secondDeriveIndices[u] = evolutionaryProcessDelegate.getSquaredInfinitesimalMatrixBufferIndex(nodeNum);
                u++;
            }
        }
        beagle.calculateEdgeDerivative(postBufferIndices, preBufferIndices, getPostOrderPartialIndex(rootNumber),
                firstDervIndices, secondDeriveIndices, 0, 0, 0, new int[]{Beagle.NONE},
                tree.getNodeCount() - 1, patternGradient, patternDiagonalHessian);
{'fraction_non_alphanumeric': 0.08793128548551463, 'fraction_numerical': 0.004610083952055127, 'mean_word_length': 2.647433628318584, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6990411', 'n_tokens_mistral': 5544, 'n_tokens_neox': 5284, 'n_words': 1438}