COVID-19 mortality prediction using ensemble learning and grey wolf optimization

Dataset

The dataset comes from a study by Altschul et al. (2020). The dataset contains 4,711 COVID-19-infected cases confirmed by a real-time analyzer (characterized by reverse transcriptase-polymerase chain reaction) in the period between March 1 to April 16, 2020. Cases evaluated in the emergency zone but not admitted or died at the site were removed from the study. The majority of patients had a single hospitalization, and the research only evaluated the most recent cases for those with repeated hospitalizations over the study period (Altschul et al., 2020).

The dataset has 85 variables, most of which are secondary variables created from the original ones. In our study, we used original variables only, so all secondary variables were removed. The “Derivation cohort” variable is a categorical variable used for grouping. Samples having “Derivation cohort = 1” were used as training samples, while samples having “Derivation cohort = 0” were used as validation (25%) samples and test samples (75%). The “Death” variable is the target class with two values: 0 (dead) and 1 (alive). Variables whose majority of values were missing were also removed. After filtering the whole set of variables, 20 original variables (features) were obtained for model development. These original variables were “Age.1” (patient’s age), “Temp” (Temperature), “OsSats” (oxygen saturation), “Lympho” (lyphocyte count), “WBC” (white blood cell count), “Plts” (platelet count), “Creatinine” (creatinine level), “MAP” (mean arterial pressure, in mmHg), “Sodium” (sodium level), “ALT” (alanine transaminase level), “AST” (aspartate transaminase level), “INR” (international normalised ratio), “BUN” (blood urea nitrogen), “Troponin” (troponin level), “CrctProtein”, “Ddimer” (D-dimer level), “Glucose” (Blood glucose level), “Ferritin” (ferritin level), “Procalcitonin” (procalcitonin level), “IL6” (interleukin-6 level). Table 1 shows descriptive statistics for 20 selected features (variables) in the two cohorts. Figure 1 visualizes the distributions of those features in the training set (Cohorts 1).

Overview of the method

Figure 2 describes the main steps in developing our ensemble model. Initially, samples were grouped into two cohorts: Cohort 0 and Cohort 1. The samples in Cohort 1 were used as training data. For Cohort 0, the validation and test sets were created using random sampling with the ratio of 25% and 75% of total samples, respectively. It is worth noting that the validation set was used to find the optimal weights for selected base classifiers only. A $5$-fold cross-validation was executed on the training set to discover the optimal hyper-parameters for each model. Five learning algorithms, including gradient boosting (GB), random forest (RF), extremely randomized trees (ERT), $k$-nearest neighbors ( $k$-NN), and support vector machines (SVM), were used to construct prediction models and these models were ranked based on their performances. Three best-performing models were selected as the base classifiers for the ensemble learning strategy. The top three classifiers were then retrained on the whole training set using their optimal hyper-parameters. Our proposed strategy is weighted ensemble learning which means each base model is allocated a weight values ( $w_1$, $w_2$, and $w_3$) explored by using the grey wolf optimization (GWO) algorithm (Rezaei, Bozorg-Haddad & Chu, 2018) on the validation set. Eventually, the test set was used to evaluate the model’s performance. The ensemble model predicts a probability for each sample in the form of $p_{ensemble}$ = $w_1$. $p_1$ + $w_2$. $p_2$ + $w_3$. $p_3$.

Machine learning algorithms

Grey wolf optimization

The grey wolf optimization (GWO) (Rezaei, Bozorg-Haddad & Chu, 2018) technique was inspired by the cooperative hunting behavior of grey wolves. It was first presented by Mirjalili, Mirjalili & Lewis (2014) as an efficient optimization technique among other novel techniques of meta-heuristic optimization family. Later in 2015, Gholizadeh (Gholizadeh, 2015) continued to develop the GWO method to address the nonlinear behavior of a double-layer grid optimization issue. His research revealed that GWO accomplished a more satisfactory job than other algorithms at discovering the most promising strategy for nonlinear double-layer grids.

GWO quantitatively simulates the social order of wolves in their hunting activities with three wolves: Alpha ( $\alpha$), Beta ( $\beta$), and Delta ( $\delta$) whose leadership are decending. These $\alpha$, $\beta$, and $\delta$ wolves represent the best solution, the second-best and third-best solutions, respectively. The remaining wolves in the herd (other potential solutions) are called Omega ( $\omega$). In a hunting activity (optimization), $\alpha$, $\beta$, and $\delta$ wolves manage the hunting while $\omega$ wolves follow their leaders. The GWO is characterized by three stages: (i) Encircling, (ii) Hunting, and (iii) Attacking. In stage Encircling, the mathematical models for encircling behavior are expressed as:

(1) $$\overrightarrow{D}=|\overrightarrow{C}.{\overrightarrow{X}}_p(t)-\overrightarrow{X}(t)|,$$

