Spaces:
Build error
Build error
| """ | |
| 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 |