Update app.py

app.py

@@ -2198,9 +2198,7 @@ with tab_predict:
     # =========================
     pipe = st.session_state.pipe
     meta = st.session_state.meta
-    feature_cols = meta["schema"]["features"]
-    num_cols = meta["schema"]["numeric"]
-    cat_cols = meta["schema"]["categorical"]
+    
 
     
 
@@ -2426,300 +2424,282 @@ with tab_predict:
             values_by_index[i] = np.nan if v.strip() == "" else v
 
                 
-        # Apply FISH/NGS selections to row
-        fish_set = set(fish_selected)
-        ngs_set = set(ngs_selected)
+    # Apply FISH/NGS selections to row
+    fish_set = set(fish_selected)
+    ngs_set = set(ngs_selected)
+    
+    for i, f in enumerate(feature_cols):
+        if f in FISH_MARKERS:
+            values_by_index[i] = 1 if f in fish_set else 0
+        if f in NGS_MARKERS:
+            values_by_index[i] = 1 if f in ngs_set else 0
         
-        for i, f in enumerate(feature_cols):
-            if f in FISH_MARKERS:
-                values_by_index[i] = 1 if f in fish_set else 0
-            if f in NGS_MARKERS:
-                values_by_index[i] = 1 if f in ngs_set else 0
-        
-        # Auto-fill marker counts if present
-        if FISH_COUNT_COL in feature_cols:
-            values_by_index[feature_cols.index(FISH_COUNT_COL)] = int(len(fish_selected))
-        if NGS_COUNT_COL in feature_cols:
-            values_by_index[feature_cols.index(NGS_COUNT_COL)] = int(len(ngs_selected))
-        st.divider()
-        st.subheader("Predict single patient")
+    # Auto-fill marker counts if present
+    if FISH_COUNT_COL in feature_cols:
+        values_by_index[feature_cols.index(FISH_COUNT_COL)] = int(len(fish_selected))
+    if NGS_COUNT_COL in feature_cols:
+        values_by_index[feature_cols.index(NGS_COUNT_COL)] = int(len(ngs_selected))
+    st.divider()
+    st.subheader("Predict single patient")
 
-        m = meta.get("metrics", {})
-        default_thr = float(m.get("best_threshold", 0.5))
-        
-        thr_single = st.slider(
-            "Classification threshold",
-            0.0, 1.0, default_thr, 0.01,
-            key="sp_thr"
-        )
-        
-        # External validation threshold
-        thr_ext = st.slider(
-            "External validation threshold",
-            0.0, 1.0, default_thr, 0.01,
-            key="thr_ext"
-        )
+    m = meta.get("metrics", {})
+    default_thr = float(m.get("best_threshold", 0.5))
+    
+    thr_single = st.slider(
+        "Classification threshold",
+        0.0, 1.0, default_thr, 0.01,
+        key="sp_thr"
+    )
+    
+    # External validation threshold
+    thr_ext = st.slider(
+        "External validation threshold",
+        0.0, 1.0, default_thr, 0.01,
+        key="thr_ext"
+    )
 
-        low_cut_s, high_cut_s = st.slider(
-            "Risk band cutoffs (low, high)",
-            0.0, 1.0, (0.2, 0.8), 0.01,
-            key="sp_risk_cuts"
-        )
-        
-        # Ensure low <= high
-        if low_cut_s > high_cut_s:
-            low_cut_s, high_cut_s = high_cut_s, low_cut_s
+    low_cut_s, high_cut_s = st.slider(
+        "Risk band cutoffs (low, high)",
+        0.0, 1.0, (0.2, 0.8), 0.01,
+        key="sp_risk_cuts"
+    )
+    
+    # Ensure low <= high
+    if low_cut_s > high_cut_s:
+        low_cut_s, high_cut_s = high_cut_s, low_cut_s
 
 
-        def band_one(p: float) -> str:
-            if p < low_cut_s:
-                return "Low"
-            if p >= high_cut_s:
-                return "High"
-            return "Intermediate"
-        # Submit button (no form needed; simpler + fewer state surprises)
-        if st.button("Predict single patient", key="sp_predict_btn"):
-            X_one = pd.DataFrame([values_by_index], columns=feature_cols).replace({pd.NA: np.nan})
+    def band_one(p: float) -> str:
+        if p < low_cut_s:
+            return "Low"
+        if p >= high_cut_s:
+            return "High"
+        return "Intermediate"
+    # Submit button (no form needed; simpler + fewer state surprises)
+    if st.button("Predict single patient", key="sp_predict_btn"):
+        X_one = pd.DataFrame([values_by_index], columns=feature_cols).replace({pd.NA: np.nan})
+    
+        for c in num_cols:
+            if c in X_one.columns:
+                X_one[c] = pd.to_numeric(X_one[c], errors="coerce")
+    
+        for c in cat_cols:
+            if c in X_one.columns:
+                X_one[c] = X_one[c].astype("object")
+                X_one.loc[X_one[c].isna(), c] = np.nan
+                X_one[c] = X_one[c].map(lambda v: v if pd.isna(v) else str(v))
+
+    
+        proba_one = float(pipe.predict_proba(X_one)[:, 1][0])
+        st.success("Prediction generated.")
+        st.metric("Predicted probability", f"{proba_one:.4f}") 
+   
+
+        # ---- Survival prediction for this patient (if survival model loaded) ----
+        # ---- Survival prediction for this patient (if survival model loaded) ----
+        bundle = st.session_state.get("surv_model", None)
         
