source: branches/2994-AutoDiffForIntervals/HeuristicLab.Problems.DataAnalysis.Symbolic/3.4/Interpreter/VectorAutoDiffEvaluator.cs @ 17296

using System;
using System.Collections.Generic;
using System.Linq;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;

namespace HeuristicLab.Problems.DataAnalysis.Symbolic {
  public sealed class VectorAutoDiffEvaluator : InterpreterBase<MultivariateDual<AlgebraicDoubleVector>> {
    private const int BATCHSIZE = 128;
    [ThreadStatic]
    private Dictionary<string, double[]> cachedData;

    [ThreadStatic]
    private IDataset dataset;

    [ThreadStatic]
    private int rowIndex;

    [ThreadStatic]
    private int[] rows;

    [ThreadStatic]
    private Dictionary<ISymbolicExpressionTreeNode, int> node2paramIdx;

    private void InitCache(IDataset dataset) {
      this.dataset = dataset;
      cachedData = new Dictionary<string, double[]>();
      foreach (var v in dataset.DoubleVariables) {
        cachedData[v] = dataset.GetDoubleValues(v).ToArray();
      }
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="tree"></param>
    /// <param name="dataset"></param>
    /// <param name="rows"></param>
    /// <param name="parameterNodes"></param>
    /// <param name="fi">Function output. Must be allocated by the caller.</param>
    /// <param name="jac">Jacobian matrix. Must be allocated by the caller.</param>
    public void Evaluate(ISymbolicExpressionTree tree, IDataset dataset, int[] rows, ISymbolicExpressionTreeNode[] parameterNodes, double[] fi, double[,] jac) {
      if (cachedData == null || this.dataset != dataset) {
        InitCache(dataset);
      }

      int nParams = parameterNodes.Length;
      node2paramIdx = new Dictionary<ISymbolicExpressionTreeNode, int>();
      for (int i = 0; i < parameterNodes.Length; i++) node2paramIdx.Add(parameterNodes[i], i);

      var code = Compile(tree);

      var remainingRows = rows.Length % BATCHSIZE;
      var roundedTotal = rows.Length - remainingRows;

      this.rows = rows;

      for (rowIndex = 0; rowIndex < roundedTotal; rowIndex += BATCHSIZE) {
        Evaluate(code);
        code[0].value.Value.CopyTo(fi, rowIndex, BATCHSIZE);

        // TRANSPOSE into JAC
        var g = code[0].value.Gradient;
        for (int j = 0; j < nParams; ++j) {
          if (g.Elements.TryGetValue(j, out AlgebraicDoubleVector v)) {
            v.CopyColumnTo(jac, j, rowIndex, BATCHSIZE);
          } else {
            for (int r = 0; r < BATCHSIZE; r++) jac[rowIndex + r, j] = 0.0;
          }
        }
      }

      if (remainingRows > 0) {
        Evaluate(code);
        code[0].value.Value.CopyTo(fi, roundedTotal, remainingRows);

        var g = code[0].value.Gradient;
        for (int j = 0; j < nParams; ++j)
          if (g.Elements.TryGetValue(j, out AlgebraicDoubleVector v)) {
            v.CopyColumnTo(jac, j, roundedTotal, remainingRows);
          } else {
            for (int r = 0; r < remainingRows; r++) jac[roundedTotal + r, j] = 0.0;
          }
      }
    }

    protected override void InitializeInternalInstruction(ref Instruction instruction, ISymbolicExpressionTreeNode node) {
      var zero = new AlgebraicDoubleVector(BATCHSIZE);
      instruction.value = new MultivariateDual<AlgebraicDoubleVector>(zero);
    }

    protected override void InitializeTerminalInstruction(ref Instruction instruction, ConstantTreeNode constant) {
      var g_arr = new double[BATCHSIZE];
      if (node2paramIdx.TryGetValue(constant, out var paramIdx)) {
        for (int i = 0; i < BATCHSIZE; i++) g_arr[i] = 1.0;
        var g = new AlgebraicDoubleVector(g_arr);
        instruction.value = new MultivariateDual<AlgebraicDoubleVector>(new AlgebraicDoubleVector(BATCHSIZE), paramIdx, g); // only a single column for the gradient
      } else {
        instruction.value = new MultivariateDual<AlgebraicDoubleVector>(new AlgebraicDoubleVector(BATCHSIZE));
      }

      instruction.dblVal = constant.Value;
      instruction.value.Value.AssignConstant(instruction.dblVal);
    }

    protected override void InitializeTerminalInstruction(ref Instruction instruction, VariableTreeNode variable) {
      double[] data;
      if (cachedData.ContainsKey(variable.VariableName)) {
        data = cachedData[variable.VariableName];
      } else {
        data = dataset.GetReadOnlyDoubleValues(variable.VariableName).ToArray();
        cachedData[variable.VariableName] = (double[])instruction.data;
      }

      var paramIdx = -1;
      if (node2paramIdx.ContainsKey(variable)) {
        paramIdx = node2paramIdx[variable];
        var f = new AlgebraicDoubleVector(BATCHSIZE);
        var g = new AlgebraicDoubleVector(BATCHSIZE);
        instruction.value = new MultivariateDual<AlgebraicDoubleVector>(f, paramIdx, g);
      } else {
        var f = new AlgebraicDoubleVector(BATCHSIZE);
        instruction.value = new MultivariateDual<AlgebraicDoubleVector>(f);
      }

      instruction.dblVal = variable.Weight;
      instruction.data = new object[] { data, paramIdx };
    }

    protected override void LoadVariable(Instruction a) {
      var paramIdx = (int)((object[])a.data)[1];
      var data = (double[])((object[])a.data)[0];

      for (int i = rowIndex; i < rows.Length && i - rowIndex < BATCHSIZE; i++) a.value.Value[i - rowIndex] = data[rows[i]];
      a.value.Scale(a.dblVal);

      if (paramIdx >= 0) {
        // update gradient with variable values
        var g = a.value.Gradient.Elements[paramIdx];
        for (int i = rowIndex; i < rows.Length && i - rowIndex < BATCHSIZE; i++) {
          g[i - rowIndex] = data[rows[i]];
        }
      }
    }
  }
}