Network-based predictive models for artificial intelligence: an interpretable application of machine learning techniques in the assessment of depression in stroke patients

Abstract

Background

Depression is a common complication after a stroke that may lead to increased disability and decreased quality of life. The objective of this study was to develop and validate an interpretable predictive model to assess the risk of depression in stroke patients using machine learning (ML) methods.

Methods

This study included 1143 stroke patients from the NHANES database between 2005 and 2020. First, risk factors for depression in stroke patients were determined by univariate and multivariate logistic regression analysis. Next, five machine learning algorithms were used to construct predictive models, and several evaluation metrics (including area under the curve (AUC)) were used to compare the predictive performance of the models. In addition, the SHAP (Shapley Additive Explanations) method was used to rank the importance of features and to interpret the final model.

Results

We screened seven features to construct a predictive model. Among the 5 machine learning models, the XGBoost (extreme gradient boosting) model showed the best discriminative ability, with an AUC of the ROC (receiver operating characteristic curve) in the test set of 0.746 and an accuracy of 0.834. In addition, the prediction results of the XGBoost model were interpreted in detail using the SHAP algorithm. We also developed a web-based calculator that provides a convenient tool for predicting the risk of depression in stroke patients at the following link: https://prediction-model-for-depression.streamlit.app.

Conclusions

Our interpretable machine learning model serves as an auxiliary tool for clinical judgment, aimed at early and effective identification of depression risk in stroke patients.

Introduction

Depression is a prevalent and serious complication among stroke patients. Studies indicate that it affects approximately 30% of individuals within five years post-stroke, a rate significantly higher than the prevalence in the general population [1, 2]. Stroke not only impairs physical functioning but also poses a substantial challenge to the mental health of patients. Those affected often experience mood swings, loss of self-efficacy, and fear of the future, which collectively contribute to the onset of depression [2]. In stroke patients, depression manifests through a range of symptoms, including persistent sadness, loss of interest, increased fatigue, insomnia or excessive sleepiness, and difficulty concentrating [3, 4]. These symptoms adversely influence the emotional state of patients and may lead to cognitive decline, hindering their recovery process and overall quality of life [5]. Research has established a significant association between depression and poor functional recovery, an increased risk of recurrent stroke, and heightened mortality rates following a stroke [6, 7]. Consequently, the high prevalence of depression among stroke patients represents a public health issue that warrants urgent attention. Furthermore, individuals suffering from depression often exhibit lower adherence and motivation during rehabilitation therapy, undermining treatment effectiveness and exacerbating their suffering. Therefore, early identification and intervention for depressive symptoms in stroke patients is essential.

Nomograms have been extensively utilized in various studies to predict the risk of depression in stroke patients. These studies offer an intuitive risk assessment tool by integrating multiple clinical variables through statistical modeling [8, 9, 10]. Nomograms are user-friendly and provide straightforward risk evaluations. Nevertheless, to enhance the accuracy of depression risk prediction, a machine learning (ML) approach has been employed [11, 12, 13]. ML algorithms can manage numerous variables and identify potential non-linear relationships, thereby demonstrating notable advantages in complex data analysis [14]. Through self-learning, ML models can be continuously refined to improve prediction accuracy. This approach not only boosts the predictive power of the models but also reinforces their clinical applicability. An additional significant advantage is that many ML models can overcome the “black boxes” limitations associated with traditional models [15]. While some ML algorithms face challenges regarding interpretability, feature significance analysis and visualization tools such as SHAP (Shapley Additive Explanations) can elucidate the model’s decision-making process, thereby enhancing transparency [16, 17]. This interpretability allows clinicians to comprehend how models derive their predictions, fostering trust in real-world applications.

This study aimed to develop and validate an interpretable ML model for the early and accurate prediction of depression risk in stroke patients using the NHANES 2005–2020 dataset. We employed the SHAP method to clarify the significance of each feature and to elucidate the model’s decision-making process. Furthermore, we assessed the model’s significance in clinical prognosis to assist healthcare professionals in better identifying and managing the risk of depression in stroke patients, ultimately improving their overall health and quality of life.

Methods

Study design and study population

