First we attach the healthcareai R package to make its functions available. If your package version is less than 2.0, none of the code here will work. You can check the package version with packageVersion("healthcareai"), and you can get the latest stable version by running install.packages("healthcareai"). If you have v1.X code that you want to use with the new version of the package, check out the Transitioning vignette.

healthcareai comes with a built in dataset documenting diabetes among adult Pima females. Once you attach the package, the dataset is available in the variable pima_diabetes. Let’s take a look at the data with the str function. There are 768 records in 10 variables including one identifier column, several nominal variables, and substantial missingness (represented in R by NA).

Easy Machine Learning

If you don’t want to fuss with details any more than necessary, machine_learn is the function for you. It makes it as easy as possible to implement machine learning models by putting all the detains in the background so that you don’t have to worry about them. Of course it might be wise to worry about them, and we’ll get to how to do that further down, but for now, you can automatically take care of problems in the data, do basic feature engineering, and tune multiple machine learning models using cross validation with machine_learn.

machine_learn always gets the name of the data frame, then any columns that should not be used by the model (uninformative columns, such as IDs), then the variable to be predicted with outcome =. If you want machine_learn to run faster, you can have that—at the expense of a bit of predictive power—by setting its tune argument to FALSE.

machine_learn has told us that it has created a recipe for data preparation (this allows us to do exactly the same data cleaning and feature engineering when you want predictions on a new dataset), is ignoring patient_id when tuning models as we told it to, is training classification algorithms because the outcome variable diabetes is categorical, and has executed cross validation for three machine learning models: random forests, XGBoost, and regularized regression. Let’s see what the models look like.

Everything looks as expected, and the best model is is a random forest that achieves performance of AUROC = 0.84. Not bad for one line of code.

Now that we have our models, we can make predictions using the predict function. If you provide a new data frame to predict it will make predictions on the new data; otherwise, it will make predictions on the training data.

We get a message about when the model was trained and how well it preformed in training, and we get back a data frame that looks sort of like the original, but has a new column predited_diabetes that contains the model-generated probability each individual has diabetes, and contains changes that were made preparing the data for model training, e.g. missingness has been filled in and weight_class has been split into a series of “dummy” variables.

We can plot how effectively the model is able to separate diabetic from non-diabetic individuals by calling the plot function on the output of predict.

If you want outcome-class predictions in addition to predicted probabilites, the outcome_groups argument accomplishes that. If it is TRUE the overall accuracy of predictions is maximized. If it is a number, it represents the relative cost of a false-negative to a false-positive outcome. The example below says that one false negative is as bad as two false positives. If you want risk groups instead, see the risk_groups argument.

Data Profiling

It is always a good idea to be aware of where there are missing values in data. The missingness function helps with that. In addition to looking for values R sees as missing, it looks for other values that might represent missing, such as "NULL", and issues a warning if it finds any. Like many healthcareai functions, it has a plot method so you can inspect the results more quickly and intuitively by passing the output to plot.

missingness(pima_diabetes) %>% 

It’s good that we don’t have any missingness in our ID or outcome columns. We’ll see how missingness in predictors is addressed further down.

Data Preparation

To get an honest picture of how well a model performs (and an accurate estimate of how well it will perform on yet-unseen data), it is wise to hide a small portion of observations from model training and assess model performance on this “validation” or “test” dataset. In fact, healthcareai does this automatically and repeatedly under the hood, so it’s not strictly necessary, but it’s still a good idea. The split_train_test function simplifies this, and it ensures the test dataset has proportionally similar characteristics to the training dataset. By default, 80% of observations are used for training; that proportion can be adjusted with the p parameter. The seed parameter controls randomness so that you can get the same split every time you run the code if you want strict reproducability.

split_data contains two data frames, named train and test.

One of the major workhorse functions in healthcareai is prep_data. It is called under-the-hood by machine_learn, so you don’t have to worry about these details if you don’t want to, but eventually you’ll want to customize how your data is prepared; this is where you do that. The helpfile ?prep_data describes what the function does and how it can be customized. Here, let’s customize preparation to scale and center numeric variables and avoid collapsing rare factor levels into “other”.

The first arguments to prep_data are the same as those to machine_learn: data frame, ignored columns, and the outcome column. Then we can specify prep details.

The “recipe” that the above message refers to is a set of instructions for how to transform a dataset the way we just transformed our training data. Any machine learning that we do (within healthcareai) on prepped_training_data will retain that recipe and apply it before making predictions on new data. That means that when you have models making predictions in production, you don’t have to figure out how to transform the data or worry about encountering missing data or new category levels.