-            for c in num_cols:
-                if c in X_one.columns:
-                    X_one[c] = pd.to_numeric(X_one[c], errors="coerce")
+        if isinstance(bundle, dict) and bundle.get("model") is not None:
+            try:
+                cph = bundle["model"]
+                surv_cols = bundle.get("columns", [])
         
-            for c in cat_cols:
-                if c in X_one.columns:
-                    X_one[c] = X_one[c].astype("object")
-                    X_one.loc[X_one[c].isna(), c] = np.nan
-                    X_one[c] = X_one[c].map(lambda v: v if pd.isna(v) else str(v))
-
+                # Build Cox input row (same preprocessing as Cox training)
+                df_one_surv = X_one[feature_cols].copy()
         
-            proba_one = float(pipe.predict_proba(X_one)[:, 1][0])
-            st.success("Prediction generated.")
-            st.metric("Predicted probability", f"{proba_one:.4f}") 
-       
-
-            # ---- Survival prediction for this patient (if survival model loaded) ----
-            cph = st.session_state.get("surv_model", None)
-            
-            if cph is not None:
-                try:
-                    # Build Cox input row with same preprocessing used in training Cox:
-                    # predictors + one-hot categoricals
-                    df_one_surv = X_one[feature_cols].copy()
-                    for c in num_cols:
+                for c in num_cols:
+                    if c in df_one_surv.columns:
                         df_one_surv[c] = pd.to_numeric(df_one_surv[c], errors="coerce")