This study utilized data from the National Health and Nutrition Examination Survey (NHANES) 2005–2020, a comprehensive cross-sectional study of the non-institutionalized civilian population in the United States. NHANES assesses adults’ and children’s health and nutritional status through household interviews and physical examinations conducted at mobile screening centers. The household interview gathers demographic, socioeconomic, dietary, and health information, while the physical examination includes medical, dental, physiological, and laboratory evaluations. The National Center for Health Statistics (NCHS) Ethics Review Board approved the study, and informed consent was given by each participant. Out of 76,496 participants, 1805 were identified as stroke patients. After excluding those with missing data on depression and other covariates, a total of 1143 participants were included in the analyses (Fig. S1).

Identification of stroke

Stroke patients were identified based on self-reported diagnostic history, as determined by the question: “Has a doctor or other health professional ever told you that you had a stroke?” Participants who answered “yes” to this question were classified as stroke patients.

Determination of depression

The Patient Health Questionnaire-9 (PHQ-9) is a validated self-report instrument for assessing depressive symptoms over the previous two weeks [18]. The PHQ-9 consists of nine items, each rated on a scale of 0 to 3 (0 = ‘not at all’, 1 = ‘a few days’, 2 = ‘more than half the days’, 3 = ‘almost every day’). The total score ranges from 0 to 27. In this study, a PHQ-9 score of ≥ 10 was defined as indicating depression, while a score of < 10 was classified as no depression [19].

Predictors

Demographic information collected in this study included gender (male, female), age, race (Mexican American, non-Hispanic white, non-Hispanic black, Hispanic, other race), education level (less than high school, high school or equivalent, college or above), marital status (married/living with partner, widowed/divorced/separated, never married), and poverty income ratio (PIR). Physical examination data provided body mass index (BMI), calculated as weight (kg) divided by height (m) squared (kg/m²). Lifestyle factors investigated included smoking behavior (whether participants had smoked at least 100 cigarettes in their lifetime), drinking habits (whether they had consumed at least 12 alcoholic beverages of any type in any given year), and moderate recreational activities. Disease history was also recorded, including diagnoses of hypertension, diabetes, arthritis, congestive heart failure (CHF), coronary heart disease (CHD), heart attack, and cancer. Additionally, data related to sleep duration and sleep disorders were collected through questionnaires, either self-reported by patients or diagnosed by physicians. Important biochemical indicators were obtained from laboratory tests, including total cholesterol (TC), high-density lipoprotein cholesterol (HDL-C), and creatinine levels.

Data preprocessing

To address the issue of missing values in the dataset, we chose to exclude all individuals with missing data to ensure the accuracy of the analysis and the reliability of the model. While this approach helps minimize potential biases associated with imputing missing values, we acknowledge that excluding participants with missing data may introduce selection bias, especially when the missing data mechanism is not completely random. Additionally, to ensure the fairness and validity of model evaluation, we randomly split the dataset into training and test sets at a 7:3 ratio, with 70% of the data used for model training and 30% for assessing model performance. This procedure ensures the scientific rigor and credibility of the model we developed.

Construction of the model

In the training set, we first analyzed risk factors for depression in stroke patients using univariate logistic regression and selected variables significantly associated with depression (P < 0.05) for inclusion in the subsequent multivariate logistic regression model. In this analysis, variables with a P-value of less than 0.05 were considered candidates for predicting depression in stroke patients. We calculated odds ratios (OR) and 95% confidence intervals (CI). Furthermore, five different ML models were employed, including random forest (RF), decision tree (DT), extreme gradient boosting (XGBoost), Naïve Bayesian (NB), and support vector machine (SVM), to identify depression in stroke patients. In the hyperparameter settings (Table S1), we provide the corresponding parameter configurations for five models. Taking the XGBoost model as an example, we performed a grid search with 5-fold cross-validation to determine the optimal hyperparameters. Specifically, we explored a range of learning rates (0.001, 0.01, 0.1), maximum depths (3, 5, 7), and the number of estimators (100, 300, 500). Through systematic experiments and evaluation, we selected a learning rate of 0.01, a maximum depth of 3, and 300 estimators as the optimal combination. The choice of a learning rate of 0.01 effectively balances training speed and model performance while reducing the risk of overfitting. A maximum depth of 3 was selected to control model complexity, ensuring that the model captures important features while minimizing overfitting. The 5-fold cross-validation results demonstrated that this depth setting achieved stable performance on both training and validation datasets. Finally, after multiple experiments, we determined that using 300 trees provided the best trade-off between predictive accuracy and computational efficiency. These hyperparameter selections aim to enhance the model’s effectiveness and reproducibility, providing a reference for other model research.

