ahadalii/Predictive_Maintenance_System

1"""
2Data Preprocessing Module for AI4I 2020 Predictive Maintenance Dataset
3Handles missing values, encoding, scaling, and train-test splitting
4"""
5
6import pandas as pd
7import numpy as np
8from sklearn.model_selection import train_test_split
9from sklearn.preprocessing import StandardScaler, LabelEncoder
10from sklearn.impute import SimpleImputer
11
12class DataPreprocessor:
13    def __init__(self, data_path='ai4i2020.csv'):
14        """Initialize the preprocessor"""
15        self.df = pd.read_csv(data_path)
16        self.scaler = StandardScaler()
17        self.label_encoder = LabelEncoder()
18        self.feature_columns = None
19        self.is_fitted = False
20        
21    def create_features(self):
22        """Create additional features"""
23        # Temperature difference
24        self.df['Temperature difference [K]'] = (
25            self.df['Process temperature [K]'] - 
26            self.df['Air temperature [K]']
27        )
28        
29        # Power calculation
30        self.df['Power [W]'] = (
31            self.df['Rotational speed [rpm]'] * 
32            self.df['Torque [Nm]'] / 9.5488
33        )
34        
35        # Tool wear rate (if we had time data, but we'll use a proxy)
36        # We can create bins for tool wear
37        self.df['Tool wear category'] = pd.cut(
38            self.df['Tool wear [min]'],
39            bins=[0, 50, 100, 150, 200, 300],
40            labels=['Very Low', 'Low', 'Medium', 'High', 'Very High']
41        )
42        
43    def handle_missing_values(self):
44        """Handle missing values"""
45        # Check for missing values
46        missing = self.df.isnull().sum()
47        
48        if missing.sum() > 0:
49            # For numerical columns, use mean imputation
50            numerical_cols = self.df.select_dtypes(include=[np.number]).columns
51            imputer = SimpleImputer(strategy='mean')
52            self.df[numerical_cols] = imputer.fit_transform(self.df[numerical_cols])
53            
54            # For categorical columns, use mode imputation
55            categorical_cols = self.df.select_dtypes(include=['object']).columns
56            for col in categorical_cols:
57                if self.df[col].isnull().sum() > 0:
58                    mode_value = self.df[col].mode()[0]
59                    self.df[col].fillna(mode_value, inplace=True)
60        else:
61            print("No missing values found in the dataset.")
62    
63    def encode_categorical_variables(self):
64        """Encode categorical variables"""
65        # Encode Type column
66        self.df['Type_encoded'] = self.label_encoder.fit_transform(self.df['Type'])
67        
68        # One-hot encode Type (alternative approach)
69        type_dummies = pd.get_dummies(self.df['Type'], prefix='Type')
70        self.df = pd.concat([self.df, type_dummies], axis=1)
71        
72        # Encode Tool wear category if it exists
73        if 'Tool wear category' in self.df.columns:
74            self.df['Tool_wear_category_encoded'] = LabelEncoder().fit_transform(
75                self.df['Tool wear category'].astype(str)
76            )
77    
78    def select_features(self):
79        """Select features for modeling"""
80        # Drop non-feature columns
81        columns_to_drop = [
82            'UDI', 'Product ID', 'Type', 'Tool wear category'
83        ]
84        
85        # Keep only relevant columns
86        feature_columns = [
87            'Air temperature [K]',
88            'Process temperature [K]',
89            'Rotational speed [rpm]',
90            'Torque [Nm]',
91            'Tool wear [min]',
92            'Temperature difference [K]',
93            'Power [W]',
94            'Type_encoded',
95            'Type_H',
96            'Type_L',
97            'Type_M'
98        ]
99        
100        # Remove columns that don't exist
101        feature_columns = [col for col in feature_columns if col in self.df.columns]
102        
103        self.feature_columns = feature_columns
104        return feature_columns
105    
106    def scale_features(self, X_train, X_test):
107        """Scale numerical features"""
108        # Scale training data
109        X_train_scaled = self.scaler.fit_transform(X_train)
110        X_test_scaled = self.scaler.transform(X_test)
111        
112        # Convert back to DataFrame
113        X_train_scaled = pd.DataFrame(
114            X_train_scaled,
115            columns=X_train.columns,
116            index=X_train.index
117        )
118        X_test_scaled = pd.DataFrame(
119            X_test_scaled,
120            columns=X_test.columns,
121            index=X_test.index
122        )
123        
124        return X_train_scaled, X_test_scaled
125    
126    def prepare_data(self, target='Machine failure', test_size=0.2, random_state=42):
127        """Complete preprocessing pipeline"""
128        print("Starting data preprocessing...")
129        
130        # Step 1: Create features
131        print("1. Creating additional features...")
132        self.create_features()
133        
134        # Step 2: Handle missing values
135        print("2. Handling missing values...")
136        self.handle_missing_values()
137        
138        # Step 3: Encode categorical variables
139        print("3. Encoding categorical variables...")
140        self.encode_categorical_variables()
141        
142        # Step 4: Select features
143        print("4. Selecting features...")
144        feature_columns = self.select_features()
145        
146        # Step 5: Prepare X and y
147        X = self.df[feature_columns]
148        y = self.df[target]
149        
150        # Step 6: Split data
151        print("5. Splitting data into train and test sets...")
152        X_train, X_test, y_train, y_test = train_test_split(
153            X, y, test_size=test_size, random_state=random_state, stratify=y
154        )
155        
156        # Step 7: Scale features
157        print("6. Scaling features...")
158        X_train_scaled, X_test_scaled = self.scale_features(X_train, X_test)
159        
160        self.is_fitted = True
161        
162        print("Preprocessing complete!")
163        print(f"Training set shape: {X_train_scaled.shape}")
164        print(f"Test set shape: {X_test_scaled.shape}")
165        print(f"Features: {feature_columns}")
166        
167        return X_train_scaled, X_test_scaled, y_train, y_test, feature_columns
168    
169    def preprocess_new_data(self, new_data):
170        """Preprocess new data for prediction (using fitted scaler and encoders)"""
171        if not self.is_fitted:
172            raise ValueError("Preprocessor must be fitted first using prepare_data()")
173        
174        # Create a copy
175        df_new = new_data.copy()
176        
177        # Create features
178        df_new['Temperature difference [K]'] = (
179            df_new['Process temperature [K]'] - 
180            df_new['Air temperature [K]']
181        )
182        df_new['Power [W]'] = (
183            df_new['Rotational speed [rpm]'] * 
184            df_new['Torque [Nm]'] / 9.5488
185        )
186        
187        # Encode Type
188        df_new['Type_encoded'] = self.label_encoder.transform(df_new['Type'])
189        type_dummies = pd.get_dummies(df_new['Type'], prefix='Type')
190        
191        # Ensure all Type columns exist
192        for col in ['Type_H', 'Type_L', 'Type_M']:
193            if col not in type_dummies.columns:
194                type_dummies[col] = 0
195        
196        df_new = pd.concat([df_new, type_dummies[['Type_H', 'Type_L', 'Type_M']]], axis=1)
197        
198        # Select features
199        X_new = df_new[self.feature_columns]
200        
201        # Scale
202        X_new_scaled = self.scaler.transform(X_new)
203        X_new_scaled = pd.DataFrame(X_new_scaled, columns=self.feature_columns)
204        
205        return X_new_scaled