Update app.py

app.py

@@ -1,5 +1,4 @@
 # app.py — ST_GeoMech_SMW
-
 import io, json, os, base64, math
 from pathlib import Path
 import streamlit as st
@@ -23,7 +22,7 @@ from sklearn.metrics import mean_squared_error, mean_absolute_error
 APP_NAME = "ST_GeoMech_SMW"
 TAGLINE  = "Real-Time Upper/Lower Mud Weight (MW) Limits For Safe Drilling"
 
-# Defaults (can be overridden by meta files)
+# Defaults (can be overridden by metas or model.feature_names_in_)
 FEATURES_DEFAULT = [
     "WOB (klbf)",
     "Torque (kft.lbf)",
@@ -45,7 +44,6 @@ BD_MODEL_PATH = MODELS_DIR / "bd_model.joblib"
 BO_META_PATH  = MODELS_DIR / "bo_meta.json"
 BD_META_PATH  = MODELS_DIR / "bd_meta.json"
 
-MODEL_FALLBACKS = [MODELS_DIR / "model.joblib", MODELS_DIR / "model.pkl"]
 COLORS = {
     "pred_bo": "#1f77b4",  # blue
     "pred_bd": "#d62728",  # red
@@ -197,10 +195,10 @@ def _build_alias_map(canonical_features: list[str], tgt_bo: str, tgt_bd: str) ->
     can_ROP    = pick(canonical_features, ["ROP (ft/h)", "ROP(ft/h)"])
     can_FR     = pick(canonical_features, ["Flow Rate (gpm)", "Flow Rate, gpm", "Flow Rate,gpm", "Flow Rate , gpm"])
 
-    # NEW: make Depth robust—map any of these to whatever your model expects
+    # Depth canonical = whatever is in training FEATURES (Depth or Depth (ft), etc.)
     can_DEPTH  = pick(canonical_features, ["Depth", "Depth (ft)", "Depth, ft", "Depth(ft)", "DEPTH, ft"])
 
-    alias = {
+    return {
         "WOB (klbf)": can_WOB, "WOB, klbf": can_WOB, "WOB(klbf)": can_WOB, "WOB( klbf)": can_WOB,
         "Torque (kft.lbf)": can_TORQUE, "Torque(kft.lbf)": can_TORQUE, "TORQUE(kft.lbf)": can_TORQUE,
         "SPP (psi)": can_SPP, "SPP(psi)": can_SPP,
@@ -208,15 +206,15 @@ def _build_alias_map(canonical_features: list[str], tgt_bo: str, tgt_bd: str) ->
         "ROP (ft/h)": can_ROP, "ROP(ft/h)": can_ROP,
         "Flow Rate (gpm)": can_FR, "Flow Rate, gpm": can_FR, "Flow Rate,gpm": can_FR, "Flow Rate , gpm": can_FR,
 
-        # Depth aliases (add plain "Depth" too)
+        # Depth aliases
         "Depth": can_DEPTH, "Depth (ft)": can_DEPTH, "Depth, ft": can_DEPTH, "Depth(ft)": can_DEPTH, "DEPTH, ft": can_DEPTH,
 
         # Targets
         "Breakout MW": tgt_bo, "BOMW": tgt_bo, "BO MW": tgt_bo,
         "Breakdown MW": tgt_bd, "BDMW": tgt_bd, "BD MW": tgt_bd,
     }
-    return alias
 
+DEPTH_CANDIDATES = ["Depth", "Depth (ft)", "Depth, ft", "Depth(ft)", "DEPTH, ft"]
 
 def _normalize_columns(df: pd.DataFrame, canonical_features: list[str], tgt_bo: str, tgt_bd: str) -> pd.DataFrame:
     out = df.copy()
@@ -225,6 +223,32 @@ def _normalize_columns(df: pd.DataFrame, canonical_features: list[str], tgt_bo:
     actual = {k: v for k, v in alias.items() if k in out.columns and k != v}
     return out.rename(columns=actual)
 
+def _coerce_depth(df: pd.DataFrame, canon_depth: str | None) -> pd.DataFrame:
+    """If the model expects a Depth-like column, rename any variant to it."""
+    if not canon_depth:
+        return df
+    if canon_depth in df.columns:
+        return df
+    for c in DEPTH_CANDIDATES:
+        if c in df.columns:
+            return df.rename(columns={c: canon_depth})
+    return df
+
+def _prepare_table(raw_df: pd.DataFrame, stage: str, features: list[str],
+                   tgt_bo: str, tgt_bd: str, canon_depth: str | None) -> pd.DataFrame:
+    """Normalize headers, coerce depth name, and validate required columns."""
+    df = _normalize_columns(raw_df, features, tgt_bo, tgt_bd)
+    df = _coerce_depth(df, canon_depth)
+    missing = [c for c in features if c not in df.columns]
+    if missing:
+        st.error(
+            f"{stage}: Missing required column(s): {missing}\n\n"
+            f"Found columns: {list(df.columns)}\n\n"
+            f"Expected features: {features}"
+        )
+        st.stop()
+    return df
+
 def ensure_cols(df: pd.DataFrame, cols: list[str]) -> bool:
     miss = [c for c in cols if c not in df.columns]
     if miss:
@@ -607,12 +631,13 @@ try:
 except Exception as e:
     st.error(f"Failed to load models: {e}"); st.stop()
 
-# Defaults
-FEATURES = FEATURES_DEFAULT[:]
+# ---------------------------
+# Resolve FEATURES & targets
+# ---------------------------
+FEATURES = FEATURES_DEFAULT[:]          # fallback
 TARGET_BO = TARGET_BO_DEFAULT
 TARGET_BD = TARGET_BD_DEFAULT
 
-# Meta overrides
 def _load_meta(p: Path):
     if not p.exists(): return {}
     try:
@@ -623,19 +648,33 @@ def _load_meta(p: Path):
 meta_bo = _load_meta(BO_META_PATH)
 meta_bd = _load_meta(BD_META_PATH)
 
-# Use BO meta as primary feature source (or BD if BO missing metas)
-if meta_bo.get("features"): FEATURES = meta_bo["features"]
-elif meta_bd.get("features"): FEATURES = meta_bd["features"]
+# Prefer the model's trained feature names (golden source)
+try:
+    if hasattr(model_bo, "feature_names_in_"):
+        FEATURES = [str(x) for x in model_bo.feature_names_in_]
+except Exception:
+    pass
+
+# If still not set, fall back to BO meta, then BD meta, then defaults
+if not FEATURES and meta_bo.get("features"):
+    FEATURES = [str(x) for x in meta_bo["features"]]
+elif not FEATURES and meta_bd.get("features"):
+    FEATURES = [str(x) for x in meta_bd["features"]]
 
-if meta_bo.get("target"): TARGET_BO = meta_bo["target"]
-if meta_bd.get("target"): TARGET_BD = meta_bd["target"]
+# Targets (from metas if provided)
+if meta_bo.get("target"): TARGET_BO = str(meta_bo["target"])
+if meta_bd.get("target"): TARGET_BD = str(meta_bd["target"])
+
+# Canonical Depth name the model expects (if any)
+CANON_DEPTH = next((c for c in FEATURES if str(c).strip().lower().startswith("depth")), None)
 
 # Session constants for easy access elsewhere
 st.session_state["FEATURES"]  = FEATURES
 st.session_state["TARGET_BO"] = TARGET_BO
 st.session_state["TARGET_BD"] = TARGET_BD
+st.session_state["CANON_DEPTH"] = CANON_DEPTH
 
-# Optional: strict version banner
+# Optional strict version banner
 if STRICT_VERSION_CHECK:
     import numpy as _np, sklearn as _skl
     msgs=[]
@@ -744,27 +783,24 @@ if st.session_state.app_step == "dev":
             st.markdown('<div class="st-message-box st-error">Workbook must include Train/Training/training2 and Test/Testing/testing2 sheets.</div>', unsafe_allow_html=True)
             st.stop()
 
-        # Normalize + ensure cols
-        tr = _normalize_columns(book[sh_train].copy(), FEATURES, TARGET_BO, TARGET_BD)
-        te = _normalize_columns(book[sh_test].copy(),  FEATURES, TARGET_BO, TARGET_BD)
-        need = FEATURES + [TARGET_BO, TARGET_BD]
-        if not (ensure_cols(tr, need) and ensure_cols(te, need)):
-            st.markdown('<div class="st-message-box st-error">Missing required columns.</div>', unsafe_allow_html=True); st.stop()
+        # Normalize + ensure cols (robust depth handling)
+        tr = _prepare_table(book[sh_train].copy(), "Training", st.session_state["FEATURES"], st.session_state["TARGET_BO"], st.session_state["TARGET_BD"], st.session_state["CANON_DEPTH"])
+        te = _prepare_table(book[sh_test].copy(),  "Testing",  st.session_state["FEATURES"], st.session_state["TARGET_BO"], st.session_state["TARGET_BD"], st.session_state["CANON_DEPTH"])
 
         # Predict with exact training feature order
-        tr[PRED_BO] = model_bo.predict(_make_X(tr, FEATURES))
-        tr[PRED_BD] = model_bd.predict(_make_X(tr, FEATURES))
-        te[PRED_BO] = model_bo.predict(_make_X(te, FEATURES))
-        te[PRED_BD] = model_bd.predict(_make_X(te, FEATURES))
+        tr[PRED_BO] = model_bo.predict(_make_X(tr, st.session_state["FEATURES"]))
+        tr[PRED_BD] = model_bd.predict(_make_X(tr, st.session_state["FEATURES"]))
+        te[PRED_BO] = model_bo.predict(_make_X(te, st.session_state["FEATURES"]))
+        te[PRED_BD] = model_bd.predict(_make_X(te, st.session_state["FEATURES"]))
 
         st.session_state.results["Train"]=tr; st.session_state.results["Test"]=te
-        st.session_state.results["m_train_bo"]={"R": pearson_r(tr[TARGET_BO], tr[PRED_BO]), "RMSE": rmse(tr[TARGET_BO], tr[PRED_BO]), "MAE": mean_absolute_error(tr[TARGET_BO], tr[PRED_BO])}
-        st.session_state.results["m_train_bd"]={"R": pearson_r(tr[TARGET_BD], tr[PRED_BD]), "RMSE": rmse(tr[TARGET_BD], tr[PRED_BD]), "MAE": mean_absolute_error(tr[TARGET_BD], tr[PRED_BD])}
-        st.session_state.results["m_test_bo"] ={"R": pearson_r(te[TARGET_BO], te[PRED_BO]), "RMSE": rmse(te[TARGET_BO], te[PRED_BO]), "MAE": mean_absolute_error(te[TARGET_BO], te[PRED_BO])}
-        st.session_state.results["m_test_bd"] ={"R": pearson_r(te[TARGET_BD], te[PRED_BD]), "RMSE": rmse(te[TARGET_BD], te[PRED_BD]), "MAE": mean_absolute_error(te[TARGET_BD], te[PRED_BD])}
+        st.session_state.results["m_train_bo"]={"R": pearson_r(tr[st.session_state["TARGET_BO"]], tr[PRED_BO]), "RMSE": rmse(tr[st.session_state["TARGET_BO"]], tr[PRED_BO]), "MAE": mean_absolute_error(tr[st.session_state["TARGET_BO"]], tr[PRED_BO])}
+        st.session_state.results["m_train_bd"]={"R": pearson_r(tr[st.session_state["TARGET_BD"]], tr[PRED_BD]), "RMSE": rmse(tr[st.session_state["TARGET_BD"]], tr[PRED_BD]), "MAE": mean_absolute_error(tr[st.session_state["TARGET_BD"]], tr[PRED_BD])}
+        st.session_state.results["m_test_bo"] ={"R": pearson_r(te[st.session_state["TARGET_BO"]], te[PRED_BO]), "RMSE": rmse(te[st.session_state["TARGET_BO"]], te[PRED_BO]), "MAE": mean_absolute_error(te[st.session_state["TARGET_BO"]], te[PRED_BO])}
+        st.session_state.results["m_test_bd"] ={"R": pearson_r(te[st.session_state["TARGET_BD"]], te[PRED_BD]), "RMSE": rmse(te[st.session_state["TARGET_BD"]], te[PRED_BD]), "MAE": mean_absolute_error(te[st.session_state["TARGET_BD"]], te[PRED_BD])}
 
-        tr_min = tr[FEATURES].min().to_dict(); tr_max = tr[FEATURES].max().to_dict()
-        st.session_state.train_ranges = {f:(float(tr_min[f]), float(tr_max[f])) for f in FEATURES}
+        tr_min = tr[st.session_state["FEATURES"]].min().to_dict(); tr_max = tr[st.session_state["FEATURES"]].max().to_dict()
+        st.session_state.train_ranges = {f:(float(tr_min[f]), float(tr_max[f])) for f in st.session_state["FEATURES"]}
         st.markdown('<div class="st-message-box st-success">Case has been built and results are displayed below.</div>', unsafe_allow_html=True)
 
     def _metrics_block(lbl, m):
@@ -782,11 +818,11 @@ if st.session_state.app_step == "dev":
         """, unsafe_allow_html=True)
         t1, t2, t3 = st.tabs(["Breakout", "Breakdown", "Combined"])
         with t1:
-            st.plotly_chart(track_plot_single(df, PRED_BO, actual_col=TARGET_BO, title_suffix="Breakout"), use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
-            st.pyplot(cross_plot_static(df[TARGET_BO], df[PRED_BO], xlabel=f"Actual {TARGET_BO}", ylabel=f"Predicted {TARGET_BO}", color=COLORS["pred_bo"]), use_container_width=False)
+            st.plotly_chart(track_plot_single(df, PRED_BO, actual_col=st.session_state["TARGET_BO"], title_suffix="Breakout"), use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
+            st.pyplot(cross_plot_static(df[st.session_state["TARGET_BO"]], df[PRED_BO], xlabel=f"Actual {st.session_state['TARGET_BO']}", ylabel=f"Predicted {st.session_state['TARGET_BO']}", color=COLORS["pred_bo"]), use_container_width=False)
         with t2:
-            st.plotly_chart(track_plot_single(df, PRED_BD, actual_col=TARGET_BD, title_suffix="Breakdown"), use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
-            st.pyplot(cross_plot_static(df[TARGET_BD], df[PRED_BD], xlabel=f"Actual {TARGET_BD}", ylabel=f"Predicted {TARGET_BD}", color=COLORS["pred_bd"]), use_container_width=False)
+            st.plotly_chart(track_plot_single(df, PRED_BD, actual_col=st.session_state["TARGET_BD"], title_suffix="Breakdown"), use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
+            st.pyplot(cross_plot_static(df[st.session_state["TARGET_BD"]], df[PRED_BD], xlabel=f"Actual {st.session_state['TARGET_BD']}", ylabel=f"Predicted {st.session_state['TARGET_BD']}", color=COLORS["pred_bd"]), use_container_width=False)
         with t3:
             st.plotly_chart(track_plot_combined(df), use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
 
@@ -822,29 +858,26 @@ if st.session_state.app_step == "validate":
     if go_btn and up is not None:
         book = read_book_bytes(up.getvalue())
         name = find_sheet(book, ["Validation","Validate","validation2","Val","val"]) or list(book.keys())[0]
-        df = _normalize_columns(book[name].copy(), FEATURES, TARGET_BO, TARGET_BD)
-        need = FEATURES + [TARGET_BO, TARGET_BD]
-        if not ensure_cols(df, need):
-            st.markdown('<div class="st-message-box st-error">Missing required columns.</div>', unsafe_allow_html=True); st.stop()
+        df = _prepare_table(book[name].copy(), "Validation", st.session_state["FEATURES"], st.session_state["TARGET_BO"], st.session_state["TARGET_BD"], st.session_state["CANON_DEPTH"])
 
-        df[PRED_BO] = model_bo.predict(_make_X(df, FEATURES))
-        df[PRED_BD] = model_bd.predict(_make_X(df, FEATURES))
+        df[PRED_BO] = model_bo.predict(_make_X(df, st.session_state["FEATURES"]))
+        df[PRED_BD] = model_bd.predict(_make_X(df, st.session_state["FEATURES"]))
         st.session_state.results["Validate"]=df
 
         ranges = st.session_state.train_ranges; oor_pct = 0.0; tbl=None
         if ranges:
-            any_viol = pd.DataFrame({f:(df[f]<ranges[f][0])|(df[f]>ranges[f][1]) for f in FEATURES}).any(axis=1)
+            any_viol = pd.DataFrame({f:(df[f]<ranges[f][0])|(df[f]>ranges[f][1]) for f in st.session_state["FEATURES"]}).any(axis=1)
             oor_pct = float(any_viol.mean()*100.0)
             if any_viol.any():
-                tbl = df.loc[any_viol, FEATURES].copy()
-                for c in FEATURES:
+                tbl = df.loc[any_viol, st.session_state["FEATURES"]].copy()
+                for c in st.session_state["FEATURES"]:
                     if pd.api.types.is_numeric_dtype(tbl[c]): tbl[c] = tbl[c].round(2)
-                tbl["Violations"] = pd.DataFrame({f:(df[f]<ranges[f][0])|(df[f]>ranges[f][1]) for f in FEATURES}).loc[any_viol].apply(
+                tbl["Violations"] = pd.DataFrame({f:(df[f]<ranges[f][0])|(df[f]>ranges[f][1]) for f in st.session_state["FEATURES"]}).loc[any_viol].apply(
                     lambda r:", ".join([c for c,v in r.items() if v]), axis=1
                 )
 
-        st.session_state.results["m_val_bo"]={"R": pearson_r(df[TARGET_BO], df[PRED_BO]), "RMSE": rmse(df[TARGET_BO], df[PRED_BO]), "MAE": mean_absolute_error(df[TARGET_BO], df[PRED_BO])}
-        st.session_state.results["m_val_bd"]={"R": pearson_r(df[TARGET_BD], df[PRED_BD]), "RMSE": rmse(df[TARGET_BD], df[PRED_BD]), "MAE": mean_absolute_error(df[TARGET_BD], df[PRED_BD])}
+        st.session_state.results["m_val_bo"]={"R": pearson_r(df[st.session_state["TARGET_BO"]], df[PRED_BO]), "RMSE": rmse(df[st.session_state["TARGET_BO"]], df[PRED_BO]), "MAE": mean_absolute_error(df[st.session_state["TARGET_BO"]], df[PRED_BO])}
+        st.session_state.results["m_val_bd"]={"R": pearson_r(df[st.session_state["TARGET_BD"]], df[PRED_BD]), "RMSE": rmse(df[st.session_state["TARGET_BD"]], df[PRED_BD]), "MAE": mean_absolute_error(df[st.session_state["TARGET_BD"]], df[PRED_BD])}
         st.session_state.results["sv_val"]={"n":len(df), "bo_min":float(df[PRED_BO].min()), "bo_max":float(df[PRED_BO].max()),
                                             "bd_min":float(df[PRED_BD].min()), "bd_max":float(df[PRED_BD].max()), "oor":oor_pct}
         st.session_state.results["oor_tbl"]=tbl
@@ -860,14 +893,14 @@ if st.session_state.app_step == "validate":
 
         t1, t2, t3 = st.tabs(["Breakout", "Breakdown", "Combined"])
         with t1:
-            st.plotly_chart(track_plot_single(df, PRED_BO, actual_col=TARGET_BO, title_suffix="Breakout"),
+            st.plotly_chart(track_plot_single(df, PRED_BO, actual_col=st.session_state["TARGET_BO"], title_suffix="Breakout"),
                             use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
-            st.pyplot(cross_plot_static(df[TARGET_BO], df[PRED_BO], f"Actual {TARGET_BO}", f"Predicted {TARGET_BO}", COLORS["pred_bo"]),
+            st.pyplot(cross_plot_static(df[st.session_state["TARGET_BO"]], df[PRED_BO], f"Actual {st.session_state['TARGET_BO']}", f"Predicted {st.session_state['TARGET_BO']}", COLORS["pred_bo"]),
                       use_container_width=False)
         with t2:
-            st.plotly_chart(track_plot_single(df, PRED_BD, actual_col=TARGET_BD, title_suffix="Breakdown"),
+            st.plotly_chart(track_plot_single(df, PRED_BD, actual_col=st.session_state["TARGET_BD"], title_suffix="Breakdown"),
                             use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
-            st.pyplot(cross_plot_static(df[TARGET_BD], df[PRED_BD], f"Actual {TARGET_BD}", f"Predicted {TARGET_BD}", COLORS["pred_bd"]),
+            st.pyplot(cross_plot_static(df[st.session_state["TARGET_BD"]], df[PRED_BD], f"Actual {st.session_state['TARGET_BD']}", f"Predicted {st.session_state['TARGET_BD']}", COLORS["pred_bd"]),
                       use_container_width=False)
         with t3:
             st.plotly_chart(track_plot_combined(df), use_container_width=False, config={"displayModeBar": False, "scrollZoom": True})
@@ -899,16 +932,15 @@ if st.session_state.app_step == "predict":
 
     if go_btn and up is not None:
         book = read_book_bytes(up.getvalue()); name = list(book.keys())[0]
-        df = _normalize_columns(book[name].copy(), FEATURES, TARGET_BO, TARGET_BD)
-        if not ensure_cols(df, FEATURES):
-            st.markdown('<div class="st-message-box st-error">Missing required columns.</div>', unsafe_allow_html=True); st.stop()
-        df[PRED_BO] = model_bo.predict(_make_X(df, FEATURES))
-        df[PRED_BD] = model_bd.predict(_make_X(df, FEATURES))
+        df = _prepare_table(book[name].copy(), "Prediction", st.session_state["FEATURES"], st.session_state["TARGET_BO"], st.session_state["TARGET_BD"], st.session_state["CANON_DEPTH"])
+
+        df[PRED_BO] = model_bo.predict(_make_X(df, st.session_state["FEATURES"]))
+        df[PRED_BD] = model_bd.predict(_make_X(df, st.session_state["FEATURES"]))
         st.session_state.results["PredictOnly"]=df
 
         ranges = st.session_state.train_ranges; oor_pct = 0.0
         if ranges:
-            any_viol = pd.DataFrame({f:(df[f]<ranges[f][0])|(df[f]>ranges[f][1]) for f in FEATURES}).any(axis=1)
+            any_viol = pd.DataFrame({f:(df[f]<ranges[f][0])|(df[f]>ranges[f][1]) for f in st.session_state["FEATURES"]}).any(axis=1)
             oor_pct = float(any_viol.mean()*100.0)
         st.session_state.results["sv_pred"]={
             "n":len(df),
@@ -962,12 +994,12 @@ if st.session_state.show_preview_modal:
             tabs = st.tabs(names)
             for t, name in zip(tabs, names):
                 with t:
-                    df = _normalize_columns(book_to_preview[name], FEATURES, TARGET_BO, TARGET_BD)
+                    df = _prepare_table(book_to_preview[name], "Preview", st.session_state["FEATURES"], st.session_state["TARGET_BO"], st.session_state["TARGET_BD"], st.session_state["CANON_DEPTH"])
                     t1, t2 = st.tabs(["Tracks", "Summary"])
                     with t1:
-                        st.pyplot(preview_tracks(df, FEATURES), use_container_width=True)
+                        st.pyplot(preview_tracks(df, st.session_state["FEATURES"]), use_container_width=True)
                     with t2:
-                        feat_present = [c for c in FEATURES if c in df.columns]
+                        feat_present = [c for c in st.session_state["FEATURES"] if c in df.columns]
                         if not feat_present:
                             st.info("No feature columns found to summarize.")
                         else: