""" 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