-                    for c in cat_cols:
-                        df_one_surv[c] = df_one_surv[c].astype("object").map(lambda v: v if pd.isna(v) else str(v))
-            
-                    df_one_surv_oh = pd.get_dummies(df_one_surv, columns=cat_cols, drop_first=True)
-            
-                    # Align columns to Cox model training columns
-                    bundle = st.session_state.get("surv_model", None)
-                    if isinstance(bundle, dict) and "model" in bundle:
-                        cph = bundle["model"]
-                        surv_cols = bundle.get("columns", [])
-                    else:
-                        cph = None
-                        surv_cols = []
-                    
-                    if cph is not None:
-                        try:
-                            df_one_surv = X_one[feature_cols].copy()
-                    
-                            for c in num_cols:
-                                if c in df_one_surv.columns:
-                                    df_one_surv[c] = pd.to_numeric(df_one_surv[c], errors="coerce")
-                            for c in cat_cols:
-                                if c in df_one_surv.columns:
-                                    df_one_surv[c] = df_one_surv[c].astype("object").map(lambda v: v if pd.isna(v) else str(v))
-                    
-                            df_one_surv_oh = pd.get_dummies(df_one_surv, columns=cat_cols, drop_first=True)
-                    
-                            # Align to training predictor columns
-                            for col in surv_cols:
-                                if col not in df_one_surv_oh.columns:
-                                    df_one_surv_oh[col] = 0
-                            df_one_surv_oh = df_one_surv_oh[surv_cols]
-                    
-                            surv_fn = cph.predict_survival_function(df_one_surv_oh)
-                            ...
-                        except Exception as e:
-                            st.warning(f"Survival prediction could not be computed: {e}")
+        
+                for c in cat_cols:
+                    if c in df_one_surv.columns:
+                        df_one_surv[c] = df_one_surv[c].astype("object").map(
+                            lambda v: v if pd.isna(v) else str(v)
+                        )
+        
+                df_one_surv_oh = pd.get_dummies(df_one_surv, columns=cat_cols, drop_first=True)
+        
+                # Align to training predictor columns
+                for col in surv_cols:
+                    if col not in df_one_surv_oh.columns:
+                        df_one_surv_oh[col] = 0
+                df_one_surv_oh = df_one_surv_oh[surv_cols]
+        
+                # Predict survival function
+                surv_fn = cph.predict_survival_function(df_one_surv_oh)
+        
+                def surv_at(days: int) -> float:
+                    idx = surv_fn.index.values
+                    j = int(np.argmin(np.abs(idx - days)))
+                    return float(surv_fn.iloc[j, 0])
+        
+                s6m = surv_at(180)
+                s1y = surv_at(365)
+                s2y = surv_at(730)
+                s3y = surv_at(1095)
+        
+                st.subheader("Predicted survival probability")
+                a, b, c, d = st.columns(4)
+                a.metric("6 months", f"{s6m*100:.1f}%")
+                b.metric("1 year",   f"{s1y*100:.1f}%")
+                c.metric("2 years",  f"{s2y*100:.1f}%")
+                d.metric("3 years",  f"{s3y*100:.1f}%")
+        
+                fig, ax = make_fig(figsize=FIGSIZE, dpi=plot_dpi_screen)
+                ax.plot(surv_fn.index, surv_fn.values)
+                ax.set_xlabel("Days from diagnosis")
+                ax.set_ylabel("Survival probability")
+                ax.set_ylim(0, 1)
+                ax.set_title("Predicted survival curve (patient)")
+        
+                render_plot_with_download(
+                    fig,
+                    title="Patient survival curve",
+                    filename="patient_survival_curve.png",
+                    export_dpi=export_dpi,
+                    key="dl_patient_surv_curve"
+                )
+        
+            except Exception as e:
+                st.warning(f"Survival prediction could not be computed: {e}")
+        
+        else:
+            st.info("Survival model not loaded/published yet (survival bundle missing).")
 
-            
-                    surv_fn = cph.predict_survival_function(df_one_surv_oh)
 
-                    def surv_at(days: int) -> float:
-                        # find nearest time index
-                        idx = surv_fn.index.values
-                        j = int(np.argmin(np.abs(idx - days)))
-                        return float(surv_fn.iloc[j, 0])
-            
-                    s6m = surv_at(180)
-                    s1y = surv_at(365)
-                    s2y = surv_at(730)
-                    s3y = surv_at(1095)
-            
-                    st.subheader("Predicted survival probability")
-                    a,b,c,d = st.columns(4)
-                    a.metric("6 months", f"{s6m*100:.1f}%")
-                    b.metric("1 year",   f"{s1y*100:.1f}%")
-                    c.metric("2 years",  f"{s2y*100:.1f}%")
-                    d.metric("3 years",  f"{s3y*100:.1f}%")
-            
-                    # Optional: plot curve
-                    fig, ax = make_fig(figsize=FIGSIZE, dpi=plot_dpi_screen)
-                    ax.plot(surv_fn.index, surv_fn.values)
-                    ax.set_xlabel("Days from diagnosis")
-                    ax.set_ylabel("Survival probability")
-                    ax.set_ylim(0, 1)
-                    ax.set_title("Predicted survival curve (patient)")
-            
-                    render_plot_with_download(
-                        fig,
-                        title="Patient survival curve",
-                        filename="patient_survival_curve.png",
-                        export_dpi=export_dpi,
-                        key="dl_patient_surv_curve"
-                    )
+        
+        out = X_one.copy()
+        out["predicted_probability"] = proba_one
+        pred_class = int(proba_one >= thr_single)
+        out["predicted_class"] = pred_class
+        
+        out["risk_band"] = band_one(proba_one)
 
-                except Exception as e:
-                    st.warning(f"Survival prediction could not be computed: {e}")
-            else:
-                st.info("Survival model not loaded/published yet (survival_model.joblib missing).")
+        
 