Evaluation of the model

Evaluation metrics such as the receiver operating characteristic curve (ROC), area under the curve (AUC), average precision score (APS), accuracy, sensitivity (recall), specificity, negative predictive value (NPV), positive predictive value (PPV), false positive rate (FPR), false negative rate (FNR), and F1 score were used to evaluate the performance of the model. Using the SHAP algorithm, we also visualized the important features influencing the risk of depression in stroke patients, analyzed the importance of individual features on model output, and elucidated the impact of key features on the final model results.

Web deployment tool based on the streamlit framework

To facilitate the clinical application of the model, we developed a user-friendly web application based on the Streamlit Python framework. This application implements our final predictive model, allowing healthcare professionals to assess the risk of depression in stroke patients conveniently. When users input the feature values associated with the final model, the application automatically calculates and returns the probability of depression for that patient and a force map for individual stroke patients based on the pre-trained model.

Statistical analysis

All ML models were constructed and validated within the Python (3.12.0) environment. For statistical analyses, we utilized R software (4.3.2). For continuous data that were not normally distributed, we reported median and interquartile ranges for descriptive purposes and employed the Mann-Whitney U-test for between-group comparisons. Categorical data were described using frequencies and percentages, with the chi-square test applied to assess differences between groups.

Results

Baseline characteristics of participants

This study included a total of 1143 participants, with the training set comprising 806 individuals and the test set consisting of 337 individuals. The median age of the overall sample was 67 years, with 562 (49.17%) males and 581 (50.83%) females; among these, 200 were diagnosed with depression. In the training set, the median age also stood at 67 years, involving 399 (49.50%) males and 407 (50.50%) females, with 145 patients diagnosed with depression. In contrast, the median age of the test set was 66 years, comprising 163 (48.37%) males and 174 (51.63%) females, with 55 patients diagnosed with depression. There were no statistically significant differences (p > 0.05) between the two datasets across all variables (Table 1).

Table 1 Baseline characteristics of the training and test sets

Construction of ML models

In the univariate logistic regression model, a total of 14 variables were associated with the risk of depression in the training group of stroke patients. Following multivariate logistic regression analysis, seven significant predictors were ultimately identified: females had an increased risk of depression compared to men (OR: 1.814, 95% CI: 1.142–2.883); age was negatively associated with the risk of depression (OR: 0.977, 95% CI: 0.961–0.993); and the PIR was also negatively associated with the risk of depression (OR: 0.829, 95% CI: 0.701–0.993). Additionally, drinking increased the risk of depression (OR: 1.976, 95% CI: 1.143–3.415); sleep disorders significantly heightened the risk of depression (OR: 3.390, 95% CI: 2.179–5.272); moderate recreational activities were negatively correlated with the risk of depression (OR: 0.503, 95% CI: 0.299–0.846); and TC levels were positively associated with the risk of depression (OR: 1.215, 95% CI: 1.023–1.445) (Table 2). Subsequently, we applied five ML models to the NHANES dataset using the training dataset that included these seven variables.

Table 2 Univariate and multivariate logistic regression analyses of the training set

Testing the performance of ML models

During the testing phase, we applied the trained model to the test set. The results indicated that the XGBoost model outperformed others in terms of AUC performance (AUC: 0.746; 95% CI: 0.674–0.810). The AUCs for the other models were as follows: 0.711 (95% CI: 0.638–0.778) for RF, 0.719 (95% CI: 0.643–0.792) for DT, 0.671 (95% CI: 0.607–0.736) for NB, and 0.703 (95% CI: 0.627–0.776) for SVM. Figure 1 illustrates the accuracy correction and AUC curves for the five ML models. The accuracy rates for RF (0.825), DT (0.837), XGBoost (0.834), NB (0.837), and SVM (0.837) demonstrated good performance in identifying depression in stroke patients. Table 3 summarizes the estimation performance of each model. Among all five models, XGBoost exhibited the highest average precision score (APS) at 0.353, indicating the best discrimination. The specificity, sensitivity/recall, NPV, PPV, FPR, FNR, FDR, and F1 scores for the five models are presented in Table 3. Fig. S2 provides the confusion matrix for the five models. Comprehensive feature-based analysis confirmed that XGBoost demonstrated the highest precision and robustness in identifying depression in stroke patients.

