R/water_feature.R

In tohweizhong/auxml: Auxiliary Functions for Machine Learning

water_feature <- function(mod, target, test_case){
    
    sm <- summary(mod)
    betas <- sm$coefficients[,1]
    nm <- names(betas)
    
    # prediction
    pred <- predict(mod, test_case)
    
    # feature vector
    feat_vec <- test_case[-which(test_case %>% names == target)] %>% as.matrix
    
    betas2 <- betas[-1] # less intercept
    
    # feature contribution
    feat_cont <- betas2*feat_vec
    feat_cont <- c(betas[1], feat_cont, pred)
    names(feat_cont) <- c(nm, "Prediction")
    
    # waterfall chart on feature contribution
    plot_data <- data.frame(coef = names(feat_cont), feat_cont = feat_cont, row.names = NULL)
    plot_data$coef <- factor(plot_data$coef, levels = plot_data$coef)
    plot_data$id <- seq_along(plot_data$coef)
    plot_data$Impact <- ifelse(plot_data$feat_cont > 0, "+ve", "-ve")
    plot_data[plot_data$coef %in% c("(Intercept)", "Prediction"), "Impact"] <- "Initial/Net"
    plot_data$end <- cumsum(plot_data$feat_cont)
    plot_data$end <- c(head(plot_data$end, -1), 0)
    plot_data$start <- c(0, head(plot_data$end, -1))
    plot_data <- plot_data[, c(3, 1, 4, 6, 5, 2)]
    
    gg <- ggplot(plot_data, aes(coef, fill = Impact)) +
        geom_rect(aes(coef,
                      xmin = id - 0.45,
                      xmax = id + 0.45,
                      ymin = end,
                      ymax = start)) +
        theme_minimal() +
        scale_fill_manual(values = c("darkred", "darkgreen", "darkblue")) +
        theme(axis.text.x = element_text(angle = 90, hjust = 1)) +
        xlab("Features") +
        ggtitle("Feature Waterfall")
    
    if(sign(plot_data$end[1]) != sign(plot_data$start[nrow(plot_data)]))
        gg <- gg + geom_hline(yintercept = 0)
    
    gg
 
}