Utilizing Machine Learning Models to Predict Acute Kidney Injury in Septic Patients from MIMIC-III Database

Data Source

The MIMIC-III database, sourced from the ICU at Beth Israel Deaconess Medical Center in Boston, contains extensive health data on over 40,000 patients from 2001 to 2012 [12]. This robust data set features detailed records from more than 58,000 hospitalizations, including physiological signals, medications, lab results, diagnostic data, and medical summaries. Focuses on a diverse range of patients in the ICU, providing vital information for clinical decision-making and patient outcomes. The database was crucial for our study as it provided real patient data used to develop our models. The MIMIC-III dataset used in this study is publicly available at https://physionet.org/content/mimiciii/1.4/.

Participant Selection

Our focus was on patients who currently have sepsis. The flow chart in Figure 1 gives a breakdown of the screening and classification requirements that are carried out to reach the final data set. Initially, we filtered the MIMIC-III database using the unique ICD-9 codes 99591, 99592, and 78552 to find patients with sepsis. This left us with 6138 septic patients. We then excluded patients not within the age range of 18 to 89 years, patients with multiple admissions, patients who stayed in the ICU for less than 48 hours, and patients who had more than 20% missing values. Lastly, patients with incomplete information were also eliminated. Initially, we encountered 0.7% of our final dataset with more than 11 missing values. Given that this was less than 1% of the data, we removed these rows, assuming that it would not affect our result. We finalized the data set for 3301 septic patients, 2410 of whom developed AKI and 891 of whom did not.

• Download figure
• Open in new tab

Figure 1.

Flow chart of patient selection from MIMIC-III dataset.

Data Extraction

Patient data from the first 24 hours after admission were extracted from the MIMIC-III database to capture the early clinical status. Demographic information provided context on patient variability and potential risk factors, while comorbidities highlighted underlying conditions that could influence disease progression. Vital signs reflected physiological stability, offering key indicators of cardiovascular, respiratory, and metabolic function. Laboratory values provided biochemical information on organ performance, electrolyte balance, and coagulation status. In addition, therapeutic interventions such as mechanical ventilation and vasopressor use served as proxy for the severity of the disease. For continuously measured variables, maximum and minimum values were recorded to account for physiological fluctuations. Detailed indicators are shown in the following table.

To reduce bias caused by missing data, features with more than 20 missing values % were excluded from the final cohort. The missing values in other features were imputed using the multiple imputation (MI) method. Multiple Imputation (MI) is a statistical method widely used for handling missing data. Address missing values by generating multiple imputed datasets, each with the missing values being replaced with plausible values. These data sets are then synthetically analyzed, obtaining more accurate estimates of the missing values.

Feature Selection

After an initial filtering of septic patients and grouping them into AKI and non-AKI, we started with 50 characteristics. To determine the main characteristics, we used a correlation-based method. Each characteristic X_i was evaluated for its correlation with the target variable Y (AKI presence: 1 for AKI, 0 for non-AKI) using the Pearson correlation coefficient:

We selected features with an absolute correlation coefficient within the range: leaving us with the 23 most correlated features to AKI. All features satisfying were considered for further analysis.Table 1 shows the 23 features selected, which are (i) age: patient age when enter ICU; (ii) minimum heart rate: patient’s minimum recorded heart rate; (iii) minimum temperature: patient’s minimum recorded temperature; (iv) minimum oxygen saturation (SpO2): patient’s minimum recorded oxygen saturation; (v) minimum systolic blood pressure (SysBP): patient’s minimum recorded systolic blood pressure; (vi) minimum diastolic blood pressure (DiasBP): patient’s minimum recorded diastolic blood pressure; (vii) minimum bilirubin: patient’s minimum recorded bilirubin from lab results; (viii) maximum bilirubin: patient’s maximum recorded bilirubin from lab results; (ix) minimum aniongap: patient’s minimum recorded aniongap from lab results; (x) maximum aniongap: patient’s maximum recorded aniongap from lab results; (xi) minimum potassium: patient’s minimum recorded potassium level from lab results; (xii) maximum potassium: patient’s maximum recorded potassium level from lab results; (xiii) minimum lactate: patient’s minimum recorded lactate level from lab results; (xiv) maximum lactate: patient’s maximum recorded lactate level from lab results; (xv) minimum creatinine: patient’s minimum recorded creatinine level from lab results; (xvi) maximum creatinine: patient’s maximum recorded creatinine level from lab results; (xvii) minimum blood urea nitrogen (BUN): patient’s minimum recorded blood urea nitrogen from lab results;(Xviii) maximum blood urea nitrogen (BUN): patient’s maximum recorded blood urea nitrogen from lab results; (xix) urine output: patient’s recorded urine output; (xx) vasopressor (vaso): whether patient took vasopressor or not; (xxi) weight: patient’s weight at ICU admission; (xxii) minimum estimated glomerular filtration rate (eGFR): patient’s minimum recorded estimated glomerular filtration rate from lab results; and (xxiii) mechanical ventilation: whether a patient uses mechanical ventilation or not.

