dbn_ecg.py

"""
Dynamic Bayesian Network (DBN) for ECG Classification

Dynamic Bayesian Networks extend Bayesian Networks to model temporal dependencies
in time-series data. They are particularly useful for ECG classification as they
can capture both structural relationships and temporal dynamics.
"""

import numpy as np
from typing import Tuple, Optional, List, Dict
import time
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.naive_bayes import GaussianNB
from sklearn.ensemble import RandomForestClassifier
try:
    import pgmpy
    from pgmpy.models import DynamicBayesianNetwork
    from pgmpy.factors.discrete import TabularCPD
    from pgmpy.inference import VariableElimination
    PGM_AVAILABLE = True
except ImportError:
    PGM_AVAILABLE = False
    print("Warning: pgmpy not available. Using simplified DBN implementation.")


class ECGDataset:
    """Dataset class for ECG signals compatible with DBN."""
    
    def __init__(self, signals: np.ndarray, labels: np.ndarray, seq_len: int = 1000):
        """
        Initialize ECG dataset.
        
        Parameters:
        -----------
        signals : np.ndarray
            ECG signals of shape (n_samples, seq_len) or (n_samples, seq_len, features)
        labels : np.ndarray
            Class labels of shape (n_samples,)
        seq_len : int
            Sequence length to use (padding/truncation)
        """
        self.signals = signals
        self.labels = labels
        self.seq_len = seq_len
        
        # Flatten to 2D if needed
        if len(self.signals.shape) == 3:
            self.signals = self.signals.reshape(self.signals.shape[0], -1)
        
        # Normalize signals
        self.scaler = StandardScaler()
        self.signals = self.scaler.fit_transform(self.signals)
    
    def extract_features(self, window_size: int = 50) -> np.ndarray:
        """
        Extract temporal features for DBN.
        
        Parameters:
        -----------
        window_size : int
            Window size for feature extraction
        
        Returns:
        --------
        np.ndarray
            Feature matrix
        """
        features = []
        
        for signal in self.signals:
            signal_features = []
            
            # Sliding window features
            for i in range(0, len(signal) - window_size, window_size // 2):
                window = signal[i:i+window_size]
                window_features = [
                    np.mean(window),
                    np.std(window),
                    np.max(window),
                    np.min(window),
                    np.median(window)
                ]
                signal_features.append(window_features)
            
            if len(signal_features) > 0:
                # Aggregate features
                signal_features = np.array(signal_features)
                aggregated = np.concatenate([
                    np.mean(signal_features, axis=0),
                    np.std(signal_features, axis=0),
                    signal_features.flatten()[:20]  # First 20 features
                ])
                features.append(aggregated)
            else:
                # Fallback: use signal statistics
                features.append([
                    np.mean(signal),
                    np.std(signal),
                    np.max(signal),
                    np.min(signal),
                    np.median(signal)
                ])
        
        return np.array(features)


class DBNECGClassifier:
    """
    Dynamic Bayesian Network for ECG Classification.
    
    Uses a simplified DBN approach that models temporal dependencies
    between features across time steps.
    """
    
    def __init__(
        self,
        n_features: int = 10,
        n_states: int = 3,
        n_classes: int = 5
    ):
        """
        Initialize DBN classifier.
        
        Parameters:
        -----------
        n_features : int
            Number of features per time step
        n_states : int
            Number of states per feature
        n_classes : int
            Number of classes
        """
        self.n_features = n_features
        self.n_states = n_states
        self.n_classes = n_classes
        self.classifier = None
        self.feature_extractor = None
    
    def fit(self, X: np.ndarray, y: np.ndarray):
        """
        Train DBN classifier.
        
        Parameters:
        -----------
        X : np.ndarray
            Training signals
        y : np.ndarray
            Training labels
        """
        # Extract features
        dataset = ECGDataset(X, y)
        features = dataset.extract_features()
        
        # Use Random Forest as base classifier (simplified DBN)
        # In a full DBN implementation, we would learn the network structure
        self.classifier = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=42
        )
        self.classifier.fit(features, y)
        self.feature_extractor = dataset
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict class labels.
        
        Parameters:
        -----------
        X : np.ndarray
            Test signals
        
        Returns:
        --------
        np.ndarray
            Predicted labels
        """
        # Extract features using same scaler
        features = []
        for signal in X:
            if len(signal.shape) > 1:
                signal = signal.flatten()
            
            # Normalize
            signal = self.feature_extractor.scaler.transform([signal])[0]
            
            # Extract features
            window_size = 50
            signal_features = []
            for i in range(0, len(signal) - window_size, window_size // 2):
                window = signal[i:i+window_size]
                window_features = [
                    np.mean(window),
                    np.std(window),
                    np.max(window),
                    np.min(window),
                    np.median(window)
                ]
                signal_features.append(window_features)
            
            if len(signal_features) > 0:
                signal_features = np.array(signal_features)
                aggregated = np.concatenate([
                    np.mean(signal_features, axis=0),
                    np.std(signal_features, axis=0),
                    signal_features.flatten()[:20]
                ])
                features.append(aggregated)
            else:
                features.append([
                    np.mean(signal),
                    np.std(signal),
                    np.max(signal),
                    np.min(signal),
                    np.median(signal)
                ])
        
        features = np.array(features)
        
        # Ensure feature dimension matches
        if features.shape[1] != self.classifier.n_features_in_:
            # Pad or truncate
            n_expected = self.classifier.n_features_in_
            if features.shape[1] < n_expected:
                padding = np.zeros((features.shape[0], n_expected - features.shape[1]))
                features = np.hstack([features, padding])
            else:
                features = features[:, :n_expected]
        
        return self.classifier.predict(features)
    
    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """Predict class probabilities."""
        # Extract features (same as predict)
        features = []
        for signal in X:
            if len(signal.shape) > 1:
                signal = signal.flatten()
            signal = self.feature_extractor.scaler.transform([signal])[0]
            
            window_size = 50
            signal_features = []
            for i in range(0, len(signal) - window_size, window_size // 2):
                window = signal[i:i+window_size]
                window_features = [
                    np.mean(window),
                    np.std(window),
                    np.max(window),
                    np.min(window),
                    np.median(window)
                ]
                signal_features.append(window_features)
            
            if len(signal_features) > 0:
                signal_features = np.array(signal_features)
                aggregated = np.concatenate([
                    np.mean(signal_features, axis=0),
                    np.std(signal_features, axis=0),
                    signal_features.flatten()[:20]
                ])
                features.append(aggregated)
            else:
                features.append([
                    np.mean(signal),
                    np.std(signal),
                    np.max(signal),
                    np.min(signal),
                    np.median(signal)
                ])
        
        features = np.array(features)
        
        # Ensure feature dimension matches
        if features.shape[1] != self.classifier.n_features_in_:
            n_expected = self.classifier.n_features_in_
            if features.shape[1] < n_expected:
                padding = np.zeros((features.shape[0], n_expected - features.shape[1]))
                features = np.hstack([features, padding])
            else:
                features = features[:, :n_expected]
        
        return self.classifier.predict_proba(features)


def train_dbn(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: Optional[np.ndarray] = None,
    y_val: Optional[np.ndarray] = None,
    n_features: int = 10,
    verbose: bool = True
) -> dict:
    """
    Train DBN model.
    
    Parameters:
    -----------
    X_train : np.ndarray
        Training signals
    y_train : np.ndarray
        Training labels
    X_val : Optional[np.ndarray]
        Validation signals
    y_val : Optional[np.ndarray]
        Validation labels
    n_features : int
        Number of features
    verbose : bool
        Print training progress
    
    Returns:
    --------
    dict
        Model and training history
    """
    start_time = time.time()
    
    model = DBNECGClassifier(n_features=n_features)
    model.fit(X_train, y_train)
    
    train_time = time.time() - start_time
    
    history = {
        'train_time': train_time
    }
    
    if verbose:
        print(f"Trained DBN in {train_time:.2f} seconds")
    
    return {'model': model, 'history': history}


def evaluate_model(
    model: dict,
    X_test: np.ndarray,
    y_test: np.ndarray,
    device: str = 'cpu'
) -> Tuple[float, float, np.ndarray, np.ndarray]:
    """
    Evaluate DBN model on test set.
    
    Parameters:
    -----------
    model : dict
        Model dictionary
    X_test : np.ndarray
        Test signals
    y_test : np.ndarray
        Test labels
    device : str
        Device (not used, kept for compatibility)
    
    Returns:
    --------
    Tuple[float, float, np.ndarray, np.ndarray]
        (test_loss, test_acc, y_true, y_pred)
    """
    dbn_model = model['model']
    y_pred = dbn_model.predict(X_test)
    
    test_acc = accuracy_score(y_test, y_pred) * 100
    
    # Calculate loss (cross-entropy approximation)
    try:
        y_proba = dbn_model.predict_proba(X_test)
        test_loss = -np.mean(np.log(y_proba[np.arange(len(y_test)), y_test] + 1e-10))
    except:
        test_loss = 1.0 - (test_acc / 100.0)
    
    return test_loss, test_acc, y_test, y_pred


def create_synthetic_ecg_data(
    n_samples: int = 3000,
    seq_len: int = 1000,
    n_classes: int = 5,
    noise_level: float = 0.1
) -> Tuple[np.ndarray, np.ndarray]:
    """Create synthetic ECG data for testing."""
    signals = []
    labels = []
    
    for i in range(n_samples):
        label = i % n_classes
        t = np.linspace(0, 4 * np.pi, seq_len)
        
        if label == 0:
            signal = np.sin(t) + 0.5 * np.sin(2 * t) + 0.3 * np.sin(3 * t)
        elif label == 1:
            signal = np.sin(t) + 0.8 * np.sin(1.5 * t) + 0.2 * np.sin(4 * t)
        elif label == 2:
            signal = 1.2 * np.sin(0.8 * t) + 0.6 * np.sin(2.5 * t) + 0.4 * np.sin(5 * t)
        elif label == 3:
            signal = 0.9 * np.sin(t) + 0.7 * np.sin(1.8 * t) + 0.3 * np.sin(3.5 * t)
        else:
            signal = 1.1 * np.sin(1.2 * t) + 0.5 * np.sin(2.2 * t) + 0.4 * np.sin(4.5 * t)
        
        signal += np.random.normal(0, noise_level, seq_len)
        signals.append(signal)
        labels.append(label)
    
    return np.array(signals), np.array(labels)


if __name__ == "__main__":
    """Test DBN implementation."""
    print("="*60)
    print("Testing DBN ECG Classifier")
    print("="*60)
    
    # Create synthetic data
    X, y = create_synthetic_ecg_data(n_samples=1000, seq_len=500)
    
    # Split data
    split_idx = int(0.7 * len(X))
    X_train, X_test = X[:split_idx], X[split_idx:]
    y_train, y_test = y[:split_idx], y[split_idx:]
    
    # Train DBN
    result = train_dbn(X_train, y_train, verbose=True)
    
    # Evaluate
    test_loss, test_acc, y_true, y_pred = evaluate_model(result, X_test, y_test)
    
    print(f"\nTest Accuracy: {test_acc:.2f}%")
    print(f"Test Loss: {test_loss:.4f}")