predict_insiders.py

#!/usr/bin/env python3
"""
Prediction script for using the trained unsupervised ensemble model to detect insiders in new data.
"""

import argparse
import numpy as np
import pandas as pd
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score, confusion_matrix

def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description='Make predictions using the trained unsupervised ensemble model')
    parser.add_argument('--input', type=str, required=True, help='Path to input data file (CSV)')
    parser.add_argument('--model', type=str, default='unsupervised_model.joblib', help='Path to model file')
    parser.add_argument('--output', type=str, default='insider_predictions.csv', help='Path to output predictions')
    parser.add_argument('--threshold', type=float, default=None, help='Custom threshold (if not using model default)')
    parser.add_argument('--label_col', type=str, default=None, help='Name of the label column if available for evaluation')
    return parser.parse_args()

def load_model(model_path):
    """Load the trained model"""
    print(f"Loading model from {model_path}")
    try:
        model = joblib.load(model_path)
        print("Model loaded successfully")
        return model
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

def load_data(input_path):
    """Load the input data"""
    print(f"Loading data from {input_path}")
    try:
        df = pd.read_csv(input_path)
        print(f"Data loaded successfully: {df.shape[0]} rows, {df.shape[1]} columns")
        return df
    except Exception as e:
        print(f"Error loading data: {e}")
        return None

def preprocess_data(df, label_col=None):
    """Preprocess the data for prediction"""
    print("Preprocessing data...")
    
    # Check if the data contains user column
    if 'user' not in df.columns:
        print("Warning: 'user' column not found. Creating a dummy user column.")
        df['user'] = [f"user_{i}" for i in range(df.shape[0])]
    
    # Extract features if this is raw data
    try:
        from unsupervised_precision_recall import extract_insider_features
        df = extract_insider_features(df)
        print("Applied insider feature extraction")
    except Exception as e:
        print(f"Warning: Could not apply feature extraction: {e}")
        print("Proceeding with raw features")
    
    # Prepare X and y (if label available)
    y = None
    if label_col and label_col in df.columns:
        y = df[label_col].values
        X = df.drop([label_col, 'user'], axis=1, errors='ignore')
    else:
        X = df.drop(['user'], axis=1, errors='ignore')
    
    # Fill NaN values
    X = X.fillna(0)
    
    print(f"Processed data shape: {X.shape}")
    return X, y, df['user']

def make_predictions(model, X, threshold=None):
    """
    Make predictions using the trained model
    
    Parameters:
    -----------
    model : object
        Trained model
    X : array-like
        Features to predict on
    threshold : float, optional
        Custom threshold for binary predictions
        
    Returns:
    --------
    results : DataFrame
        DataFrame with predictions and scores
    """
    print("Making predictions...")
    
    try:
        # Get anomaly scores
        anomaly_scores = model.predict_proba(X)
        
        # Use model's threshold if not provided
        if threshold is None and hasattr(model, 'threshold'):
            threshold = model.threshold
            print(f"Using model threshold: {threshold:.4f}")
        elif threshold is None:
            threshold = 0.5
            print(f"No threshold found. Using default: {threshold:.4f}")
        else:
            print(f"Using provided threshold: {threshold:.4f}")
        
        # Make binary predictions
        predictions = (anomaly_scores >= threshold).astype(int)
        
        # Create results DataFrame
        results = pd.DataFrame({
            'anomaly_score': anomaly_scores,
            'is_insider': predictions
        })
        
        # Summary statistics
        print(f"\nPrediction Summary:")
        print(f"Total records: {len(predictions)}")
        print(f"Predicted insiders: {np.sum(predictions)} ({np.mean(predictions)*100:.2f}%)")
        print(f"Predicted normal: {len(predictions) - np.sum(predictions)} ({(1-np.mean(predictions))*100:.2f}%)")
        
        return results
    
    except Exception as e:
        print(f"Error making predictions: {e}")
        return None

def evaluate_predictions(y_true, results):
    """
    Evaluate predictions against ground truth
    
    Parameters:
    -----------
    y_true : array-like
        True labels
    results : DataFrame
        DataFrame with predictions
        
    Returns:
    --------
    metrics : dict
        Evaluation metrics
    """
    if y_true is None:
        print("No labels provided. Skipping evaluation.")
        return None
    
    print("\nEvaluating predictions...")
    
    y_pred = results['is_insider'].values
    
    # Calculate metrics
    metrics = {
        'precision': precision_score(y_true, y_pred),
        'recall': recall_score(y_true, y_pred),
        'f1': f1_score(y_true, y_pred),
        'accuracy': accuracy_score(y_true, y_pred)
    }
    
    # Print metrics
    print(f"Precision: {metrics['precision']:.4f}")
    print(f"Recall: {metrics['recall']:.4f}")
    print(f"F1-Score: {metrics['f1']:.4f}")
    print(f"Accuracy: {metrics['accuracy']:.4f}")
    
    # Confusion matrix
    cm = confusion_matrix(y_true, y_pred)
    print("\nConfusion Matrix:")
    print(cm)
    
    # Plot confusion matrix
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
               xticklabels=['Normal', 'Insider'],
               yticklabels=['Normal', 'Insider'])
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.savefig('prediction_confusion_matrix.png')
    plt.close()
    
    # Calculate ROC curve if possible
    try:
        from sklearn.metrics import roc_curve, auc
        
        fpr, tpr, _ = roc_curve(y_true, results['anomaly_score'].values)
        roc_auc = auc(fpr, tpr)
        
        # Plot ROC curve
        plt.figure(figsize=(8, 6))
        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('Receiver Operating Characteristic (ROC) Curve')
        plt.legend(loc='lower right')
        plt.savefig('prediction_roc_curve.png')
        plt.close()
    except Exception as e:
        print(f"Warning: Could not generate ROC curve: {e}")
    
    return metrics

def main():
    """Main function"""
    # Parse arguments
    args = parse_args()
    
    # Load model
    model = load_model(args.model)
    if model is None:
        return
    
    # Load data
    df = load_data(args.input)
    if df is None:
        return
    
    # Preprocess data
    X, y, users = preprocess_data(df, args.label_col)
    
    # Make predictions
    results = make_predictions(model, X, args.threshold)
    if results is None:
        return
    
    # Add user column to results
    results['user'] = users.values
    
    # Move user column to front
    cols = results.columns.tolist()
    cols = ['user'] + [col for col in cols if col != 'user']
    results = results[cols]
    
    # Sort by anomaly score (descending)
    results = results.sort_values('anomaly_score', ascending=False)
    
    # Save predictions
    results.to_csv(args.output, index=False)
    print(f"\nPredictions saved to {args.output}")
    
    # Evaluate if labels are available
    if y is not None:
        metrics = evaluate_predictions(y, results)
        if metrics:
            # Save metrics
            pd.DataFrame([metrics]).to_csv('prediction_metrics.csv', index=False)
            print(f"Evaluation metrics saved to prediction_metrics.csv")

if __name__ == "__main__":
    main()