Our main characteristics were urine output, minimum eGFR, maximum eGFR, minimum SysBP, mechanical ventilation, vasopressor, weight, and minimum creatinine. The existing literature includes urine output, mechanical ventilation, vasopressor, and creatinine as the main characteristics, but the high correlation of minimum eGFR, maximum eGFR, minimum SysBP, and weight to predict AKI in septic patients is unique to our model.

Urine output is one of the most common signs of AKI. Studies have been done to show the ideal amount and time to collect urine sample measurements [13]. A decrease in urine content is typically a strong sign of AKI, making its importance as a characteristic a high factor in prediction. Regarding eGFR levels, it has been seen to decrease after patients are diagnosed with AKI [14]. Although there is still ongoing research on the effect of eGFR and AKI, it is a notable feature that should be considered for prediction. Systolic blood pressure is another indicator that should be considered when determining an AKI diagnosis because, based on previous studies, it is vital for organ perfusion, where blood flow provides oxygen and nutrients to the organs [15, 16]. Vasopressor is a drug that is used to increase blood pressure. It can be used to regulate hypotension by constricting blood vessels [17]. In our dataset, it was recorded whether a patient took vasopressor or not. Weight is another significant factor in the prediction of AKI. Studies have stated that obesity is associated with an increased risk of AKI in areas such as acute respiratory distress syndrome or after cardiac surgery [18, 19]. In a recent study involving critically ill septic patients in the ICU, obesity patients were found to have a higher chance of developing early sepsis-associated AKI [20]. The literature supports our findings on the characteristics that will be used in the development of a model to predict AKI in septic patients.

Statistical Analysis for AKI and Non-AKI Patients

To compare the characteristics of patients with AKI and non-AKI, we used statistical tests with an alpha set to 0.05 to compare the two cohorts and see if there are statistically significant differences. The results can be found in Table ??. Quantitative values were compared using a T test and qualitative values were compared using a chi-squared test. All features were statistically significant between AKI patients and non-AKI patients.

Model Development

Our filtered data set was imbalanced between AKI patients and non-AKI patients. There were 2 410 AKI patients and 891 non-AKI patients. To fix this, we used the synthetic minority over-sampling technique (SMOTE), which creates additional data from the minority class using the K Nearest Neighbors approach [21]. SMOTE was added to all the features that have been selected to balance the data and prevent overfitting. The data was then divided into a training, testing, and validation cohort. The six models used in this study were Logistic Regression (LR), XGBoost, K Nearest Neighbors (KNN), Support Vector Machines (SVM), Random Forest (RF), and LightGBM.

Logistic regression has been studied to have comparable performance results when compared to machine learning models and has been used in studies for predicting AKI [22, 23]. For its parameters, we used a C of 100, a maximum iteration of 200, and a penalty of ‘l2’. XGBoost has been used in previous studies for predicting AKI in patients with specific prior health conditions [24, 25]. For the parameters, we used the ‘binary:logistic’ objective, a ‘reg_lambda’ of 100, a ‘reg_alpha’ of 120, and a ‘max_depth’ of 2. KNN is another machine learning algorithm that utilizes classification and has been used for medical predictions, such as predicting development of AKI for heart failure patients [26, 27]. For its parameters, we used an n_neighbors of 40. SVM is another well-known machine learning model known to provide high accuracy [28]. For SVM, we used a C of 0.1 and a gamma of 0.02. Random forest is a widely utilized classification machine learning model that has been used in several acute kidney injury prediction studies [29, 30]. We used an n_estimator of 150, a max_depth of 12, a min_samples_split of 128, and a min_samples_leaf of 10. LightGBM is another commonly used machine learning model in the literature for predicting AKI in septic patients. We used the following hyperparameters: λ_l1 = 1, λ_l2 = 1, num_leaves = 20, max_depth = 4, random_state = 42, learning_rate = 0.01, n_estimators = 1050, and verbose =− 1.

The best model was determined on the basis of its respective AUC value. Although we also looked at the accuracy, F1 score, and recall of each model, we decided to use AUC as the metric of comparison since it is more suitable for handling unbalanced datasets and when comparing multiple models.

