scripts/builtin/decisionTreePredict.dml

#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

# This script implements random forest prediction for recoded and binned
# categorical and numerical input features.
# Hummingbird paper (https://www.usenix.org/system/files/osdi20-nakandala.pdf).
#
# INPUT:
# ------------------------------------------------------------------------------
# X               Feature matrix in recoded/binned representation
# y               Label matrix in recoded/binned representation,
#                 optional for accuracy evaluation
# ctypes          Row-Vector of column types [1 scale/ordinal, 2 categorical]
# M               Matrix M holding the learned tree in linearized form
#                 see decisionTree() for the detailed tree representation.
# strategy        Prediction strategy, can be one of ["GEMM", "TT", "PTT"],
#                 referring to "Generic matrix multiplication",
#                 "Tree traversal", and "Perfect tree traversal", respectively
# verbose         Flag indicating verbose debug output
# ------------------------------------------------------------------------------
#
# OUTPUT:
# ------------------------------------------------------------------------------
# yhat            Label vector of predictions
# ------------------------------------------------------------------------------

m_decisionTreePredict = function(Matrix[Double] X, Matrix[Double] y = matrix(0,0,0),
    Matrix[Double] ctypes, Matrix[Double] M, String strategy="TT", Boolean verbose = FALSE)
  return (Matrix[Double] yhat)
{
  if( strategy == "TT" )
    yhat = predict_TT(M, X);
  else if( strategy == "GEMM" )
    yhat = predict_GEMM(M, X);
  else {
    print ("No such strategy" + strategy)
    yhat = matrix("0", rows=0, cols=0)
  }
}

predict_TT = function (Matrix[Double] M, Matrix[Double] X) 
  return (Matrix[Double] yhat)
{
  # initialization of model tensors and parameters
  [N, N_L, N_R, N_F, N_T, C] = createTTNodeTensors(M)
  nr = nrow(X); n = ncol(N_L);
  tree_depth = ceiling(log(max(N)+1,2)) # max depth


  Ti = matrix(1, nr, 1); # current nodes (start at root)
  noChange = FALSE; i = 1;
  while( !noChange & i <= tree_depth) {
    P = table(seq(1,nr), Ti, nr, n);
    TF = P %*% t(N_F); # get node feature indexes
    Tv = rowSums(X * table(seq(1,nr),TF,nr,ncol(X))); # get feature values
    Tt = P %*% t(N_T); # get node thresholds
    TL = P %*% t(N_L); # get node left paths
    TR = P %*% t(N_R); # get node right paths
    # pick left or right path for each record separately
    Ti_new = ifelse(Tv <= Tt, TL, TR);
    noChange = (sum(Ti != Ti_new) == 0);
    i = i + 1;
    Ti = Ti_new;
  }

  # extract classes
  yhat = table(seq(1,nr), Ti, nr, n) %*%  C;
}

predict_GEMM = function (Matrix[Double] M, Matrix[Double] X)
  return (Matrix[Double] Y)
{
  # initialization of model tensors and parameters
  [A, B, C, D, E] = createGEMMNodeTensors(M, ncol(X));

  # scoring pipline, evaluating all nodes in parallel
  Y = rowIndexMax(((((X %*% A) < B) %*% C) == D) %*% E);
}

createTTNodeTensors = function( Matrix[Double] M )
  return (Matrix[Double] N, Matrix[Double] N_L, Matrix[Double] N_R,
          Matrix[Double] N_F, Matrix[Double] N_T, Matrix[Double] C)
{
  # all tree nodes (inner and leaf nodes)
  M2 = matrix(M, rows=ncol(M)/2, cols=2);
  NID = seq(1, nrow(M2));
  nI = (M2[,1]!=0 | M2[,2]!=0)
  cnI = cumsum(nI);
  N = t(removeEmpty(target=NID, margin="rows", select=nI));
  n_nodes = ncol(N)

  # left/right child node ids, default self-id
  N_L = t(removeEmpty(target=ifelse(M2[,1]!=0, 2*NID, NID), margin="rows", select=nI));
  N_R = t(removeEmpty(target=ifelse(M2[,1]!=0, 2*NID+1, NID), margin="rows", select=nI));
  # recoding to node vector positions (removed non-existing nodes)
  N_L = t(table(seq(1,n_nodes), t(N_L), n_nodes, nrow(M2)) %*% cnI);
  N_R = t(table(seq(1,n_nodes), t(N_R), n_nodes, nrow(M2)) %*% cnI);

  # node feature IDs (positions) and threshold values
  N_F = t(removeEmpty(target=ifelse(M2[,1]!=0, M2[,1], 1), margin="rows", select=nI));
  N_T = t(removeEmpty(target=ifelse(M2[,1]!=0, M2[,2], 0), margin="rows", select=nI));

  C = removeEmpty(target=M2[,2], margin="rows", select=nI);
}

createGEMMNodeTensors = function( Matrix[Double] M, Int m )
  return (Matrix[Double] A, Matrix[Double] B, Matrix[Double] C,
  Matrix[Double] D, Matrix[Double] E)
{
  #TODO update for new model layout and generalize
  stop("GEMM not fully supported yet");

  nin = sum(M[2,]!=0); # num inner nodes

  # predicate map [#feat x #inodes] and values [1 x #inodes]
  I1 = removeEmpty(target=M[3,], margin="cols");
  A = table(I1, seq(1,nin), m, nin);
  B = removeEmpty(target=M[6,], margin="cols", select=M[2,]!=0);

  # bucket paths [#inodes x #paths] and path sums
  I2 = (M[2,] == 0)
  np = ncol(M) - nin;
  C = matrix("1 -1", rows=1, cols=2); # TODO general case
  D = colSums(max(C, 0));

  # class map [#paths x #classes]
  E = table(seq(1,ncol(C)), t(M[4,(ncol(M)-ncol(C)+1):ncol(M)]));
}