Update src/train_model.py

src/train_model.py

@@ -588,6 +588,11 @@ def generate_earthquake_data(n: int = 3000):
         "focal_depth_km", "tectonic_stress_index", "building_vulnerability",
         "population_density_norm", "bedrock_amplification",
     ]
+        # Add immediately before the assert
+    print(f"[Earthquake] Columns available before assert: {list(base.columns)}")
+    print(f"[Earthquake] Required: {required}")
+    print(f"[Earthquake] Missing: {[c for c in required if c not in base.columns]}")
+
 
     assert features == list(EARTHQUAKE_FEATURES), (
         f"Feature mismatch!\n  train:     {features}\n"
@@ -622,22 +627,24 @@ DATA_GENERATORS = {
 def evaluate_model(model: FuzzyNeuralNetwork, X: torch.Tensor, y: torch.Tensor) -> dict:
     model.eval()
     with torch.no_grad():
-        preds = model(X).numpy()
+        preds = model(X).numpy().flatten()
     y_np = y.numpy()
 
-    # Binarize at 0.5 for AUC
+    threshold = float(np.median(y_np))
     try:
-        auc = roc_auc_score((y_np > 0.5).astype(int), preds)
-    except Exception:
-        auc = float('nan')
+        auc = roc_auc_score((y_np > threshold).astype(int), preds)
+    except ValueError as e:
+        auc = float("nan")
+        print(f"  [Warning] AUC undefined: {e}")
 
     mae = mean_absolute_error(y_np, preds)
 
     return {
-        "MAE": round(float(mae), 4),
-        "AUC-ROC": round(float(auc), 4),
+        "MAE":             round(float(mae), 4),
+        "AUC-ROC":         round(float(auc), 4) if not np.isnan(auc) else "nan",
         "Mean Prediction": round(float(preds.mean()), 4),
-        "Std Prediction": round(float(preds.std()), 4),
+        "Mean Label":      round(float(y_np.mean()), 4),
+        "Std Prediction":  round(float(preds.std()), 4),
     }
 
 
@@ -672,7 +679,9 @@ def train_all(epochs: int = 200):
     print("  TRAINING SUMMARY")
     print("="*60)
     for dt, metrics in results.items():
-        print(f"  {dt.upper():12s} | MAE: {metrics['MAE']:.4f} | AUC: {metrics['AUC-ROC']:.4f}")
+        auc = metrics["AUC-ROC"]
+        auc_str = f"{auc:.4f}" if isinstance(auc, float) and not np.isnan(auc) else str(auc)
+        print(f"  {dt.upper():12s} | MAE: {metrics['MAE']:.4f} | AUC: {auc_str}")
     print("="*60)