(2) $$\overrightarrow{X}(t+1)={\overrightarrow{X}}_p(t)-\overrightarrow{A}.\overrightarrow{D},$$where $t$ is the present iteration, $\overrightarrow{X}$ and ${\overrightarrow{X}}_p$ refer to the location vectors of the grey wolf and the prey, respectively. $\overrightarrow{A}$ and $\overrightarrow{C}$ are coefficient vectors which are calculated as:

(3) $$\overrightarrow{A}=2.\overrightarrow{a}.r_1-\overrightarrow{a},$$

(4) $$\overrightarrow{C}=2.{\overrightarrow{r}}_2,$$where ${\overrightarrow{r}}_1$ and ${\overrightarrow{r}}_2$ are vectors that are randomly initialized in the interval of between 0 and 1, and vector $\overrightarrow{a}$ is set to be decreased linearly from 2 to 0. A grey wolf at location (X,Y) may move to the location of its prey ( $X^{\ast }$, $Y^{\ast }$) using the preceding equations. In stage Hunting, after the grey wolves recognize their prey’s location and encircle them, the $\alpha$ wolf drives the hunting activity. The $\beta$ and $\delta$ wolves sometimes participate in the activity. The hunting behavior of the grey wolves is formulated as:

(5) $$\overrightarrow{D_{\alpha }}=|\overrightarrow{C_1}.{\overrightarrow{X}}_{\alpha }-\overrightarrow{X}|,$$

(6) $$\overrightarrow{D_{\beta }}=|\overrightarrow{C_2}.{\overrightarrow{X}}_{\beta }-\overrightarrow{X}|,$$

(7) $$\overrightarrow{D_{\delta }}=|\overrightarrow{C_3}.{\overrightarrow{X}}_{\delta }-\overrightarrow{X}|,$$

(8) $$\overrightarrow{X_1}=\overrightarrow{X_{\alpha }}-\overrightarrow{A_1}.{\overrightarrow{D}}_{\alpha },$$

(9) $$\overrightarrow{X_2}=\overrightarrow{X_{\beta }}-\overrightarrow{A_2}.{\overrightarrow{D}}_{\beta },$$

(10) $$\overrightarrow{X_3}=\overrightarrow{X_{\delta }}-\overrightarrow{A_3}.{\overrightarrow{D}}_{\delta },$$

(11) $$\overrightarrow{X}(t+1)=\frac{\overrightarrow{X_1}+\overrightarrow{X_2}+\overrightarrow{X_3}}{3}.$$

In an $n$-dimensional investigation space, the search agent uses these equations to adjust its location based on $\alpha$, $\beta$, and $\delta$. The investigation space parameters $\alpha$, $\beta$, and $\delta$ would constitute a circle, and the ultimate location is estimated inside that circle. By the other way of explanation, $\alpha$, $\beta$, and $\delta$ wolves estimate the prey’s location, whereas other wolves in the herd update their locations to approach their prey. In stage Attacking, after the prey (objective) stops moving, the grey wolves finish their hunting game by attacking the prey.

Ensemble learning strategy

To construct the ensemble model, the three best-performing classifiers (from $5$-fold cross-validation) were selected and termed as ‘base classifiers’. For each data point, the ensemble model’s predicted probability is derived by successively multiplying the probabilities predicted by the three base classifiers by their respective weights.

(12) $$p_{ensemble,i}=w_1.p_{1,i}+w_2.p_{2,i}+w_3.p_{3,i},$$where $p_{ensemble,i}$ is the obtained ensemble model’s probability; $w_1$, $w_2$, and $w_3$ are the weights of the three base classifiers; $p_{1,i}$, $p_{2,i}$, and $p_{3,i}$ are the predicted probabilities computed by the three base classifiers. The sum of the weights is equal to 1 and these weights are computed by the GWO algorithm to maximize the area under the receiver operating characteristic (ROC) curve (AUC) value on the validation set.

Five-fold cross-validation

To discover the optimal hyperparameters for individual model, we used a $5$-fold cross-validated randomized search across parameter settings. During the $5$-fold cross-validation, the first four folds were employed as the training fold, while the left fold was employed as the validation fold. Each fold was iteratively used as a test fold until no unused fold was left. Table 3 provides results of the hyperparameter tuning using $5$-fold cross-validation. The results indicate that the GB model achieves the highest cross-validated AUC value, followed by the RF and ERT classifiers. The $k$-NN and SVM classifiers have similar cross-validated performance. Hence, the GB, RF, and ERT classifiers were selected as base classifiers for the ensemble learning strategy. The GB, RF, and ERT classifiers were finally retrained on the training set using their tuned hyperparameters.

Ensemble model

