preprocessing.py

"""
Data Preprocessing Module for AI4I 2020 Predictive Maintenance Dataset
Handles missing values, encoding, scaling, and train-test splitting
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.impute import SimpleImputer

class DataPreprocessor:
    def __init__(self, data_path='ai4i2020.csv'):
        """Initialize the preprocessor"""
        self.df = pd.read_csv(data_path)
        self.scaler = StandardScaler()
        self.label_encoder = LabelEncoder()
        self.feature_columns = None
        self.is_fitted = False
        
    def create_features(self):
        """Create additional features"""
        # Temperature difference
        self.df['Temperature difference [K]'] = (
            self.df['Process temperature [K]'] - 
            self.df['Air temperature [K]']
        )
        
        # Power calculation
        self.df['Power [W]'] = (
            self.df['Rotational speed [rpm]'] * 
            self.df['Torque [Nm]'] / 9.5488
        )
        
        # Tool wear rate (if we had time data, but we'll use a proxy)
        # We can create bins for tool wear
        self.df['Tool wear category'] = pd.cut(
            self.df['Tool wear [min]'],
            bins=[0, 50, 100, 150, 200, 300],
            labels=['Very Low', 'Low', 'Medium', 'High', 'Very High']
        )
        
    def handle_missing_values(self):
        """Handle missing values"""
        # Check for missing values
        missing = self.df.isnull().sum()
        
        if missing.sum() > 0:
            # For numerical columns, use mean imputation
            numerical_cols = self.df.select_dtypes(include=[np.number]).columns
            imputer = SimpleImputer(strategy='mean')
            self.df[numerical_cols] = imputer.fit_transform(self.df[numerical_cols])
            
            # For categorical columns, use mode imputation
            categorical_cols = self.df.select_dtypes(include=['object']).columns
            for col in categorical_cols:
                if self.df[col].isnull().sum() > 0:
                    mode_value = self.df[col].mode()[0]
                    self.df[col].fillna(mode_value, inplace=True)
        else:
            print("No missing values found in the dataset.")
    
    def encode_categorical_variables(self):
        """Encode categorical variables"""
        # Encode Type column
        self.df['Type_encoded'] = self.label_encoder.fit_transform(self.df['Type'])
        
        # One-hot encode Type (alternative approach)
        type_dummies = pd.get_dummies(self.df['Type'], prefix='Type')
        self.df = pd.concat([self.df, type_dummies], axis=1)
        
        # Encode Tool wear category if it exists
        if 'Tool wear category' in self.df.columns:
            self.df['Tool_wear_category_encoded'] = LabelEncoder().fit_transform(
                self.df['Tool wear category'].astype(str)
            )
    
    def select_features(self):
        """Select features for modeling"""
        # Drop non-feature columns
        columns_to_drop = [
            'UDI', 'Product ID', 'Type', 'Tool wear category'
        ]
        
        # Keep only relevant columns
        feature_columns = [
            'Air temperature [K]',
            'Process temperature [K]',
            'Rotational speed [rpm]',
            'Torque [Nm]',
            'Tool wear [min]',
            'Temperature difference [K]',
            'Power [W]',
            'Type_encoded',
            'Type_H',
            'Type_L',
            'Type_M'
        ]
        
        # Remove columns that don't exist
        feature_columns = [col for col in feature_columns if col in self.df.columns]
        
        self.feature_columns = feature_columns
        return feature_columns
    
    def scale_features(self, X_train, X_test):
        """Scale numerical features"""
        # Scale training data
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)
        
        # Convert back to DataFrame
        X_train_scaled = pd.DataFrame(
            X_train_scaled,
            columns=X_train.columns,
            index=X_train.index
        )
        X_test_scaled = pd.DataFrame(
            X_test_scaled,
            columns=X_test.columns,
            index=X_test.index
        )
        
        return X_train_scaled, X_test_scaled
    
    def prepare_data(self, target='Machine failure', test_size=0.2, random_state=42):
        """Complete preprocessing pipeline"""
        print("Starting data preprocessing...")
        
        # Step 1: Create features
        print("1. Creating additional features...")
        self.create_features()
        
        # Step 2: Handle missing values
        print("2. Handling missing values...")
        self.handle_missing_values()
        
        # Step 3: Encode categorical variables
        print("3. Encoding categorical variables...")
        self.encode_categorical_variables()
        
        # Step 4: Select features
        print("4. Selecting features...")
        feature_columns = self.select_features()
        
        # Step 5: Prepare X and y
        X = self.df[feature_columns]
        y = self.df[target]
        
        # Step 6: Split data
        print("5. Splitting data into train and test sets...")
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=random_state, stratify=y
        )
        
        # Step 7: Scale features
        print("6. Scaling features...")
        X_train_scaled, X_test_scaled = self.scale_features(X_train, X_test)
        
        self.is_fitted = True
        
        print("Preprocessing complete!")
        print(f"Training set shape: {X_train_scaled.shape}")
        print(f"Test set shape: {X_test_scaled.shape}")
        print(f"Features: {feature_columns}")
        
        return X_train_scaled, X_test_scaled, y_train, y_test, feature_columns
    
    def preprocess_new_data(self, new_data):
        """Preprocess new data for prediction (using fitted scaler and encoders)"""
        if not self.is_fitted:
            raise ValueError("Preprocessor must be fitted first using prepare_data()")
        
        # Create a copy
        df_new = new_data.copy()
        
        # Create features
        df_new['Temperature difference [K]'] = (
            df_new['Process temperature [K]'] - 
            df_new['Air temperature [K]']
        )
        df_new['Power [W]'] = (
            df_new['Rotational speed [rpm]'] * 
            df_new['Torque [Nm]'] / 9.5488
        )
        
        # Encode Type
        df_new['Type_encoded'] = self.label_encoder.transform(df_new['Type'])
        type_dummies = pd.get_dummies(df_new['Type'], prefix='Type')
        
        # Ensure all Type columns exist
        for col in ['Type_H', 'Type_L', 'Type_M']:
            if col not in type_dummies.columns:
                type_dummies[col] = 0
        
        df_new = pd.concat([df_new, type_dummies[['Type_H', 'Type_L', 'Type_M']]], axis=1)
        
        # Select features
        X_new = df_new[self.feature_columns]
        
        # Scale
        X_new_scaled = self.scaler.transform(X_new)
        X_new_scaled = pd.DataFrame(X_new_scaled, columns=self.feature_columns)
        
        return X_new_scaled