lgb.model.dt.tree.R 6.81 KB
Newer Older
1
2
3
#' @name lgb.model.dt.tree
#' @title Parse a LightGBM model json dump
#' @description Parse a LightGBM model json dump into a \code{data.table} structure.
4
#' @param model object of class \code{lgb.Booster}
James Lamb's avatar
James Lamb committed
5
#' @param num_iteration number of iterations you want to predict with. NULL or
James Lamb's avatar
James Lamb committed
6
#'                      <= 0 means use best iteration
7
8
#' @return
#' A \code{data.table} with detailed information about model trees' nodes and leafs.
James Lamb's avatar
James Lamb committed
9
#'
10
#' The columns of the \code{data.table} are:
James Lamb's avatar
James Lamb committed
11
#'
12
#' \itemize{
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
#'  \item{\code{tree_index}: ID of a tree in a model (integer)}
#'  \item{\code{split_index}: ID of a node in a tree (integer)}
#'  \item{\code{split_feature}: for a node, it's a feature name (character);
#'                              for a leaf, it simply labels it as \code{"NA"}}
#'  \item{\code{node_parent}: ID of the parent node for current node (integer)}
#'  \item{\code{leaf_index}: ID of a leaf in a tree (integer)}
#'  \item{\code{leaf_parent}: ID of the parent node for current leaf (integer)}
#'  \item{\code{split_gain}: Split gain of a node}
#'  \item{\code{threshold}: Splitting threshold value of a node}
#'  \item{\code{decision_type}: Decision type of a node}
#'  \item{\code{default_left}: Determine how to handle NA value, TRUE -> Left, FALSE -> Right}
#'  \item{\code{internal_value}: Node value}
#'  \item{\code{internal_count}: The number of observation collected by a node}
#'  \item{\code{leaf_value}: Leaf value}
#'  \item{\code{leaf_count}: The number of observation collected by a leaf}
28
#' }
James Lamb's avatar
James Lamb committed
29
#'
30
#' @examples
James Lamb's avatar
James Lamb committed
31
#'
32
#' data(agaricus.train, package = "lightgbm")
33
34
35
#' train <- agaricus.train
#' dtrain <- lgb.Dataset(train$data, label = train$label)
#'
36
37
38
#' params <- list(
#'   objective = "binary"
#'   , learning_rate = 0.01
39
40
41
42
#'   , num_leaves = 63L
#'   , max_depth = -1L
#'   , min_data_in_leaf = 1L
#'   , min_sum_hessian_in_leaf = 1.0
43
#' )
44
#' model <- lgb.train(params, dtrain, 10L)
45
46
#'
#' tree_dt <- lgb.model.dt.tree(model)
James Lamb's avatar
James Lamb committed
47
#'
48
#' @importFrom data.table := rbindlist
James Lamb's avatar
James Lamb committed
49
#' @importFrom jsonlite fromJSON
50
51
#' @export
lgb.model.dt.tree <- function(model, num_iteration = NULL) {
James Lamb's avatar
James Lamb committed
52

53
  # Dump json model first
54
  json_model <- lgb.dump(model, num_iteration = num_iteration)
James Lamb's avatar
James Lamb committed
55

56
  # Parse json model second
57
58
59
60
61
62
63
  parsed_json_model <- jsonlite::fromJSON(
    json_model
    , simplifyVector = TRUE
    , simplifyDataFrame = FALSE
    , simplifyMatrix = FALSE
    , flatten = FALSE
  )
James Lamb's avatar
James Lamb committed
64

65
  # Parse tree model third
66
  tree_list <- lapply(parsed_json_model$tree_info, single.tree.parse)
James Lamb's avatar
James Lamb committed
67

68
  # Combine into single data.table fourth
69
  tree_dt <- data.table::rbindlist(tree_list, use.names = TRUE)
James Lamb's avatar
James Lamb committed
70

71
72
73
74
  # Substitute feature index with the actual feature name

  # Since the index comes from C++ (which is 0-indexed), be sure
  # to add 1 (e.g. index 28 means the 29th feature in feature_names)
75
  split_feature_indx <- tree_dt[, split_feature] + 1L
76
77
78
79
80

  # Get corresponding feature names. Positions in split_feature_indx
  # which are NA will result in an NA feature name
  feature_names <- parsed_json_model$feature_names[split_feature_indx]
  tree_dt[, split_feature := feature_names]
James Lamb's avatar
James Lamb committed
81

82
  # Return tree
83
  return(tree_dt)
James Lamb's avatar
James Lamb committed
84

85
86
}

