Update app.py

app.py

@@ -6,11 +6,14 @@ from inference import (
     FEATURE_NAMES,
     REPORTING_OUTCOMES,
     OUTCOME_DESCRIPTIONS,
+    OUTCOMES,
+    SHAP_OUTCOMES,
     predict_with_comparison,
     create_all_shap_plots,
-    icon_array,
+    create_all_icon_arrays,
 )
 
+
 # ---------------------------------------------------------------------------
 # Choice lists
 # ---------------------------------------------------------------------------
@@ -52,6 +55,7 @@ SCATXRSN_CHOICES  = [
     "Hodgkin lymphoma",
 ]
 
+
 # ---------------------------------------------------------------------------
 # Grouped published-regimen dropdown
 # ---------------------------------------------------------------------------
@@ -132,6 +136,7 @@ PUBLISHED_PRESETS = {
     },
 }
 
+
 # ---------------------------------------------------------------------------
 # Feature groupings
 # ---------------------------------------------------------------------------
@@ -142,6 +147,7 @@ DONOR_FEATURES   = ["DONORF", "GRAFTYPE", "HLA_FINAL",
 DISEASE_FEATURES = ["NACS2YR", "EXCHTFPR", "VOC2YPR", "VOCFRQPR", "SCATXRSN"]
 ALL_FEATURES     = PATIENT_FEATURES + DONOR_FEATURES + DISEASE_FEATURES
 
+
 # ---------------------------------------------------------------------------
 # Utility callbacks
 # ---------------------------------------------------------------------------
@@ -173,6 +179,7 @@ def vocfrqpr_from_voc2ypr(voc_status):
 
 
 def apply_grouped_preset(selected_value):
+    # Header row clicked — reset dropdown, leave all fields unchanged
     if not selected_value or selected_value in HEADER_VALUES:
         return [gr.update(value=None)] + [gr.update()] * 6
 
@@ -181,7 +188,7 @@ def apply_grouped_preset(selected_value):
         return [gr.update()] * 7
 
     return [
-        gr.update(),
+        gr.update(),                             # leave dropdown showing selection
         gr.update(value=preset["DONORF"]),
         gr.update(value=preset["CONDGRPF"]),
         gr.update(value=preset["CONDGRP_FINAL"]),
@@ -190,6 +197,7 @@ def apply_grouped_preset(selected_value):
         gr.update(value=preset["HLA_FINAL"]),
     ]
 
+
 # ---------------------------------------------------------------------------
 # Component factory
 # ---------------------------------------------------------------------------
@@ -238,6 +246,7 @@ def make_component(name: str):
     else:
         return gr.Textbox(label=name)
 
+
 # ---------------------------------------------------------------------------
 # Prediction callback
 # ---------------------------------------------------------------------------
@@ -255,7 +264,7 @@ def predict_gradio(*values):
                 f"Please fill in all fields before predicting.\nMissing: {', '.join(missing)}"
             )
 
-        calibrated, _ = predict_with_comparison(user_vals)
+        calibrated, uncalibrated = predict_with_comparison(user_vals)
         calibrated_probs, calibrated_intervals = calibrated
 
         rows = []
@@ -270,19 +279,12 @@ def predict_gradio(*values):
             })
         df = pd.DataFrame(rows)
 
