DecisionTreeModel02.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import joblib
import os
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.model_selection import train_test_split

# Define data directory
data_path = r'C:\Users\jsntg\OneDrive\Desktop\3YP\code\data\test_data_damn.csv'
graph_path = r'C:\Users\jsntg\OneDrive\Desktop\3YP\code\graph'
result_path = r'C:\Users\jsntg\OneDrive\Desktop\3YP\code\result'

class SepsisDecisionTreeExtended:
    # Initialize
    def __init__(self):  
        self.data_path = data_path
        self.graph_path = graph_path
        self.result_path = result_path
        self.model = None
        self.feature_names = None
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None

    # Load data    
    def load_data(self):
        print(f"Loading data from: {self.data_path}")
        
        if not os.path.exists(self.data_path):
            raise FileNotFoundError(f"\nError! Data file not found")

        df = pd.read_csv(self.data_path)
        print(f"Data shape: {df.shape}")
        print(f"Columns: {df.columns.tolist()}")
        
        return df
    
    # Calculate simple risk score
    def calculate_risk_score(self, row):
        risk_score = 0
        
        if row['respiratory_rate'] > 20 or row['respiratory_rate'] < 12:
            risk_score += 1        
        if row['oxygen_saturations'] < 94:
            risk_score += 1        
        #if row['supplemental_oxygen'] == 1:
            #risk_score += 2        
        if row['temperature'] > 100.4 or row['temperature'] < 96.8:
            risk_score += 1        
        if row['systolic_bp'] < 100 or row['systolic_bp'] > 140:
            risk_score += 1        
        if row['diastolic_bp'] < 60 or row['diastolic_bp'] > 90:
            risk_score += 1        
        if row['heart_rate'] > 100 or row['heart_rate'] < 60:
            risk_score += 1    
        #if row['pain'] > 7:
            #risk_score += 1        
        #if row['proteinuria'] == 1:
            #risk_score += 1        
        return risk_score

    # Preprocess data
    def preprocess_data(self, df, lookback=5):
        print(f"Preprocessing data: current features + risk level from past {lookback} hours")
        
        all_patients = sorted(df['patient_id'].unique())
        split_idx = int(len(all_patients) * 0.8)
        train_patients = all_patients[:split_idx]
        test_patients = all_patients[split_idx:]
        
        print(f"Training Set: {train_patients}")
        print(f"Testing Set: {test_patients}")
        
        X_train, y_train = self._create_risk_score_features(df[df['patient_id'].isin(train_patients)], lookback)
        X_test, y_test = self._create_risk_score_features(df[df['patient_id'].isin(test_patients)], lookback)
        
        # Create feature names
        feature_names = []
        current_features = [
            'respiratory_rate', 'oxygen_saturations', 'supplemental_oxygen',
            'temperature', 'systolic_bp', 'diastolic_bp', 'heart_rate',
            'pain', 'proteinuria', 'gender_encoded', 'race_encoded'
        ]
        
        # Check for missing columns
        missing_cols = [col for col in current_features if col not in df.columns]
        if missing_cols:
            print(f"Warning! The following colums are missing in the data: {missing_cols}")
            current_features = [col for col in current_features if col not in missing_cols]
        
        for feature in current_features:
            feature_names.append(f"{feature}_current")
        
        for i in range(1, lookback + 1):
            feature_names.append(f"risk_score_t-{i}")
        
        print(f"Number of features: {len(feature_names)} = {len(current_features)} current features + {lookback} past risk scores")
        print(f"Training set feature shape: {X_train.shape}")
        print(f"Testing set feature shape: {X_test.shape}")
        
        # Save to instance variables
        self.X_train = X_train
        self.X_test = X_test
        self.y_train = y_train
        self.y_test = y_test
        self.feature_names = feature_names
        
        return X_train, X_test, y_train, y_test, feature_names
    
    def _create_risk_score_features(self, df, lookback):
        features_list = []
        labels_list = []
        
        for patient_id in df['patient_id'].unique():
            patient_data = df[df['patient_id'] == patient_id].copy()
            
            # Order by time
            if 'timestamp' in patient_data.columns:
                patient_data = patient_data.sort_values('timestamp')
            
            patient_data['risk_score'] = patient_data.apply(self.calculate_risk_score, axis=1)
            current_features = [
                'respiratory_rate', 'oxygen_saturations', 'supplemental_oxygen',
                'temperature', 'systolic_bp', 'diastolic_bp', 'heart_rate',
                'pain', 'proteinuria', 'gender_encoded', 'race_encoded'
            ]
            
            current_features = [col for col in current_features if col in patient_data.columns]
            
            # Create sliding window features
            for i in range(lookback, len(patient_data)):
                current_time_features = patient_data.iloc[i][current_features].values
                past_risk_scores = patient_data.iloc[i-lookback:i]['risk_score'].values
                combined_features = np.concatenate([current_time_features, past_risk_scores])
                label = patient_data.iloc[i]['sepsis']
                
                features_list.append(combined_features)
                labels_list.append(label)
        
        return np.array(features_list), np.array(labels_list)

    # Train model
    def train_model(self, max_depth=5, min_samples_split=10, min_samples_leaf=5):
        print("\nTraining decision tree model...")
        
        self.model = DecisionTreeClassifier(
            random_state=42,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            class_weight='balanced'
        )
        
        self.model.fit(self.X_train, self.y_train)
        
        print(f"Tree depth: {self.model.get_depth()}")
        print(f"Number of leaves: {self.model.get_n_leaves()}")
        
        return self.model
    
    # Evaluate model accuracy
    def evaluate_model(self):
        if self.model is None:
            raise ValueError("\nError! Please complete training before evaluation.")
        
        y_train_pred = self.model.predict(self.X_train)
        y_test_pred = self.model.predict(self.X_test)
        
        train_accuracy = accuracy_score(self.y_train, y_train_pred)
        test_accuracy = accuracy_score(self.y_test, y_test_pred)
        
        print(f"Training set accuracy: {train_accuracy:.4f} ({train_accuracy*100:.2f}%)")
        print(f"Testing set accuracy: {test_accuracy:.4f} ({test_accuracy*100:.2f}%)")
        
        print("\nTraining set classification report:")
        print(classification_report(self.y_train, y_train_pred))
        
        print("\nTesting set classification report:")
        print(classification_report(self.y_test, y_test_pred))
        
        results = {
            'train_accuracy': train_accuracy,
            'test_accuracy': test_accuracy,
            'train_report': classification_report(self.y_train, y_train_pred, output_dict=True),
            'test_report': classification_report(self.y_test, y_test_pred, output_dict=True),
            'confusion_matrix_train': confusion_matrix(self.y_train, y_train_pred),
            'confusion_matrix_test': confusion_matrix(self.y_test, y_test_pred)
        }
        
        return results
    
    # Visualize feature importance
    def plot_feature_importance(self, top_n=20):
        if self.model is None:
            raise ValueError("\nError! Please complete training before visualization.")
        
        feature_importance = self.model.feature_importances_
        sorted_idx = np.argsort(feature_importance)[::-1]
        
        # Display the top n features only
        if top_n < len(feature_importance):
            display_idx = sorted_idx[:top_n]
        else:
            display_idx = sorted_idx
        
        # Plot feature importance
        plt.figure(figsize=(12, 8))
        plt.bar(range(len(display_idx)), feature_importance[display_idx])
        plt.xticks(range(len(display_idx)), 
                  np.array(self.feature_names)[display_idx], 
                  rotation=45, ha='right')
        plt.xlabel('Feature')
        plt.ylabel('Importance')
        plt.title(f'Decision Tree Feature Importance (Top {len(display_idx)})')
        plt.tight_layout()
        
        # Save image
        save_path = os.path.join(self.graph_path, 'feature_importance_model02_damn.png')
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
        
        print(f"Feature importance plot saved to: {save_path}")
        
        # Print feature importance
        print(f"Top {len(display_idx)} important features")
        for i, idx in enumerate(display_idx):
            print(f"{i+1:2d}. {self.feature_names[idx]:30s} : {feature_importance[idx]:.4f}")
    
    # Visualize confusion matrices
    def plot_confusion_matrices(self):
        if self.model is None:
            raise ValueError("Error! Please complete training before visualization.")
        
        y_train_pred = self.model.predict(self.X_train)
        y_test_pred = self.model.predict(self.X_test)
        
        cm_train = confusion_matrix(self.y_train, y_train_pred)
        cm_test = confusion_matrix(self.y_test, y_test_pred)
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 6))
        
        # subplot 1: training set confusion matrix
        sns.heatmap(cm_train, annot=True, fmt='d', cmap='Blues', ax=axes[0])
        axes[0].set_title('Training Set Confusion Matrix')
        axes[0].set_xlabel('Predicted Label')
        axes[0].set_ylabel('True Label')
        
        # subplot 2: testing set confusion matrix
        sns.heatmap(cm_test, annot=True, fmt='d', cmap='Blues', ax=axes[1])
        axes[1].set_title('Testing Set Confusion Matrix')
        axes[1].set_xlabel('Predicted Label')
        axes[1].set_ylabel('True Label')
        
        plt.tight_layout()
        
        # Save image
        save_path = os.path.join(self.graph_path, 'confusion_matrices_model02_damn.png')
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
        
        print(f"Confusion matrices plot saved to: {save_path}")
    
    # Save model as pickle file
    def save_model(self, model_path=None):
        if self.model is None:
            raise ValueError("Error! No existing model")
        
        if model_path is None:
            model_path = os.path.join(self.result_path, 'model02_damn.pkl')
        
        joblib.dump(self.model, model_path)
        print(f"Model saved to: {model_path}")
        
        # Save feature names
        feature_names_path = os.path.join(self.result_path, 'feature_names_model02_damn.txt')
        with open(feature_names_path, 'w', encoding='utf-8') as f:
            for name in self.feature_names:
                f.write(name + '\n')
        print(f"Feature names saved to: {feature_names_path}")

# Execute main
def main():
    sepsis_model = SepsisDecisionTreeExtended()
    
    try:
        df = sepsis_model.load_data()
        sepsis_model.preprocess_data(df, lookback=5)
        sepsis_model.train_model(max_depth=5, min_samples_split=10, min_samples_leaf=5)
        results = sepsis_model.evaluate_model()
        sepsis_model.plot_feature_importance(top_n=15)
        sepsis_model.plot_confusion_matrices()
        sepsis_model.save_model()
        
    except Exception as e:
        print(f"Error! {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()