-            
-            out = X_one.copy()
-            out["predicted_probability"] = proba_one
-            pred_class = int(proba_one >= thr_single)
-            out["predicted_class"] = pred_class
-            
-            out["risk_band"] = band_one(proba_one)
+        # ---- SHAP compute only (cache) ----
+        X_one_t = transform_before_clf(pipe, X_one)
+        
+        explainer = st.session_state.get("explainer")
+        explainer_sig = st.session_state.get("explainer_sig")
+        
+        current_sig = (
+            selected,
+            None if st.session_state.get("X_bg_for_shap") is None else int(len(st.session_state["X_bg_for_shap"]))
+        )
+        
+        if explainer is None or explainer_sig != current_sig:
+            X_bg = st.session_state.get("X_bg_for_shap")
+            if X_bg is None:
+                st.error("SHAP background not available. Admin must publish latest/background.csv.")
+                st.stop()
+        
+            st.session_state.explainer = build_shap_explainer(pipe, X_bg)
+            st.session_state.explainer_sig = current_sig
+            explainer = st.session_state.explainer
+        
+        shap_vals = explainer.shap_values(X_one_t)
+        if isinstance(shap_vals, list):
+            shap_vals = shap_vals[1]
+        
+        names = get_final_feature_names(pipe)
+        
+        try:
+            x_dense = X_one_t.toarray()[0]
+        except Exception:
+            x_dense = np.array(X_one_t)[0]
+        
+        base = explainer.expected_value
+        if not np.isscalar(base):
+            base = float(np.array(base).reshape(-1)[0])
+        
+        exp = shap.Explanation(
+            values=shap_vals[0],
+            base_values=float(base),
+            data=x_dense,
+            feature_names=names,
+        )
+        
+        # CACHE ONLY
+        st.session_state.shap_single_exp = exp
 
-            
+    
+        st.dataframe(out, use_container_width=True)
+    
+        st.download_button(
+            "Download single patient result (CSV)",
+            out.to_csv(index=False).encode("utf-8"),
+            file_name="single_patient_prediction.csv",
+            mime="text/csv",
+            key="dl_sp_csv",
+        )
 
