Update app.py

app.py

@@ -271,10 +271,14 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
 
     X = torch.FloatTensor(X).to(device)
 
+    # Ensure X is 2D
+    if X.dim() == 1:
+        X = X.unsqueeze(0)
+
     train_size = int(0.85 * len(X))
     X_train, X_val = X[:train_size], X[train_size:]
 
-    model = LSTMAutoencoder(input_size=len(feature_columns)).to(device)
+    model = LSTMAutoencoder(input_size=X.shape[1]).to(device)
     criterion = nn.MSELoss()
     optimizer = optim.Adam(model.parameters())
 
@@ -295,7 +299,7 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
     with torch.no_grad():
         reconstructed = model(X.unsqueeze(0)).squeeze(0).cpu().numpy()
 
-# Compute anomalies for all features
+    # Compute anomalies for all features
     mse_all = np.mean(np.power(X.cpu().numpy() - reconstructed, 2), axis=1)
     top_indices_all = mse_all.argsort()[-num_anomalies:][::-1]
     anomalies_all = np.zeros(len(mse_all), dtype=bool)
@@ -304,7 +308,12 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
     # Compute anomalies for components only
     component_columns = [col for col in feature_columns if col.startswith('Comp')]
     component_indices = [feature_columns.index(col) for col in component_columns]
-    mse_comp = np.mean(np.power(X.cpu().numpy()[:, component_indices] - reconstructed[:, component_indices], 2), axis=1)
+    
+    if len(component_indices) > 0:
+        mse_comp = np.mean(np.power(X.cpu().numpy()[:, component_indices] - reconstructed[:, component_indices], 2), axis=1)
+    else:
+        mse_comp = mse_all  # If no components, use all features
+
     top_indices_comp = mse_comp.argsort()[-num_anomalies:][::-1]
     anomalies_comp = np.zeros(len(mse_comp), dtype=bool)
     anomalies_comp[top_indices_comp] = True