Table 3 Comparison of the characteristics of five ML models

Fig. 1

Precision-recall curves and ROC for ML models. ROC, receiver operating characteristic curve; ML, machine leaning; RF, random forest; DT, decision tree; XGBoost, extreme gradient boosting; NB, Naïve Bayesian; SVM, support vector machine; APS, average precision score; AUC, the area under the curve; CI, confidence interval

Visualization of feature importance

Using the SHAP algorithm, we evaluated the importance of each feature for the XGBoost model in predicting the risk of depression among stroke patients. As shown in Fig. 2A, the feature importance plot ranks the most significant features associated with depression in descending order. The horizontal position of each feature in the plot indicates its positive or negative effect on the predicted value, where red reflects a high positive contribution, and blue indicates a low negative contribution (Fig. 2B). The analysis revealed that sleep disorders exhibited the strongest average predictive power among all features, followed by age, PIR, moderate recreational activities, TC, gender, and drinking. In Fig. 2C, the lines represent individual participants in the decision diagram, with characteristics ranked in descending order of importance based on the observed data. Additionally, Fig. 3A and C display individual force diagrams for depressed and non-depressed individuals, respectively. Figure 3B and D illustrate waterfall plots for depressed and non-depressed patients, respectively, to provide insights into the impact of individual characteristics on model predictions. The SHAP values illustrate the predicted characteristics of individual patients and their contributions to the risk of developing depression, where red features indicate an increased risk and blue features indicate a decreased risk. The length and direction of the arrows visualize the degree of influence of each predictive feature.

Fig. 2

Global model interpretation using the SHAP method. (A) SHAP summary bar plot. (B) SHAP summary dot plot. (C) SHAP decision plot. In the model, there is a dot for each patient’s SHAP value and therefore a dot for each feature for each patient. All variables are in descending order of importance. The color of the dots indicates the actual value of each patient feature, with red indicating a higher feature value and blue indicating a lower feature value

Fig. 3

Local model interpretation using the SHAP method. Fig. 3 A and 3 C show individual force diagrams for depressed and non-depressed patients, respectively. Figure 3B and D, on the other hand, show waterfall plots for depressed and non-depressed patients, respectively. Each patient is represented by the x-axis, while the contribution of features is represented by the y-axis: the larger the red part of each patient, the more likely it is to be judged as ‘depression’

Facilitating clinical applications

As illustrated in Fig. 4, the XGBoost-based prediction model has been integrated into a web application to facilitate its use in clinical practice. When the actual values of the seven features required by the model are entered, the application automatically predicts the risk of depression in stroke patients. Furthermore, the application displays a graph of the characteristics of a single stroke patient, highlighting the key factors influencing the depression prediction: the blue characteristics on the right contribute to a “non-depression” prediction, while the red characteristics on the left indicate a tendency towards “depression”. The web application is accessible online at the following link: https://prediction-model-for-depression.streamlit.app.

Fig. 4

Application of a web-based predictor on the risk of depression in stroke patients

The final XGBoost model developed in this study is based on seven features that can effectively predict the risk of depression in stroke patients. After inputting the actual values of these seven features, the application automatically calculates and displays the probability that the patient will develop depression. Meanwhile, the force diagram for a single stroke patient shows the features that help determine ‘depression’: the red features on the left are those that push the prediction into the ‘depressed’ category, while the blue features on the right are those that push the prediction into the ‘non-depressed’ category. XGBoost, extreme gradient boosting. The website that predicts the risk of depression in stroke patients is https://prediction-model-for-depression.streamlit.app

Discussion

