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.

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.

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.

We can examine how the model performs across hyperparameters by plotting the model object. It looks like extratrees is a superior split rule for this model, and larger values of minimum node size tend to do better.

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


Let’s see how the predictive performance of our untuned random forest compares to our extensively tuned models. The evaluate function gives you performance metrics for the best performing model by default, but if you set all_models = TRUE you can see how well all trained performed.

In this the untuned model comes in just 0 AUPR units ahead of the tuned models. In our experience, that’s on the small side of typical, with the difference in performance growing with the size of the dataset. A workflow that often works well is to do all of your development using flash_models, and as a final step before putting a model into production, retrain the model using tune_models.


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.


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.

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-07-11 22:12:24
# > 
# > Models tuned via 5-fold cross validation over 9 combinations of hyperparameter values.
# > Best performance: RMSE = 8.9, MAE = 6.4, Rsquared = 0.43
# > By Random Forest with hyperparameters:
# >   mtry = 7
# >   splitrule = extratrees
# >   min.node.size = 12
# > 
# > Out-of-fold performance of all trained models:
# > 
# > $`Random Forest`
# > # A tibble: 9 x 9
# >    mtry splitrule  min.node.size  RMSE Rsquared   MAE RMSESD RsquaredSD
# > * <int> <chr>              <int> <dbl>    <dbl> <dbl>  <dbl>      <dbl>
# > 1     7 extratrees            12  8.95    0.428  6.39  0.606     0.0460
# > 2     3 variance               3  9.04    0.418  6.54  0.617     0.0465
# > 3     2 extratrees             2 10.00    0.389  7.68  0.564     0.0365
# > 4     2 extratrees            10 10.0     0.385  7.73  0.554     0.0381
# > 5     1 variance              20 10.6     0.385  8.40  0.542     0.0404
# > # ... with 4 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.290          1  8.05            0.800            3.82      0.517
# > 2 0.0122         5  8.30            0.882            2.46      0.938
# > 3 0.0789         7  4.84            0.590           10.9       0.983
# > 4 0.0290         7  7.01            0.778            3.88      0.830
# > 5 0.245          8  9.34            0.523            0.495     0.965
# > # ... 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     0 0.00192  9.40    0.369  6.78  0.713     0.0754 0.429
# > 2     0 0.00624  9.40    0.369  6.78  0.713     0.0754 0.429
# > 3     0 0.00631  9.40    0.369  6.78  0.713     0.0754 0.429
# > 4     0 0.0212   9.40    0.369  6.78  0.713     0.0754 0.429
# > 5     0 0.0699   9.40    0.369  6.78  0.713     0.0754 0.429
# > # ... 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).