For this study, the data was imported to BigQuery and extracted using SQL. Data preprocessing, feature selection, and training were performed using Python 3.10.12.

Shapley Analysis

To interpret the predictive power of each characteristic within our model, we employed Shapley analysis. Shapley values assign a contribution score to each feature based on its marginal impact on the model’s output, taking into account all possible feature combinations. This approach provides a comprehensive measure of the importance of features that is both rigorous and interpretable, making it especially suited to healthcare applications where understanding the influence of each variable is crucial. Shapley analysis helps us identify the most critical predictors, such as urine output, bilirubin, and blood urea nitrogen levels. These values are particularly useful for determining how each characteristic individually, as well as in combination with others, influences the likelihood of AKI, thus allowing for more informed clinical insights.

The analysis was chosen because of its ability to illustrate the contributions of the characteristics in a transparent way, which is consistent with the need for interpretability in clinical decision making tools [31]. This analysis enabled us to pinpoint the most impactful characteristics, visualized in our results section, that support the application of our model in real world settings where clinicians can focus on key indicators of risk of AKI.

Training and Validation Comparison

After an initial filtering of patients who did not meet the analysis requirements, we were left with 3301 patients and 56 characteristics. After narrowing our characteristics down to 23 using correlation, the data was then randomly divided into 40% for the training set, 10% for the test set, and 50% for the validation set, resulting in 1980 patients for training, 661 patients for testing, and 660 for validation. Table 2 gives a comprehensive comparison between the training and validation sets. The comparison of the training and validation set is to confirm that our null hypothesis of no difference between feature distributions holds.

Model Performance Evaluation

After running each of the six models, we have documented our results in Tables 3 and 4. The Logistic Regression model had a training set AUC of 0.898 (95% CI = [0.888 - 0.909]), a validation set AUC of 0.857 (95% CI = [0.829 – 0.885]), and a testing set AUC of 0.887 (95% CI = [0.861 - 0.915]). Regarding the other five models, LightGBM 0.885, SVM had an AUC of 0.876, random forest 0.867, XGBoost 0.862, and KNN 0.725.

Figure 2 shows the AUC curves of each of the 6 models tested. Logistic regression exhibits the closest shape to an ideal AUC curve, showcasing its superior performance amongst the other models.

• Download figure
• Open in new tab

Figure 2.

ROC curves of Baseline and Proposed models.

In addition, we applied rigorous calibration techniques to ensure the reliability of the predicted probabilities of each model. Model calibration was assessed using the Brier score, which quantifies the mean squared difference between predicted probabilities and actual outcomes, a lower score indicating better calibration. Among the models tested, logistic regression achieved the lowest Brier score (0.134), followed by LightGBM (0.144), SVM (0.141), Random Forest (0.143), XGBoost (0.147) and KNN (0.229), demonstrating its superior calibration performance. To further refine the models, we applied isotonic regression, a non-parametric technique, to adjust the probability estimates. As shown in Figure 3, the predicted probabilities align closely with the diagonal line, indicating an effective calibration. The combination of a low Brier score and a high AUC underscores the accuracy and reliability of logistic regression.

• Download figure
• Open in new tab

Figure 3.

Calibration curves of the 6 models used: Logistic Regression, SVM, XGBoost, Random Forest, KNN, and LightGBM.

Shapley Analysis for Logistic Regression

Shapley analysis is a common measure to see which features contribute the most to the model’s overall performance. Since Logistic Regression was our best model, we applied Shapley analysis to see which features most affected the outcome of our testing set.

Figure 4 visualizes the Shapley analysis for our Logistic Regression model. Urine output, maximum bilirubin, minimum bilirubin, and weight are the most important features used in the Logistic Regression prediction model. From Figure 5, we see that urine output has the most significant impact on the Logistic Regression model. Low urine output increases the risk of AKI, displayed by its negative SHAP value. This is consistent with clinical understanding where reduced urine output is a key indicator of acute kidney injury (AKI) [32]. Elevated bilirubin levels push predictions toward higher risk, reflecting liver dysfunction or sepsis-induced cholestasis, conditions often seen in critical care patients [33].

• Download figure
• Open in new tab

Figure 4.

Shapley Analysis of Logistic Model features.

• Download figure
• Open in new tab

Figure 5.

Feature importance for logistic regression test set.

Compared to the importance of the original feature, it is interesting to note that bilirubin now plays an important role in the predictive process. While in the original feature correlation results it was not among the top features, it now holds a higher predictive power in the Logistic Regression model. This highlights the importance of performing additional analysis to understand what specific features the model looks at to make predictions.