implement ec prediction in the app and interface

ProtHGT_app.py

@@ -15,6 +15,7 @@ st.set_page_config(
 from rapidfuzz import fuzz, process
 
 from run_prothgt_app import KG_PATH, generate_prediction_df, load_kg, load_model, load_name_info
+from run_prothgt_ec_app import EC_KG_PATH, generate_ec_prediction_df, load_ec_kg, load_ec_model
 from visualize_kg import visualize_protein_subgraph
 from ui import GITHUB_URL, KG_VIEW_GUIDE, STRICTNESS, render_header, render_stepper
 
@@ -38,7 +39,7 @@ CATEGORY_TO_KEY = {"Molecular Function": "GO_term_F", "Biological Process": "GO_
 # Widget values that must survive moving between steps (Streamlit drops the state of
 # widgets that are not rendered; re-assigning keeps it).
 PERSISTENT_KEYS = ["proteins", "aspect_GO_term_F", "aspect_GO_term_P", "aspect_GO_term_C",
-                   "strictness", "use_custom_threshold", "custom_threshold"]
+                   "predict_ec", "strictness", "use_custom_threshold", "custom_threshold"]
 for _k in PERSISTENT_KEYS:
     if _k in st.session_state:
         st.session_state[_k] = st.session_state[_k]
@@ -49,6 +50,7 @@ DEFAULTS = {
     "aspect_GO_term_F": True,
     "aspect_GO_term_P": True,
     "aspect_GO_term_C": True,
+    "predict_ec": False,
     "strictness": "Balanced",
     "use_custom_threshold": False,
     "custom_threshold": 0.50,
@@ -89,6 +91,16 @@ def get_model(go_cat):
     return load_model(go_cat, get_kg())
 
 
+@st.cache_resource(show_spinner=False)
+def get_ec_kg():
+    return load_ec_kg()
+
+
+@st.cache_resource(show_spinner=False)
+def get_ec_model():
+    return load_ec_model(get_ec_kg())
+
+
 @st.cache_data(show_spinner=False, max_entries=200)
 def search_proteins(query):
     """Exact ID → ID prefix → name substring (closest names first) → fuzzy fallback."""
@@ -201,8 +213,9 @@ hero_slot = st.empty()  # always created so the elements below keep a stable pos
 if step < 3:
     hero_slot.markdown(
         '<div class="phgt-hero"><h1>Protein function prediction with knowledge graphs</h1>'
-        "<p>Gene Ontology term predictions from a heterogeneous graph transformer that combines "
-        "protein language model embeddings with a large-scale biological knowledge graph.</p></div>",
+        "<p>Gene Ontology term and Enzyme Commission number predictions from a heterogeneous graph "
+        "transformer that combines protein language model embeddings with a large-scale biological "
+        "knowledge graph.</p></div>",
         unsafe_allow_html=True,
     )
 else:
@@ -354,12 +367,23 @@ def render_step_2():
     with main:
         with st.container(border=True):
             st.subheader("Which functions are you interested in?")
-            cols = st.columns(3)
-            for col, (key, (label, desc)) in zip(cols, ASPECTS.items()):
-                with col:
-                    with st.container(border=True):
-                        st.checkbox(f"**{label}**", key=f"aspect_{key}")
-                        st.caption(desc)
+            # 2x2 grid: the three GO aspects plus EC numbers. Fixed-height cards so
+            # every box is the same size regardless of how much text it holds.
+            items = [(f"aspect_{k}", "Gene Ontology", label, desc) for k, (label, desc) in ASPECTS.items()]
+            items.append((
+                "predict_ec", "Enzyme Commission", "EC number",
+                "EC numbers for enzymes, predicted across the four hierarchy levels (L1–L4) as a "
+                "consistent top-down chain. Non-enzymes return no EC number.",
+            ))
+            for i in range(0, len(items), 2):
+                cols = st.columns(2)
+                for col, (key, eyebrow, title, desc) in zip(cols, items[i:i + 2]):
+                    with col:
+                        with st.container(border=True, height=158):
+                            st.markdown(f'<span class="phgt-eyebrow">{eyebrow}</span>',
+                                        unsafe_allow_html=True)
+                            st.checkbox(f"**{title}**", key=key)
+                            st.caption(desc)
 
             st.subheader("How strict should predictions be?")
             st.radio(
@@ -382,13 +406,13 @@ def render_step_2():
                     st.caption("A very low threshold returns almost every GO term for each protein; "
                                "expect a large table.")
 
-        aspects = selected_aspects()
+        anything = selected_aspects() or st.session_state.predict_ec
         back, _, run = st.columns([1, 1.7, 1.3])
         back.button("← Back", use_container_width=True, on_click=go_to, args=(1,))
         run.button("Run prediction →", type="primary", use_container_width=True,
-                   disabled=not aspects, on_click=go_to, args=(3,))
-        if not aspects:
-            st.caption("Select at least one function category.")
+                   disabled=not anything, on_click=go_to, args=(3,))
+        if not anything:
+            st.caption("Select at least one function category or EC numbers.")
     with aside:
         render_limit_aside()
 
@@ -396,8 +420,8 @@ def render_step_2():
 # ---------------------------------------------------------------------------
 # Step 3 — run & review
 # ---------------------------------------------------------------------------
-def run_prediction(proteins, aspects, threshold):
-    with st.status("Running ProtHGT…", expanded=True) as status:
+def run_go_prediction(proteins, aspects, threshold):
+    with st.status("Running ProtHGT (GO terms)…", expanded=True) as status:
         if os.path.exists(KG_PATH):
             st.write("Loading the knowledge graph…")
         else:
@@ -423,6 +447,27 @@ def run_prediction(proteins, aspects, threshold):
     return df.sort_values("Probability", ascending=False, ignore_index=True)
 
 
+def run_ec_prediction(proteins, threshold):
+    with st.status("Running ProtHGT (EC numbers)…", expanded=True) as status:
+        if os.path.exists(EC_KG_PATH):
+            st.write("Loading the EC knowledge graph…")
+        else:
+            st.write("Downloading EC dataset files for the first time, this may take a few minutes…")
+        heterodata = get_ec_kg()
+        model = get_ec_model()
+        progress = st.progress(0.0)
+
+        def on_protein(i, n):
+            progress.progress(i / n, text=f"Scoring EC numbers ({i + 1}/{n})")
+
+        _, df = generate_ec_prediction_df(
+            proteins, heterodata, model, threshold=threshold, progress_callback=on_protein)
+        progress.progress(1.0, text="Done")
+        status.update(label=f"Predicted {len(df):,} protein–EC number associations", state="complete",
+                      expanded=False)
+    return df.sort_values("Probability", ascending=False, ignore_index=True)
+
+
 @st.fragment
 def render_predictions(df, threshold):
     if df.empty:
@@ -479,6 +524,69 @@ def render_predictions(df, threshold):
                        use_container_width=True)
 
 
+def _expasy_url(ec):
+    """ExPASy ENZYME URL for an EC number or class prefix (padded to four fields with '-')."""
+    parts = ec.split(".")
+    parts += ["-"] * (4 - len(parts))
+    return "https://enzyme.expasy.org/EC/" + ".".join(parts)
+
+
+@st.fragment
+def render_ec_predictions(df, threshold):
+    if df.empty:
+        st.info("No EC numbers passed the score threshold — none of these proteins were predicted as "
+                "enzymes at this strictness. Try a lower strictness in step 2.")
+        return
+
+    f1, f2, f3 = st.columns([1.2, 1.2, 1.6])
+    proteins = ["All proteins"] + sorted(df["UniProt_ID"].unique())
+    protein = f1.selectbox("Protein", proteins, key="ec_protein",
+                           format_func=lambda p: p if p == "All proteins"
+                           else f"{p} — {df.loc[df.UniProt_ID == p, 'Protein'].iat[0]}")
+    all_levels = ["L1", "L2", "L3", "L4"]
+    available = [lv for lv in all_levels if lv in set(df["EC_level"])]
+    levels = f2.pills("Level", available, selection_mode="multi", default=available, key="ec_levels")
+    text = f3.text_input("Search EC number", placeholder="e.g. 2.7.11", key="ec_search")
+    min_score = st.slider("Minimum score", float(threshold), 1.0, float(threshold), step=0.01, key="ec_min")
+
+    view = df
+    if protein != "All proteins":
+        view = view[view["UniProt_ID"] == protein]
+    view = view[view["EC_level"].isin(levels or [])]
+    if text.strip():
+        view = view[view["EC"].str.contains(text.strip(), case=False, regex=False)]
+    view = view[view["Probability"] >= min_score]
+
+    display = view.assign(
+        UniProt_ID="https://www.uniprot.org/uniprotkb/" + view["UniProt_ID"] + "/entry",
+        EC=view["EC"].map(_expasy_url),
+    )
+    st.dataframe(
+        display,
+        hide_index=True,
+        use_container_width=True,
+        height=520,
+        column_config={
+            "UniProt_ID": st.column_config.LinkColumn("UniProt ID", display_text=r"uniprotkb/([A-Z0-9]+)/entry"),
+            "Protein": st.column_config.TextColumn("Protein", width="large"),
+            "EC": st.column_config.LinkColumn("EC number", display_text=r"EC/([\d.]*\d)",
+                                              help="L1 (class) → L4 (full EC number). Opens the ExPASy ENZYME entry."),
+            "EC_name": st.column_config.TextColumn("Name", width="large",
+                                                   help="Enzyme class (L1–L3) or enzyme name (L4)."),
+            "EC_level": st.column_config.TextColumn("Level", width="small"),
+            "Probability": st.column_config.NumberColumn("Score", format="%.3f", width="small"),
+        },
+    )
+    st.caption("Predictions form a consistent hierarchy: every predicted L4 number has its predicted "
+               "L3/L2/L1 ancestors. A protein can stop at an intermediate level when no deeper level "
+               "passes the threshold.")
+    c1, c2 = st.columns([3, 1.3], vertical_alignment="center")
+    c1.caption(f"Showing {len(view):,} of {len(df):,} predictions")
+    c2.download_button("Download CSV", view.to_csv(index=False).encode("utf-8"),
+                       file_name="prothgt_ec_predictions.csv", mime="text/csv",
+                       use_container_width=True, key="ec_download")
+
+
 @st.fragment
 def render_knowledge_graph(df, proteins, aspects):
     if len(proteins) > MAX_KG_VIZ_PROTEINS:
@@ -532,31 +640,35 @@ def render_knowledge_graph(df, proteins, aspects):
 def render_step_3():
     proteins = list(st.session_state.proteins)
     aspects = selected_aspects()
-    if not proteins or not aspects:
+    predict_ec = st.session_state.predict_ec
+    if not proteins or not (aspects or predict_ec):
         go_to(1 if not proteins else 2)
         st.rerun()
 
     threshold = current_threshold()
-    run_key = (tuple(proteins), tuple(aspects), threshold)
+    run_key = (tuple(proteins), tuple(aspects), predict_ec, threshold)
     results = st.session_state.results
     if results is None or results["key"] != run_key:
         clear_visualizations()
         cache = st.session_state.results_cache
         if run_key in cache:  # switching back to settings already run in this session is instant
-            df = cache[run_key]
+            payload = cache[run_key]
         else:
             try:
-                df = run_prediction(proteins, aspects, threshold)
+                df_go = run_go_prediction(proteins, aspects, threshold) if aspects else None
+                df_ec = run_ec_prediction(proteins, threshold) if predict_ec else None
             except Exception as exc:  # surface failures instead of a blank page
                 st.error(f"Prediction failed: {exc}")
                 st.button("← Back to settings", on_click=go_to, args=(2,))
                 return
-            cache[run_key] = df
+            payload = {"go": df_go, "ec": df_ec}
+            cache[run_key] = payload
             while len(cache) > 3:  # keep memory per visitor bounded
                 cache.pop(next(iter(cache)))
-        st.session_state.results = {"key": run_key, "df": df}
+        st.session_state.results = {"key": run_key, "payload": payload}
         results = st.session_state.results
-    df = results["df"]
+    payload = results["payload"]
+    df_go, df_ec = payload["go"], payload["ec"]
 
     head, b1, b2 = st.columns([3, 1.1, 1.1], vertical_alignment="bottom")
     head.subheader("Results")
@@ -565,15 +677,21 @@ def render_step_3():
 
     m1, m2, m3, m4 = st.columns(4)
     m1.metric("Proteins", len(proteins))
-    m2.metric("Predicted GO terms", f"{len(df):,}")
-    m3.metric("Aspects", " · ".join(ASPECT_SHORT[a] for a in aspects))
+    m2.metric("Predicted GO terms", f"{len(df_go):,}" if df_go is not None else "—")
+    m3.metric("Predicted EC numbers", f"{len(df_ec):,}" if df_ec is not None else "—")
     m4.metric("Score threshold", f"{threshold:.2f}")
 
-    pred_tab, kg_tab = st.tabs(["Predictions", "Knowledge graph"])
-    with pred_tab:
-        render_predictions(df, threshold)
-    with kg_tab:
-        render_knowledge_graph(df, proteins, aspects)
+    tab_specs = []
+    if aspects:
+        tab_specs.append(("GO terms", lambda: render_predictions(df_go, threshold)))
+    if predict_ec:
+        tab_specs.append(("EC numbers", lambda: render_ec_predictions(df_ec, threshold)))
+    if aspects:
+        tab_specs.append(("Knowledge graph", lambda: render_knowledge_graph(df_go, proteins, aspects)))
+
+    for tab, (_, render) in zip(st.tabs([label for label, _ in tab_specs]), tab_specs):
+        with tab:
+            render()
 
 
 # Render the step inside one st.empty() slot so the previous step's widgets are replaced

ui.py

@@ -110,6 +110,8 @@ div[data-testid="stDataFrame"] .stDataFrameGlideDataEditor { border-radius: 16px
 .phgt-chips { display: flex; flex-wrap: wrap; gap: 6px; margin: 4px 0 10px; }
 .phgt-chip { padding: 4px 9px; border-radius: 8px; background: #F3F2EE; font: 500 13px 'DM Mono', monospace; }
 .phgt-muted { color: #4A4E5C; font-size: 0.9rem; }
+/* Small category label above a card title (e.g. "Gene Ontology" over "Molecular function") */
+.phgt-eyebrow { display: block; font-size: 0.68rem; font-weight: 700; letter-spacing: 0.07em; text-transform: uppercase; color: ACCENT; margin-bottom: 2px; }
 .phgt-meter { height: 6px; border-radius: 3px; background: #EDECE7; margin: 6px 0 4px; }
 .phgt-meter > div { height: 6px; border-radius: 3px; background: ACCENT; }
 .phgt-count { display: block; padding-bottom: 0.6rem; }  /* breathing room below the protein count */