-        shap_plots = create_all_shap_plots(user_vals, max_display=10)
-
-        icon_outcomes = ["DEAD", "GF", "AGVHD", "CGVHD", "VOCPSHI", "STROKEHI"]
-        icon_plots = {o: icon_array(calibrated_probs[o], o) for o in icon_outcomes}
+        shap_plots  = create_all_shap_plots(user_vals, max_display=10)
+        icon_arrays = create_all_icon_arrays(calibrated_probs)
 
         return (
             df,
-            icon_plots["DEAD"],
-            icon_plots["GF"],
-            icon_plots["AGVHD"],
-            icon_plots["CGVHD"],
-            icon_plots["VOCPSHI"],
-            icon_plots["STROKEHI"],
+            icon_arrays["__grid__"],      # single combined 4×2 grid HTML
             shap_plots["DEAD"],
             shap_plots["GF"],
             shap_plots["AGVHD"],
@@ -301,6 +303,7 @@ def predict_gradio(*values):
         print("=" * 60)
         raise gr.Error(f"{type(e).__name__}: {str(e)}\n\nSee terminal for full traceback.")
 
+
 # ---------------------------------------------------------------------------
 # CSS
 # ---------------------------------------------------------------------------
@@ -335,13 +338,13 @@ with gr.Blocks(title="HCT Outcome Prediction Model", css=custom_css) as demo:
     inputs_dict = {}
 
     with gr.Row():
-        # Patient
+        # ── Patient Characteristics ──────────────────────────────────────
         with gr.Column(scale=1):
             gr.Markdown("### Patient Characteristics")
             for f in PATIENT_FEATURES:
                 inputs_dict[f] = make_component(f)
 
-        # Transplant
+        # ── Transplant Characteristics ───────────────────────────────────
         with gr.Column(scale=1):
             gr.Markdown("### Transplant Characteristics")
 
@@ -361,13 +364,13 @@ with gr.Blocks(title="HCT Outcome Prediction Model", css=custom_css) as demo:
             gvhd_final    = inputs_dict["GVHD_FINAL"]    = make_component("GVHD_FINAL")
             hla_final     = inputs_dict["HLA_FINAL"]     = make_component("HLA_FINAL")
 
-        # Disease
+        # ── Disease Characteristics ──────────────────────────────────────
         with gr.Column(scale=1):
             gr.Markdown("### Disease Characteristics")
             for f in DISEASE_FEATURES:
                 inputs_dict[f] = make_component(f)
 
-    # reactive callbacks
+    # ── Reactive callbacks ───────────────────────────────────────────────
     inputs_dict["AGE"].change(
         fn=get_age_group,
         inputs=inputs_dict["AGE"],
@@ -380,6 +383,7 @@ with gr.Blocks(title="HCT Outcome Prediction Model", css=custom_css) as demo:
         outputs=inputs_dict["VOCFRQPR"],
     )
 
+    # outputs[0] is the dropdown itself so clicking a header resets it to None
     grouped_regimen_dropdown.change(
         fn=apply_grouped_preset,
         inputs=grouped_regimen_dropdown,
@@ -402,27 +406,18 @@ with gr.Blocks(title="HCT Outcome Prediction Model", css=custom_css) as demo:
             headers=["Outcome", "Probability", "95% CI"],
             label="",
             row_count=(len(REPORTING_OUTCOMES), "dynamic"),
-            column_count=(3, "fixed"),
+            col_count=(3, "fixed"),
         )
 
     gr.Markdown("---")
-    gr.Markdown("## Icon Arrays")
-
-    with gr.Row():
-        with gr.Column():
-            icon_dead    = gr.Plot(label="Death")
-        with gr.Column():
-            icon_gf      = gr.Plot(label="Graft Failure")
-        with gr.Column():
-            icon_agvhd   = gr.Plot(label="Acute Graft-versus-Host Disease")
+    gr.Markdown("## Outcome Probability — Icon Arrays")
+    gr.Markdown(
+        "_Each figure represents 1 patient out of 100 with similar characteristics. "
+        "Colored figures indicate the predicted probability of each outcome._"
+    )
 
-    with gr.Row():
-        with gr.Column():
-            icon_cgvhd   = gr.Plot(label="Chronic Graft-versus-Host Disease")
-        with gr.Column():
-            icon_vocpshi = gr.Plot(label="Vaso-Occlusive Crisis Post-HCT")
-        with gr.Column():
-            icon_stroke  = gr.Plot(label="Stroke Post-HCT")
+    # Single HTML component holds the entire 4×2 grid
+    icon_array_grid = gr.HTML(label="")
 
     gr.Markdown("---")
     gr.Markdown("## SHAP - Feature Importance")
@@ -452,13 +447,12 @@ with gr.Blocks(title="HCT Outcome Prediction Model", css=custom_css) as demo:
         inputs=inputs_list,
         outputs=[
             output_table,
-            icon_dead, icon_gf, icon_agvhd, icon_cgvhd, icon_vocpshi, icon_stroke,
+            icon_array_grid,
             shap_dead, shap_gf, shap_agvhd, shap_cgvhd,
             shap_vocpshi, shap_efs, shap_stroke, shap_os,
         ],
     )
 
+
 if __name__ == "__main__":
-    demo.launch(
-        ssr_mode=False,
-    )
+    demo.launch(ssr_mode=False)