Model Training

machine_learn takes care of data preparation and model training for you, but if you want more precise control, tune_models and flash_models are the model-training function you’re looking for. They differ in that tune_models searches over hyperparameters to optimize model performance, while flash_models trains models at set hyperparameter values. So, tune_models produces better models, but takes longer (approaching 10x longer at default settings).

Let’s tune all three available models: random forests (“RF”), regularized regression (i.e. lasso and ridge, “GLM”), and gradient-boosted decision trees (i.e. XGBoost, “XGB”). To optimize model performance, let’s crank tune_depth up a little from its default value of ten. That will tune the models over more combinations of hyperparameter values in the search for the best model. This will increasing training time, so be cautious with it at first, but for this modest-sized dataset, the entire process takes less than a minute to complete on a laptop.

Let’s also select “PR” as our model metric. That optimizes for area under the precision-recall curve rather than the default of area under the receiver operating characteristic curve (“ROC”). This is usually a good idea when one outcome category is much more common than the other category.

You can compare performance across models with evaluate.

For more detail, you can examine how models perform across hyperparameters by plotting the model object. Here we plot only the best model’s performance over hyperparameter by extracting it by name. It looks like extratrees is a superior split rule for this model.

Faster Model Training

If you’re feeling the need for speed, flash_models is the function for you. It uses fixed sets of hyperparameter values to train the models, so you still get a model customized to your data, but without burning the electricity and time to precisely optimize all the details. Here we’ll use models = "RF" to train only a random forest.

If you want to train a model on fixed hyperparameter values, but you want to choose those values, you can pass them to the hyperparameters argument of tune_models. Run get_hyperparameter_defaults() to see the default values and get a list you can customize.

Model Interpretation


If you trained a GLM model, you can extract model coefficients from it with the interpret function. These are coefficient estimates from a regularized logistic or linear regression model. If you didn’t scale your predictors (which is the default in prep_data), these will be in natural units (e.g. in the plot below, a unit increase in plasma glucose corresponds to an expected log-odds increase of diabetes of just over one). Importantly, natural units mean that you can’t interpret the size of the coefficients as the importance of the predictor. To get that interpretation, scale your features during data preparation by calling prep_data with scale = TRUE and then flash_models or tune_models.

In this plot, the low value of weight_class_normal signifies that people with normal weight are less likely to have diabetes. Similarly, plasma glucose is associated with increased risk of diabetes after accounting for other variables.

Variable Importance

Tree based methods such as random forest and boosted decision trees can’t provide coefficients like regularized regression models can, but they can provide information about how important each feature (aka predictor, aka variable) is for making accurate predictions. You can see these “variable importances” by calling get_variable_importance on your model object. Like interpret and many other functions in healthcareai, you can plot the output of get_variable_importance with a simple plot call.


The explore function reveals how a model makes its predictions. It takes the most important features in a model, and uses a variety of “counterfactual” observations across those features to see what predictions the model would make at various combinations of the features. To see the effect of more features adjust the n_use argument to plot, or for different features, specify x_var and color_var.


predict will automatically use the best-performing model from training (evaluated out-of-fold in cross validation). If no new data is passed to predict it will return out-of-fold predictions from training. The predicted probabilities appear in the predicted_diabetes column.

To get predictions on a new dataset, pass the new data to predict, and it will automatically be prepared based on the recipe generated on the training data. We can plot the predictions to see how well our model is doing, and we see that it’s separating diabetic from non-diabetic individuals pretty well, although there a fair number of non-diabetics with high predicted probabilities of diabetes. This may be due to optimizing for precision recall, or may indicate pre-diabetic patients.

Above, we saw how to make outcome-class predictions. Here, we make risk-group predictions, defining four risk groups (low, moderate, high, and extreme) containing 30%, 40%, 20% and 10% of patients, respectively.

Saving, Moving, and Loading Models

Everything we have done above happens “in memory”. It’s all within one R session, so there’s no need to save anything to disk or load anything back into R. Putting a machine learning model in production typically means moving the model into a production environment. To do that, save the model with save_models function.

save_models(models, file = "my_models.RDS")

The above code will store the models object with all its metadata in the my_models.RDS file in the working directory, which you can identify with getwd(). You can move that file to any other directory or machine, even across operating systems, and pull it back into R with the load_models function.