In this study, we selected seven easily accessible clinical variables—gender, age, PIR, drinking, sleep disorders, moderate recreational activities, and TC—to construct a prediction model for the early identification of depression in stroke patients using ML algorithms. Utilizing the XGBoost algorithm, our results demonstrated stable and satisfactory performance, achieving an AUC value of 0.746. This indicated that the predictive model possessed good discriminatory ability, effectively distinguishing between high-risk and low-risk patients. Additionally, we employed the SHAP approach to quantify the importance of each selected feature for model predictions. Finally, we implemented the prediction model as a web application to facilitate its practical application in clinical scenarios.

Our findings revealed a higher prevalence of depression in women compared to men among stroke patients, corroborating previous studies that identified female gender as a significant risk factor for post-stroke depression [20]. Specifically, the annual diagnosis rate for depression was notably higher in women than in men after stroke (HR: 1.53, 95% CI: 1.51–1.55) [21]. Possible explanations for this disparity include the relative psychological vulnerability of women and their comparatively weaker coping mechanisms. The life and work implications of stroke, combined with familial and social pressures, may exacerbate negative emotions. Furthermore, women tend to have poorer prognoses post-stroke, which intensifies both physical and psychological distress along with financial burdens [22]. This study also found a negative correlation between age and depression, indicating that younger stroke patients are more likely to experience depression than older patients. This aligns with a review that noted a significant increase in depression prevalence among adolescents [23]. Young individuals often face greater family and social responsibilities, which can diminish psychological resilience and exacerbate stress [24]. Economic status has a significant impact on reducing the incidence of depression in stroke patients. Patients with better economic conditions generally have access to superior medical resources and social support. One study identified a correlation between the severity of post-stroke depression and patients’ economic status (χ² = 11.198, P = 0.024) [25]. While our study did not include educational level as a predictor, existing literature suggested that higher education correlated with a reduced risk of depression. For instance, a Chinese population-based study found that stroke patients with a high school education or above had a lower risk of depression compared to those with only primary education (OR: 0.50, 95% CI: 0.28–0.88, P = 0.016) [26]. Future interventions should thus consider enhancing both economic and educational resources to more effectively mitigate depression risk in stroke patients.

We observed a significant association between drinking and depression risk. Specifically, the prevalence of depression was markedly higher among alcohol drinkers compared to non-drinkers. Research-based on the Korean Community Health Survey (KCHS) showed that individuals consuming less than 5 g of alcohol daily had a 20% increased risk of depression (OR: 1.20, 95% CI: 1.07–1.35), while those consuming between 5 and 14.9 g per day had a 39% increased risk (OR: 1.39, 95% CI: 1.13–1.70) [27]. Conversely, a Mendelian randomization study indicated that alcohol consumption did not causally affect depression in older men, possibly due to the stress-relieving and mood-enhancing effects of low to moderate alcohol intake [28, 29]. These discrepancies might arise from differences in study design, sample selection, or genetic variation. Therefore, further research is needed to elucidate the complex relationship between alcohol consumption and depression in this demographic. In addition, stroke patients often contend with multiple comorbidities following acute treatment, making the relationship between sleep disorders and depression particularly important. Studies have indicated that the prevalence of post-stroke depression has been higher in patients with poor sleep quality than in those with good sleep quality [30]. Moderate to severe obstructive sleep apnea has been identified as a significant factor influencing post-stroke anxiety during the acute phase [31]. A retrospective study indicated that severe obstructive sleep apnea was significantly associated with an increased risk of post-stroke depression within three months (OR: 4.04, 95% CI: 1.38–9.62) [32]. Consequently, early identification and intervention for sleep disorders in stroke patients are crucial for reducing depression risk.

Our study found that moderate recreational activities significantly lowered the risk of depression among stroke patients. This finding aligns with existing literature, a meta-analysis of nine studies highlighted the potential benefits of home exercise in alleviating post-stroke depression, particularly emphasizing physical and mental exercises, such as tai chi, as effective treatments [2]. Clinical guidelines recommend non-pharmacological interventions, such as physical exercise, for stroke survivors experiencing mild depressive symptoms [33]. Although the precise mechanisms by which exercise alleviates depression are not fully understood, several plausible explanations exist. In the short term, physical activity activates the endorphin system, providing immediate mood enhancement; in the long term, regular exercise promotes neuroplasticity, improves brain function, and enhances the body’s stress resistance. Furthermore, exercise positively impacts depression by fostering social interactions and boosting self-esteem [34, 35]. In our analysis, TC emerged as a significant predictor of post-stroke depression. Consistent with existing studies, elevated TC concentrations were associated with a higher risk of depression [36]. Cholesterol plays a critical role in brain function as a vital component of nerve membranes, influencing neurotransmitter synthesis and release, which in turn affects mood and behavior [37]. However, a study conducted on a Japanese population found that elevated cholesterol levels during pregnancy were linked to a reduced risk of postpartum depression [38]. This discrepancy may stem from the unique psychological experiences associated with pregnancy, contrasting with the long-term rehabilitation challenges faced by stroke patients.

