using System; using System.Collections.Generic; using System.Linq; using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding; namespace HeuristicLab.Problems.DataAnalysis.Symbolic { public sealed class VectorAutoDiffEvaluator : InterpreterBase> { private const int BATCHSIZE = 128; [ThreadStatic] private Dictionary cachedData; [ThreadStatic] private IDataset dataset; [ThreadStatic] private int rowIndex; [ThreadStatic] private int[] rows; [ThreadStatic] private Dictionary node2paramIdx; private void InitCache(IDataset dataset) { this.dataset = dataset; cachedData = new Dictionary(); foreach (var v in dataset.DoubleVariables) { cachedData[v] = dataset.GetDoubleValues(v).ToArray(); } } /// /// /// /// /// /// /// /// Function output. Must be allocated by the caller. /// Jacobian matrix. Must be allocated by the caller. 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(); 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(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(new AlgebraicDoubleVector(BATCHSIZE), paramIdx, g); // only a single column for the gradient } else { instruction.value = new MultivariateDual(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(f, paramIdx, g); } else { var f = new AlgebraicDoubleVector(BATCHSIZE); instruction.value = new MultivariateDual(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]]; } } } } }