The trained GB, RF, and ERT models were termed ‘model’ numbered 1, 2, and 3, respectively. The validation set was used to optimize the model’s weights ( $w_1$, $w_2$, and $w_3$) that contributed to the ensemble model’s prediction. The ensemble model’s prediction has the form of $p_{ensemble}$ = $w_1$. $p_1$ + $w_2$. $p_2$ + $w_3$. $p_3$ where $p_1$, $p_2$, and $p_3$ are the predicted probabilities of ‘model’ 1, 2, and 3 respectively. The ensemble model was created with three scenarios using the all-features set and two selected-features sets. An ensemble model constructed from a smaller feature set while still giving a similar performance is preferable. Table 4 gives information on base classifiers and their tuned weights in three ensemble learning scenarios.

Ensemble model with all features

All 20 features were used in the construction of three base classifiers and an ensemble model. Table 5 displays the performance of the three base classifiers and the ensemble model. The ensemble model obtains AUC values of 0.7801 and MCC values of 0.3898, which are greater than those of its base classifiers. The GB model achieves higher balanced precision and sensitivity than the others. The ERT model outperforms other classifiers in terms of accuracy and specificity. Based on the tuned values, the GB model contributes the most to the ensemble model’s predictive ability, followed by the ERT and RF classifiers (Table 4). Figure 3 visualizes the feature importance ranking of three basic classifiers across 20 features. Variable ‘MAP’ has the greatest effect on all classifiers. Moreover, the ‘OsSats’ variable is ranked among the top five key features in three classifiers. While the ‘Age.1’ variable is regarded as the third most significant feature by the GB and ERT classifiers, it has less of an effect on the RF classifiers. The GB and RF classifiers rank ‘Troponin’ among their top five key features, but the ERT model mostly disregards it. In addition, there are variables whose contributions may be ignored by the RF and ERT classifiers.

Table 5:

The performance on the test set of the models developed using 20 selected features.

Model	AUC	BA	SN	SP	PR	MCC
GB	0.7768	0.6828	0.4734	0.8921	0.5582	0.3885
RF	0.7598	0.6377	0.3519	0.9235	0.5697	0.3325
ERT	0.7695	0.6156	0.2785	0.9526	0.6286	0.3223
Ensemble	0.7801	0.6726	0.4304	0.9147	0.5923	0.3898

DOI: 10.7717/peerj-cs.1209/table-5

Ensemble model with 10 selected features

The ensemble model was developed by selecting the best five features of each base model (Fig. 3). Since the top-five features of the three base classifiers are somewhat different from each other, all features in the top five were preserved. Hence, 10 features were used to build three basic classifiers and an ensemble model. In terms of the AUC value, the ensemble model remains superior to its individual classifiers. The GB model exhibits a superior balanced accuracy of 0.6844, a sensitivity of 0.4759, and an MCC of 0.3920 (Table 6). Comparatively, the ERT model achieves greater specificity and accuracy compared to the other classifiers. Based on the tuned values, the GB model still contributes the most to the ensemble model’s predictive ability, followed by the ERT and RF classifiers (Table 4). However, the contribution of the ERT model decreases while the contribution of RF increases. Figure 4 visualizes the feature importance ranking of three basic classifiers across 10 selected features. The top-three features of the three base classifiers in the first and second scenarios are unchanged. The performance of ensemble models developed with 10 and 20 features has equivalent predictive efficiency.

Table 6:

The performance on the test set of the models developed using 10 selected features.

Model	AUC	BA	SN	SP	PR	MCC
GB	0.7783	0.6844	0.4759	0.8929	0.5612	0.3920
RF	0.7626	0.6489	0.3772	0.9206	0.5775	0.3513
ERT	0.7694	0.6135	0.2810	0.9461	0.6000	0.3090
Ensemble	0.7802	0.6725	0.4405	0.9045	0.5705	0.3804

DOI: 10.7717/peerj-cs.1209/table-6

Model’s robustness and stability

To assess the model’s robustness and stability, we repeated the experiments 10 times for each model. Table 8 provides information on the model performance of three models over ten random trials. Results show that the performances of the three models have small variations and high repeatability. The performances of the model developed with 10 features and the model developed with 20 features are equivalent while the performance of the model developed with five features is smaller. The repeated experimental outcomes confirm the model’s robustness and stability. Besides, the model developed with 10 features is the best-performing model with a small number of variables used but still has high predictive power. Figure 5 visualizes the ROC curves of base classifiers and ensemble model.

Table 8:

The performance on the test set of three models over ten random trials.

Trial	AUCs of models using
	20 Features	10 Features	5 Features
1	0.7801	0.7802	0.7782
2	0.7800	0.7799	0.7766
3	0.7800	0.7798	0.7766
4	0.7799	0.7798	0.7766
5	0.7800	0.7798	0.7766
6	0.7800	0.7799	0.7766
7	0.7800	0.7798	0.7766
8	0.7799	0.7800	0.7766
9	0.7800	0.7798	0.7767
10	0.7799	0.7798	0.7766
Mean	0.7800	0.7800	0.7768
SD	0.0001	0.0001	0.0005
95% CI	[0.7799–0.7800]	[0.7798–0.7800]	[0.7764–0.7771]

DOI: 10.7717/peerj-cs.1209/table-8