In recent years, relatively few studies have employed machine learning methods to predict depression risk in stroke patients. While one study constructed a machine learning model utilizing ten features and demonstrated high predictive performance [39], a significant limitation was the lack of application of the SHAP method to explain the model’s decision-making process. Machine learning models, particularly deep learning and ensemble methods, are often viewed as “black boxes”, rendering their internal mechanisms difficult to interpret. This opacity can lead to confusion among clinicians regarding how predictions are derived, which not only undermines clinical confidence but also hampers the implementation of personalized treatments. To address this challenge, our study employs the SHAP method to provide comprehensive explanations of model outputs. SHAP quantifies the contribution of each feature to the model’s predictions, allowing us to identify which factors are pivotal in predicting depression risk in stroke patients. The global interpretation reveals the primary features influencing depression risk, thus offering physicians valuable insights that facilitate the identification of high-risk patients and the development of tailored interventions. Concurrently, local explanations shed light on the predicted outcomes for individual patients, enabling clinicians to understand the specific sources of depression risk for each case—information that is essential for informed clinical decision-making. Additionally, to enhance the usability and convenience of our machine learning model, we developed a tool based on the Streamlit framework, making the predictive model easily accessible via a web interface. This user-friendly platform allows clinicians to conveniently input patient information and instantly retrieve prediction results along with SHAP interpretations. Through these enhancements, we not only improved model interpretability but also promoted knowledge-sharing and collaboration among clinicians. The shareability of this tool enables broader access to and utilization of the predictive model, thereby fostering improved early identification and intervention for depression risk in stroke patients. It is noteworthy that our study did not perform subgroup analyses for different patient populations (such as age, and gender) in terms of model predictions. Therefore, while the model demonstrated good predictive performance in the overall population, its applicability and effectiveness in specific subgroups have yet to be validated. This limitation restricts our comprehensive understanding of the model’s practical application and efficacy in various clinical contexts.

There are some limitations of this study that need to be noted. First, the cross-sectional design only identified significant associations between variables, precluding the establishment of causality. Future longitudinal studies or randomized controlled trials could better verify causal relationships. Second, due to the reliance on respondents’ subjective judgment and memory in self-reported data, the information reported may be inaccurate or incomplete. This bias can affect the reliability of the study’s findings, particularly in the diagnosis of stroke. Therefore, we recommend that future research employ more objective measurement tools to reduce the bias introduced by self-reported data, thereby enhancing the credibility of the research outcomes. Third, the dataset may not comprehensively cover certain key confounding factors, such as the severity of strokes, rehabilitation compliance, and levels of cognitive impairment, which could affect the accuracy of the predictions. Fourth, we recommend considering the use of estimation methods to handle missing data, to make more comprehensive use of the available information, and to reduce potential biases caused by a reduced sample size. This approach will provide more data support for model training, further enhancing the model’s generalization ability and performance. Lastly, this study lacks external validation, which may impact the generalizability of the results. When other researchers apply this model to patient populations with different clinical characteristics, the model’s performance might differ from the results observed in the NHANES dataset. Therefore, conducting external validation studies is essential to confirm the model’s efficacy and reliability in broader populations. Future studies should conduct external validation across diverse independent samples and settings to ascertain the model’s reliability and applicability.

Conclusions

We successfully developed an interpretable ML model aimed at predicting depression risk in stroke patients based on clinical data. Following rigorous validation, our XGBoost model demonstrated superior predictive capabilities, establishing a strong foundation for its future application in clinical settings. We hope that this predictive model can serve as an auxiliary tool to help develop more accurate and personalized treatment plans, thereby enhancing the mental health and overall quality of life of stroke patients.