Learn R Programming

dtGAP (version 0.0.2)

train_tree: Fit a Decision Tree Model

Description

Fits a decision tree to training data using one of several supported tree implementations (rpart, party, C50, or via caret) and returns a standardized party object along with variable importance scores.

Usage

train_tree(
  data_train = NULL,
  data = NULL,
  target_lab = NULL,
  model = c("rpart", "party", "C50", "caret"),
  task = c("classification", "regression"),
  control = NULL
)

Value

A list with elements:

fit

A party object representing the fitted tree.

var_imp

A named numeric vector of relative variable importance (scaled to sum to 1 and rounded to two decimals).

Arguments

data_train

Data frame. Explicit training set. If NULL, will be subset from data by data_type == 'train'.

data

Data frame. Combined dataset with a data_type column when data_train is NULL.

target_lab

Character. Name of the target column to predict.

model

Character. Which implementation to use: one of "rpart", "party", "C50", or "caret".

task

Character. Type of task: "classification" or "regression".

control

List or control object. Optional control parameters passed to the chosen tree function.

Examples

Run this code

library(partykit)
library(C50)
library(caret)

data(train_covid)
train_tree(data_train = train_covid, target_lab = "Outcome", model = "rpart")
train_tree(data_train = train_covid, target_lab = "Outcome", model = "C50")
train_tree(data_train = train_covid, target_lab = "Outcome", model = "caret")

data(Psychosis_Disorder)
data <- add_data_type(data_all = Psychosis_Disorder)
data <- prepare_features(data, target_lab = "UNIQID", task = "classification")
train_tree(
  data = data, target_lab = "UNIQID", model = "party",
  control = ctree_control(minbucket = 15)
)

Run the code above in your browser using DataLab