James Lamb's avatar
James Lamb committed
87

88
#' @importFrom data.table := data.table rbindlist
89
single.tree.parse <- function(lgb_tree) {
James Lamb's avatar
James Lamb committed
90

91
  # Traverse tree function
Yachen Yan's avatar
Yachen Yan committed
92
  pre_order_traversal <- function(env = NULL, tree_node_leaf, current_depth = 0L, parent_index = NA_integer_) {
James Lamb's avatar
James Lamb committed
93

Yachen Yan's avatar
Yachen Yan committed
94
95
96
    if (is.null(env)) {
      # Setup initial default data.table with default types
      env <- new.env(parent = emptyenv())
97
      env$single_tree_dt <- data.table::data.table(
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
        tree_index = integer(0L)
        , depth = integer(0L)
        , split_index = integer(0L)
        , split_feature = integer(0L)
        , node_parent = integer(0L)
        , leaf_index = integer(0L)
        , leaf_parent = integer(0L)
        , split_gain = numeric(0L)
        , threshold = numeric(0L)
        , decision_type = character(0L)
        , default_left = character(0L)
        , internal_value = integer(0L)
        , internal_count = integer(0L)
        , leaf_value = integer(0L)
        , leaf_count = integer(0L)
113
      )
Yachen Yan's avatar
Yachen Yan committed
114
115
116
      # start tree traversal
      pre_order_traversal(env, tree_node_leaf, current_depth, parent_index)
    } else {
James Lamb's avatar
James Lamb committed
117

Yachen Yan's avatar
Yachen Yan committed
118
119
      # Check if split index is not null in leaf
      if (!is.null(tree_node_leaf$split_index)) {
James Lamb's avatar
James Lamb committed
120

Yachen Yan's avatar
Yachen Yan committed
121
122
123
124
125
126
127
        # update data.table
        env$single_tree_dt <- data.table::rbindlist(l = list(env$single_tree_dt,
                                                             c(tree_node_leaf[c("split_index",
                                                                                "split_feature",
                                                                                "split_gain",
                                                                                "threshold",
                                                                                "decision_type",
128
                                                                                "default_left",
Yachen Yan's avatar
Yachen Yan committed
129
130
131
132
133
134
                                                                                "internal_value",
                                                                                "internal_count")],
                                                               "depth" = current_depth,
                                                               "node_parent" = parent_index)),
                                                    use.names = TRUE,
                                                    fill = TRUE)
James Lamb's avatar
James Lamb committed
135

Yachen Yan's avatar
Yachen Yan committed
136
        # Traverse tree again both left and right
137
138
139
140
141
142
143
144
145
146
147
148
        pre_order_traversal(
          env
          , tree_node_leaf$left_child
          , current_depth = current_depth + 1L
          , parent_index = tree_node_leaf$split_index
        )
        pre_order_traversal(
          env
          , tree_node_leaf$right_child
          , current_depth = current_depth + 1L
          , parent_index = tree_node_leaf$split_index
        )
James Lamb's avatar
James Lamb committed
149

Yachen Yan's avatar
Yachen Yan committed
150
      } else if (!is.null(tree_node_leaf$leaf_index)) {
James Lamb's avatar
James Lamb committed
151

Yachen Yan's avatar
Yachen Yan committed
152
153
154
155
156
157
158
159
160
        # update data.table
        env$single_tree_dt <- data.table::rbindlist(l = list(env$single_tree_dt,
                                                             c(tree_node_leaf[c("leaf_index",
                                                                                "leaf_value",
                                                                                "leaf_count")],
                                                               "depth" = current_depth,
                                                               "leaf_parent" = parent_index)),
                                                    use.names = TRUE,
                                                    fill = TRUE)
James Lamb's avatar
James Lamb committed
161

Yachen Yan's avatar
Yachen Yan committed
162
      }
James Lamb's avatar
James Lamb committed
163

164
    }
Yachen Yan's avatar
Yachen Yan committed
165
    return(env$single_tree_dt)
166
  }
James Lamb's avatar
James Lamb committed
167

168
  # Traverse structure
Yachen Yan's avatar
Yachen Yan committed
169
  single_tree_dt <- pre_order_traversal(tree_node_leaf = lgb_tree$tree_structure)
James Lamb's avatar
James Lamb committed
170

171
  # Store index
172
  single_tree_dt[, tree_index := lgb_tree$tree_index]
James Lamb's avatar
James Lamb committed
173

174
  # Return tree
175
  return(single_tree_dt)
James Lamb's avatar
James Lamb committed
176

177
}