-            # ---- SHAP compute only (cache) ----
-            X_one_t = transform_before_clf(pipe, X_one)
-            
-            explainer = st.session_state.get("explainer")
-            explainer_sig = st.session_state.get("explainer_sig")
-            
-            current_sig = (
-                selected,
-                None if st.session_state.get("X_bg_for_shap") is None else int(len(st.session_state["X_bg_for_shap"]))
-            )
-            
-            if explainer is None or explainer_sig != current_sig:
-                X_bg = st.session_state.get("X_bg_for_shap")
-                if X_bg is None:
-                    st.error("SHAP background not available. Admin must publish latest/background.csv.")
-                    st.stop()
-            
-                st.session_state.explainer = build_shap_explainer(pipe, X_bg)
-                st.session_state.explainer_sig = current_sig
-                explainer = st.session_state.explainer
-            
-            shap_vals = explainer.shap_values(X_one_t)
-            if isinstance(shap_vals, list):
-                shap_vals = shap_vals[1]
-            
-            names = get_final_feature_names(pipe)
-            
-            try:
-                x_dense = X_one_t.toarray()[0]
-            except Exception:
-                x_dense = np.array(X_one_t)[0]
-            
-            base = explainer.expected_value
-            if not np.isscalar(base):
-                base = float(np.array(base).reshape(-1)[0])
-            
-            exp = shap.Explanation(
-                values=shap_vals[0],
-                base_values=float(base),
-                data=x_dense,
-                feature_names=names,
+        # ---- Always render cached SHAP ----
+        if "shap_single_exp" in st.session_state:
+            exp = st.session_state.shap_single_exp
+        
+            max_display_single = st.slider(
+                "Top features to display (single patient)",
+                5, 40, 20, 1,
+                key="sp_single_max_display"
             )
-            
-            # CACHE ONLY
-            st.session_state.shap_single_exp = exp
-
         
-            st.dataframe(out, use_container_width=True)
+            c1, c2 = st.columns(2)
         
-            st.download_button(
-                "Download single patient result (CSV)",
-                out.to_csv(index=False).encode("utf-8"),
-                file_name="single_patient_prediction.csv",
-                mime="text/csv",
-                key="dl_sp_csv",
-            )
+            with c1:
+                plt.figure(figsize=FIGSIZE, dpi=plot_dpi_screen)
+                shap.plots.waterfall(exp, show=False, max_display=max_display_single)
+                fig_w = plt.gcf()
+                render_plot_with_download(
+                    fig_w,
+                    title="Single-patient SHAP waterfall",
+                    filename="single_patient_shap_waterfall.png",
+                    export_dpi=export_dpi,
+                    key="dl_sp_wf"
+                )
+        
+            with c2:
+                plt.figure(figsize=FIGSIZE, dpi=plot_dpi_screen)
+                shap.plots.bar(exp, show=False, max_display=max_display_single)
+                fig_b = plt.gcf()
+                render_plot_with_download(
+                    fig_b,
+                    title="Single-patient SHAP bar",
+                    filename="single_patient_shap_bar.png",
+                    export_dpi=export_dpi,
+                    key="dl_sp_bar"
+                )
 
-            # ---- Always render cached SHAP ----
-            if "shap_single_exp" in st.session_state:
-                exp = st.session_state.shap_single_exp
+            # --- After SHAP plots are displayed ---
+            with st.expander("How to interpret the SHAP explanation plots", expanded=True):
+                st.markdown(r"""
+            **What is SHAP?**  
+            SHAP (SHapley Additive exPlanations) decomposes a model prediction into **feature-wise contributions**.  
+            Each feature pushes the prediction **toward higher risk** or **toward lower risk**, relative to the model’s baseline.
             
-                max_display_single = st.slider(
-                    "Top features to display (single patient)",
-                    5, 40, 20, 1,
-                    key="sp_single_max_display"
-                )
+            **What is \(E[f(X)]\)? (baseline / expected model output)**  
+            - \(E[f(X)]\) is the model’s **average output** over the reference population used by the explainer (typically the training set).  
+            - It is the **starting point** of the explanation before any patient-specific information is applied.
             
-                c1, c2 = st.columns(2)
+            **What is \(f(x)\)? (this patient’s model output)**  
+            - \(f(x)\) is the model’s **final output for this patient**, after adding all feature contributions to the baseline.  
+            - For many classifiers (including logistic regression), SHAP waterfall plots often use the **log-odds (logit) scale** rather than raw probability.
             
-                with c1:
-                    plt.figure(figsize=FIGSIZE, dpi=plot_dpi_screen)
-                    shap.plots.waterfall(exp, show=False, max_display=max_display_single)
-                    fig_w = plt.gcf()
-                    render_plot_with_download(
-                        fig_w,
-                        title="Single-patient SHAP waterfall",
-                        filename="single_patient_shap_waterfall.png",
-                        export_dpi=export_dpi,
-                        key="dl_sp_wf"
-                    )
+            **How the waterfall plot should be read**  
+            - The plot starts at **\(E[f(X)]\)** and then adds/subtracts feature effects to arrive at **\(f(x)\)**.  
+            - **Red bars** push the prediction **toward higher risk** (increase the model output).  
+            - **Blue bars** push the prediction **toward lower risk / protective direction** (decrease the model output).  
+            - The **bar length** indicates the **strength** of that feature’s influence for this patient.  
+            - “**Other features**” is the combined net effect of features not shown individually.
             
-                with c2:
-                    plt.figure(figsize=FIGSIZE, dpi=plot_dpi_screen)
-                    shap.plots.bar(exp, show=False, max_display=max_display_single)
-                    fig_b = plt.gcf()
-                    render_plot_with_download(
-                        fig_b,
-                        title="Single-patient SHAP bar",
-                        filename="single_patient_shap_bar.png",
-                        export_dpi=export_dpi,
-                        key="dl_sp_bar"
-                    )
-
-                # --- After SHAP plots are displayed ---
-                with st.expander("How to interpret the SHAP explanation plots", expanded=True):
-                    st.markdown(r"""
-                **What is SHAP?**  
-                SHAP (SHapley Additive exPlanations) decomposes a model prediction into **feature-wise contributions**.  
-                Each feature pushes the prediction **toward higher risk** or **toward lower risk**, relative to the model’s baseline.
-                
-                **What is \(E[f(X)]\)? (baseline / expected model output)**  
-                - \(E[f(X)]\) is the model’s **average output** over the reference population used by the explainer (typically the training set).  
-                - It is the **starting point** of the explanation before any patient-specific information is applied.
-                
-                **What is \(f(x)\)? (this patient’s model output)**  
-                - \(f(x)\) is the model’s **final output for this patient**, after adding all feature contributions to the baseline.  
-                - For many classifiers (including logistic regression), SHAP waterfall plots often use the **log-odds (logit) scale** rather than raw probability.
-                
-                **How the waterfall plot should be read**  
-                - The plot starts at **\(E[f(X)]\)** and then adds/subtracts feature effects to arrive at **\(f(x)\)**.  
-                - **Red bars** push the prediction **toward higher risk** (increase the model output).  
-                - **Blue bars** push the prediction **toward lower risk / protective direction** (decrease the model output).  
-                - The **bar length** indicates the **strength** of that feature’s influence for this patient.  
-                - “**Other features**” is the combined net effect of features not shown individually.
-                
-                **How the bar plot should be read (Top contributors)**  
-                - Features are ranked by **absolute SHAP value** (largest impact first).  
-                - **Positive SHAP** (right) increases predicted risk; **negative SHAP** (left) decreases predicted risk.
-                
-                **Clinical cautions**  
-                - SHAP explains the model’s behavior; it does **not** prove causality.  
-                - Ensure variable definitions and patient population match the model’s intended use.
-                """)
+            **How the bar plot should be read (Top contributors)**  
+            - Features are ranked by **absolute SHAP value** (largest impact first).  
+            - **Positive SHAP** (right) increases predicted risk; **negative SHAP** (left) decreases predicted risk.
+            
+            **Clinical cautions**  
+            - SHAP explains the model’s behavior; it does **not** prove causality.  
+            - Ensure variable definitions and patient population match the model’s intended use.
+            """)
 
 
     # -----------------------------
@@ -2915,29 +2895,51 @@ with tab_predict:
                     cph = None
                     surv_cols = []
                 
-                if cph is not None:
+                # ---- Batch survival probabilities (if survival model loaded) ----
+                bundle = st.session_state.get("surv_model", None)
+                
+                if isinstance(bundle, dict) and bundle.get("model") is not None:
                     try:
-                        df_one_surv = X_one[feature_cols].copy()
+                        cph = bundle["model"]
+                        surv_cols = bundle.get("columns", [])
+                
+                        df_surv_in = X_inf[feature_cols].copy()
                 
                         for c in num_cols:
-                            if c in df_one_surv.columns:
-                                df_one_surv[c] = pd.to_numeric(df_one_surv[c], errors="coerce")
+                            if c in df_surv_in.columns:
+                                df_surv_in[c] = pd.to_numeric(df_surv_in[c], errors="coerce")
+                
                         for c in cat_cols:
-                            if c in df_one_surv.columns:
-                                df_one_surv[c] = df_one_surv[c].astype("object").map(lambda v: v if pd.isna(v) else str(v))
+                            if c in df_surv_in.columns:
+                                df_surv_in[c] = df_surv_in[c].astype("object").map(
+                                    lambda v: v if pd.isna(v) else str(v)
+                                )
                 
-                        df_one_surv_oh = pd.get_dummies(df_one_surv, columns=cat_cols, drop_first=True)
+                        df_surv_in_oh = pd.get_dummies(df_surv_in, columns=cat_cols, drop_first=True)
                 
                         # Align to training predictor columns
                         for col in surv_cols:
-                            if col not in df_one_surv_oh.columns:
-                                df_one_surv_oh[col] = 0
-                        df_one_surv_oh = df_one_surv_oh[surv_cols]
+                            if col not in df_surv_in_oh.columns:
+                                df_surv_in_oh[col] = 0
+                        df_surv_in_oh = df_surv_in_oh[surv_cols]
+                
+                        surv_fn_all = cph.predict_survival_function(df_surv_in_oh)
+                
+                        def surv_vec_at(days: int):
+                            idx = surv_fn_all.index.values
+                            j = int(np.argmin(np.abs(idx - days)))
+                            return surv_fn_all.iloc[j, :].values.astype(float)
+                
+                        df_out["survival_6m"] = surv_vec_at(180)
+                        df_out["survival_1y"] = surv_vec_at(365)
+                        df_out["survival_2y"] = surv_vec_at(730)
+                        df_out["survival_3y"] = surv_vec_at(1095)
                 
-                        surv_fn = cph.predict_survival_function(df_one_surv_oh)
-                        ...
                     except Exception as e:
-                        st.warning(f"Survival prediction could not be computed: {e}")
+                        st.warning(f"Batch survival probabilities could not be computed: {e}")
+                else:
+                    st.info("Survival model not loaded/published yet (survival bundle missing).")
+
 
         
                 surv_fn_all = cph.predict_survival_function(df_surv_in_oh)