##### _Advanced Data Analysis from an Elementary Point of View_ #####
# R code for the chapter "Model Evaluation"
# Please do not re-distribute or use without attribution
# http://www.stat.cmu.edu/~cshalizi/ADAfaEPoV/
#### Re-usable functions only, NOT the examples/demos






## ----kfold-cv-for-linear-models-----------------------------------------------
# General function to do k-fold CV for a bunch of linear models
  # Inputs: dataframe to fit all models on,
    # list or vector of model formulae,
    # number of folds of cross-validation
  # Output: vector of cross-validated MSEs for the models
cv.lm <- function(data, formulae, nfolds=5) {
  # Strip data of NA rows
    # ATTN: Better to check whether NAs are in variables used by the models
  data <- na.omit(data)
  # Make sure the formulae have type "formula"
  formulae <- sapply(formulae, as.formula)
  n <- nrow(data)
  # Assign each data point to a fold, at random
    # see ?sample for the effect of sample(x) on a vector x
  fold.labels <- sample(rep(1:nfolds, length.out=n))
  mses <- matrix(NA, nrow=nfolds, ncol=length(formulae))
  colnames <- as.character(formulae)
  # EXERCISE: Replace the double for() loop below by defining a new
  # function and then calling outer()
  for (fold in 1:nfolds) {
    test.rows <- which(fold.labels == fold)
    train <- data[-test.rows,]
    test <- data[test.rows,]
    for (form in 1:length(formulae)) {
       # Fit the model on the training data
       current.model <- lm(formula=formulae[[form]], data=train)
       # Generate predictions on the testing data
       predictions <- predict(current.model, newdata=test)
       # Get the responses on the testing data, using the formula and eval()
       # a formula is, internally, a list with attributes.  The first element of
        # the list is always "~", and then the second element of the list is
        # the response term, including the transformation
       # eval() takes an expression and then evaluates in an environment, such
        # as a data frame
       test.responses <- eval(formulae[[form]][[2]], envir=test)
       # Calculate errors
       test.errors <- test.responses - predictions
       # Calculate the MSE on that fold
       mses[fold, form] <- mean(test.errors^2)
    }
  }
  return(colMeans(mses))
}