The only tricky thing here is you have to direct load_models to the directory that the model file is in. If you don’t provide a filepath, i.e. call load_models(), you’ll get a dialog box from which you can choose your model file. Otherwise, you can provide load_models an absolute path to the file, e.g. load_models("C:/Users/"), or a path relative to your working directory, which again you can find with getwd(), e.g. load_models("data/my_models.RDS"). If you put the models in the same directory as your R script or project, you can load the models without any file path.

models <- load_models("my_models.RDS")

That will reestablish the models object in your R session. You can confirm this by clicking on the “Environment” tab in R Studio or running ls() to list all objects in your R session.

A Regression Example

All the examples above have been classification tasks, predicting a yes/no outcome. Here’s an example of a full regression modeling pipeline on a silly problem: predicting individuals’ ages. The code is very similar to classification.

regression_models <- machine_learn(pima_diabetes, patient_id, outcome = age)
# > Training new data prep recipe...
# > Variable(s) ignored in prep_data won't be used to tune models: patient_id
# > 
# > age looks numeric, so training regression algorithms.
# > 
# > After data processing, models are being trained on 14 features with 768 observations.
# > Based on n_folds = 5 and hyperparameter settings, the following number of models will be trained: 50 rf's, 50 xgb's, and 100 glm's
# > Training with cross validation: Random Forest
# > Training with cross validation: eXtreme Gradient Boosting
# > Training with cross validation: glmnet
# > 
# > *** Models successfully trained. The model object contains the training data minus ignored ID columns. ***
# > *** If there was PHI in training data, normal PHI protocols apply to the model object. ***
# > Models trained: 2018-09-01 18:27:53
# > 
# > Models tuned via 5-fold cross validation over 10 combinations of hyperparameter values.
# > Best performance: RMSE = 8.9, MAE = 6.3, Rsquared = 0.43
# > By Random Forest with hyperparameters:
# >   mtry = 13
# >   splitrule = extratrees
# >   min.node.size = 12
# > 
# > Out-of-fold performance of all trained models:
# > 
# > $`Random Forest`
# > # A tibble: 10 x 9
# >    mtry splitrule min.node.size  RMSE Rsquared   MAE RMSESD RsquaredSD
# > * <int> <chr>             <int> <dbl>    <dbl> <dbl>  <dbl>      <dbl>
# > 1    13 extratre…            12  8.91    0.429  6.29  0.952     0.0698
# > 2     9 extratre…            12  8.92    0.430  6.31  0.939     0.0672
# > 3     4 variance             11  8.93    0.431  6.40  0.952     0.0678
# > 4     9 variance             20  9.00    0.418  6.32  1.00      0.0740
# > 5     3 extratre…             7  9.31    0.422  6.88  0.851     0.0539
# > # ... with 5 more rows, and 1 more variable: MAESD <dbl>
# > 
# > $`eXtreme Gradient Boosting`
# > # A tibble: 10 x 13
# >      eta max_depth gamma colsample_bytree min_child_weight subsample
# > *  <dbl>     <int> <dbl>            <dbl>            <dbl>     <dbl>
# > 1 0.0136         2 8.68             0.597             5.49     0.889
# > 2 0.0405         6 3.95             0.521             1.73     0.747
# > 3 0.307          1 8.40             0.531             3.51     0.963
# > 4 0.136          6 0.209            0.701             1.76     0.905
# > 5 0.197          7 3.38             0.806             1.85     0.646
# > # ... with 5 more rows, and 7 more variables: nrounds <int>, RMSE <dbl>,
# > #   Rsquared <dbl>, MAE <dbl>, RMSESD <dbl>, RsquaredSD <dbl>, MAESD <dbl>
# > 
# > $glmnet
# > # A tibble: 20 x 8
# >   alpha lambda  RMSE Rsquared   MAE RMSESD RsquaredSD MAESD
# > * <dbl>  <dbl> <dbl>    <dbl> <dbl>  <dbl>      <dbl> <dbl>
# > 1     1 0.0318  9.40    0.363  6.76  0.743     0.0520 0.453
# > 2     1 0.0303  9.40    0.363  6.76  0.743     0.0520 0.453
# > 3     1 0.0288  9.40    0.363  6.76  0.743     0.0520 0.453
# > 4     1 0.0795  9.40    0.363  6.75  0.746     0.0532 0.442
# > 5     1 0.0208  9.40    0.363  6.76  0.742     0.0517 0.455
# > # ... with 15 more rows

Let’s make a prediction on a hypothetical new patient. Note that the model handles missingness in insulin and a new category level in weight_class without a problem (but warns about it).