miyuiu/predictability-gradient
0
1from __future__ import annotations2 3import json4from pathlib import Path5 6import altair as alt7import pandas as pd8import streamlit as st9 10 11ROOT = Path(__file__).resolve().parent12ASSET_PATH = ROOT / "assets" / "predictability_gradient.json"13 14 15def install_css() -> None:16 st.markdown(17 """18<link rel="preconnect" href="https://fonts.googleapis.com">19<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>20<link href="https://fonts.googleapis.com/css2?family=IBM+Plex+Sans:wght@400;500;600&family=IBM+Plex+Mono:wght@400;500&family=IBM+Plex+Serif:ital,wght@0,400;0,500;1,400;1,500&display=swap" rel="stylesheet">21<style>22:root {23 --paper: #f5f1e8;24 --paper-deep: #ece6d6;25 --ink: #1f1d18;26 --ink-soft: #5a554a;27 --ink-faint: #94907f;28 --rule: #d6cdb6;29 --rule-soft: #e6dfca;30 --accent: #a8521a;31 --accent-tint: #fdf6e8;32 --pos: #3f6b3a;33 --o2: #3a7d6e;34 --focused-strip: #ede4cd;35 --serif: 'IBM Plex Serif', Georgia, serif;36 --sans: 'IBM Plex Sans', system-ui, sans-serif;37 --mono: 'IBM Plex Mono', ui-monospace, monospace;38}39html, body, [data-testid="stAppViewContainer"], [data-testid="stMain"] { background: var(--paper) !important; }40[data-testid="stHeader"] { background: transparent !important; }41[data-testid="stToolbar"], footer, #MainMenu { visibility: hidden; }42.stApp { background: var(--paper); }43.block-container { max-width: 100% !important; padding: 0 0 4rem 0 !important; }44.main .block-container > div:first-child { padding-top: 0; }45body, p, div, span, li, label, .stMarkdown { font-family: var(--sans) !important; color: var(--ink); }46h1, h2, h3, h4 { font-family: var(--serif) !important; color: var(--ink); letter-spacing: -0.01em; }47code, pre { font-family: var(--mono) !important; color: var(--accent); }48.pg-kicker { font-family: var(--mono); font-size: 11px; color: var(--ink-faint); letter-spacing: .05em; text-transform: uppercase; }49.pg-hero { padding: 38px 34px 24px; border-bottom: 1px solid var(--rule); margin-bottom: 0; background: var(--paper); }50.pg-hero h1 { font-size: clamp(2.5rem, 6vw, 5.4rem); line-height: .9; margin: .22rem 0 .75rem; max-width: 1040px; font-weight: 500; }51.pg-hero p { max-width: 860px; color: var(--ink-soft); font-size: 1.06rem; line-height: 1.55; }52.pg-card { border: 1px solid var(--rule); background: var(--paper); padding: 16px 18px; border-radius: 2px; min-height: 112px; transition: border-color 120ms; }53.pg-card:hover { border-color: var(--ink); }54.pg-card .num { font-family: var(--serif); color: var(--ink); font-size: 32px; font-weight: 500; font-variant-numeric: tabular-nums; line-height: 1; }55.pg-card .label { color: var(--ink-soft); font-family: var(--mono); font-size: 10px; letter-spacing: .05em; text-transform: uppercase; margin-top: .5rem; }56.pg-pill { display: inline-block; border: 1px solid var(--accent); color: var(--accent); padding: 1px 6px; margin: 0 6px 7px 0; border-radius: 2px; font-size: 11px; font-family: var(--mono); }57.pg-callout { border: 1px solid var(--rule); background: var(--focused-strip); padding: 14px 16px; margin: 1rem 0; border-radius: 2px; color: var(--ink); font-family: var(--serif); font-style: italic; }58.pg-table { width: 100%; border-collapse: collapse; font-family: var(--sans); background: var(--paper); border: 1px solid var(--rule); }59.pg-table th { background: var(--paper-deep); border-bottom: 1px solid var(--rule); padding: 8px 12px; text-align: left; font-family: var(--mono); font-size: 10px; font-weight: 500; color: var(--ink-soft); letter-spacing: .05em; text-transform: uppercase; white-space: nowrap; }60.pg-table td { padding: 10px 12px; border-bottom: 1px solid var(--rule-soft); font-size: 12px; color: var(--ink); vertical-align: middle; }61.pg-table tr:last-child td { border-bottom: none; }62.pg-table tr:hover { background: #ede5cd; }63[data-baseweb="tab-list"] { background: var(--paper) !important; border-bottom: 1px solid var(--rule) !important; padding: 0 28px !important; gap: 0 !important; }64[data-baseweb="tab"] { font-family: var(--sans) !important; font-size: 13px !important; color: var(--ink-faint) !important; padding: 12px 18px !important; height: auto !important; background: transparent !important; }65[data-baseweb="tab"] p { color: var(--ink-faint) !important; }66[data-baseweb="tab"][aria-selected="true"] p { color: var(--ink) !important; font-weight: 500 !important; }67[data-baseweb="tab-highlight"] { background: var(--ink) !important; height: 2px !important; }68[data-baseweb="tab-border"] { display: none !important; }69[data-baseweb="tab-panel"] { padding: 0 28px !important; }70.stCaptionContainer, .stCaptionContainer p { color: var(--ink-soft) !important; font-family: var(--sans) !important; }71</style>72 """,73 unsafe_allow_html=True,74 )75 76 77def load_asset() -> dict:78 if ASSET_PATH.exists():79 return json.loads(ASSET_PATH.read_text())80 return {81 "paper": {"n_genomes": 19592, "n_proteins": 82000000, "n_traits": 21, "embedding": "ESM-2 t30 150M"},82 "headline_gradient": [83 {"split": "species", "class": "compositional", "delta_f1": 0.021, "std": 0.002},84 {"split": "species", "class": "machinery", "delta_f1": 0.083, "std": 0.012},85 {"split": "genus", "class": "compositional", "delta_f1": 0.016, "std": 0.004},86 {"split": "genus", "class": "machinery", "delta_f1": 0.067, "std": 0.010},87 {"split": "family", "class": "compositional", "delta_f1": 0.009, "std": 0.002},88 {"split": "family", "class": "machinery", "delta_f1": 0.010, "std": 0.003},89 ],90 "trait_classes": {"compositional": [], "machinery": [], "excluded": []},91 "head_deltas": [],92 "attention": {"animal": {}, "human": {}, "genes": []},93 "comparison_rows": [],94 }95 96 97def metric_card(value: str, label: str) -> str:98 return f"<div class='pg-card'><div class='num'>{value}</div><div class='label'>{label}</div></div>"99 100 101def pct(x: float) -> str:102 return f"{100 * x:.1f}%"103 104 105def gradient_chart(df: pd.DataFrame):106 base = (107 alt.Chart(df)108 .mark_bar(cornerRadiusTopLeft=3, cornerRadiusTopRight=3)109 .encode(110 x=alt.X("split:N", sort=["species", "genus", "family"], title=None),111 xOffset=alt.XOffset("class:N"),112 y=alt.Y("delta_f1:Q", title="Attention - mean macro-F1"),113 color=alt.Color(114 "class:N",115 scale=alt.Scale(domain=["compositional", "machinery"], range=["#3a7d6e", "#a8521a"]),116 legend=alt.Legend(orient="top"),117 ),118 tooltip=["split", "class", alt.Tooltip("delta_f1:Q", format=".3f"), alt.Tooltip("std:Q", format=".3f")],119 )120 .properties(height=330)121 )122 err = (123 alt.Chart(df)124 .mark_errorbar()125 .encode(126 x=alt.X("split:N", sort=["species", "genus", "family"]),127 xOffset=alt.XOffset("class:N"),128 y=alt.Y("low:Q"),129 y2="high:Q",130 color=alt.Color("class:N", scale=alt.Scale(domain=["compositional", "machinery"], range=["#3a7d6e", "#a8521a"])),131 )132 )133 return (base + err).configure_view(stroke="#d6cdb6").configure_axis(134 labelColor="#5a554a",135 titleColor="#5a554a",136 gridColor="#e6dfca",137 domainColor="#d6cdb6",138 tickColor="#d6cdb6",139 ).configure_legend(labelColor="#1f1d18", titleColor="#5a554a")140 141 142def render_table(rows: list[dict], cols: list[str]) -> None:143 html_rows = []144 for row in rows:145 html_rows.append("<tr>" + "".join(f"<td>{row.get(c, '')}</td>" for c in cols) + "</tr>")146 st.markdown(147 "<table class='pg-table'><tr>"148 + "".join(f"<th>{c.replace('_', ' ')}</th>" for c in cols)149 + "</tr>"150 + "".join(html_rows)151 + "</table>",152 unsafe_allow_html=True,153 )154 155 156def main() -> None:157 st.set_page_config(page_title="predictability-gradient", page_icon=None, layout="wide")158 install_css()159 asset = load_asset()160 paper = asset["paper"]161 162 st.markdown(163 """164<div class='pg-hero'>165 <div class='pg-kicker'>Paper companion / mechanism exhibit</div>166 <h1>When does attention help?</h1>167 <p>A compact interactive view of the predictability-gradient result:168 attention-pooling helps gene-localized machinery traits, contributes little to169 diffuse compositional traits, and loses its advantage under family-level shift.</p>170</div>171 """,172 unsafe_allow_html=True,173 )174 175 c1, c2, c3, c4 = st.columns(4)176 c1.markdown(metric_card(f"{paper['n_genomes']:,}", "BacDive genomes"), unsafe_allow_html=True)177 c2.markdown(metric_card("82M", "per-protein embeddings"), unsafe_allow_html=True)178 c3.markdown(metric_card(str(paper["n_traits"]), "prediction heads"), unsafe_allow_html=True)179 c4.markdown(metric_card("3 x 3", "splits x seeds"), unsafe_allow_html=True)180 181 tabs = st.tabs(["Gradient", "Trait Classes", "Attention Spotlight", "VFDB + Ablation", "Benchmark Context"])182 183 with tabs[0]:184 df = pd.DataFrame(asset["headline_gradient"])185 df["low"] = df["delta_f1"] - df["std"]186 df["high"] = df["delta_f1"] + df["std"]187 st.altair_chart(gradient_chart(df), use_container_width=True)188 st.markdown(189 "<div class='pg-callout'>The species-level machinery gain is about four times the compositional gain "190 "(+0.083 vs +0.021 F1). At family holdout, the gap shrinks to +0.001, localizing the bottleneck to "191 "cross-clade generalization.</div>",192 unsafe_allow_html=True,193 )194 if asset.get("head_deltas"):195 top = pd.DataFrame(asset["head_deltas"]).head(10)196 top["delta_f1_mean"] = top["delta_f1_mean"].map(lambda x: f"{x:+.3f}")197 top["delta_f1_std"] = top["delta_f1_std"].map(lambda x: f"{x:.3f}")198 render_table(top.to_dict("records"), ["split", "head", "class", "delta_f1_mean", "delta_f1_std", "n_seeds"])199 200 with tabs[1]:201 classes = asset["trait_classes"]202 left, right = st.columns(2)203 with left:204 st.subheader("Compositional")205 st.markdown("".join(f"<span class='pg-pill'>{x}</span>" for x in classes["compositional"]), unsafe_allow_html=True)206 st.caption("Diffuse or bulk cellular signals where mean-pooling should be competitive.")207 with right:208 st.subheader("Machinery")209 st.markdown("".join(f"<span class='pg-pill'>{x}</span>" for x in classes["machinery"]), unsafe_allow_html=True)210 st.caption("Gene-localized or pathway-localized signals where attention has something to select.")211 st.markdown(212 "<div class='pg-callout'>The partition is biological, not result-driven: two metadata heads "213 "(isolation source and country) are excluded from the gradient analysis.</div>",214 unsafe_allow_html=True,215 )216 217 with tabs[2]:218 animal = asset["attention"]["animal"]219 cols = st.columns(4)220 cols[0].markdown(metric_card(f"{animal.get('auroc', 0):.2f}", "animal pathogenicity AUROC"), unsafe_allow_html=True)221 cols[1].markdown(metric_card(f"{animal.get('median_entropy', 0):.2f}", "median attention entropy"), unsafe_allow_html=True)222 cols[2].markdown(metric_card(pct(animal.get("top5_attention_mass", 0)), "attention in top 5 proteins"), unsafe_allow_html=True)223 cols[3].markdown(metric_card("~3,800", "proteins per genome"), unsafe_allow_html=True)224 st.markdown("".join(f"<span class='pg-pill'>{gene}</span>" for gene in asset["attention"].get("genes", [])), unsafe_allow_html=True)225 st.markdown(226 "<div class='pg-callout'>The spotlighted genes are coherent adherence and invasion machinery, including "227 "fimbrial ushers, filamentous hemagglutinin, invasion loci, type-IV pili, and flagellar proteins.</div>",228 unsafe_allow_html=True,229 )230 231 with tabs[3]:232 animal = asset["attention"]["animal"]233 human = asset["attention"]["human"]234 rows = [235 {236 "test": "within genome",237 "head": "animal",238 "top_attended": pct(animal.get("within_top", 0)),239 "control": pct(animal.get("within_random", 0)),240 "statistic": f"Wilcoxon p={animal.get('within_p', 'n/a')}",241 },242 {243 "test": "between class",244 "head": "animal",245 "top_attended": pct(animal.get("between_top_pathogenic", 0)),246 "control": pct(animal.get("between_top_non_pathogenic", 0)),247 "statistic": f"OR={animal.get('between_or', 'n/a')}, p={animal.get('between_p', 'n/a')}",248 },249 {250 "test": "between class",251 "head": "human",252 "top_attended": "replicates",253 "control": "non-pathogenic genomes",254 "statistic": f"OR={human.get('between_or', 'n/a')}, p={human.get('between_p', 'n/a')}",255 },256 {257 "test": "top-5 ablation",258 "head": "animal",259 "top_attended": pct(animal.get("ablation_flip", 0)),260 "control": "random removal near zero",261 "statistic": f"flip p={animal.get('ablation_p', 'n/a')}",262 },263 {264 "test": "top-5 ablation",265 "head": "human",266 "top_attended": pct(human.get("ablation_flip", 0)),267 "control": "random removal near zero",268 "statistic": f"flip p={human.get('ablation_p', 'n/a')}",269 },270 ]271 render_table(rows, ["test", "head", "top_attended", "control", "statistic"])272 st.markdown(273 "<div class='pg-callout'>The mechanism claim rests on the conjunction: enrichment against an external "274 "virulence-factor database and causal dependence under ablation.</div>",275 unsafe_allow_html=True,276 )277 278 with tabs[4]:279 rows = asset.get("comparison_rows", [])280 if rows:281 render_table(rows, ["trait", "metric", "best_ours", "run", "prior", "prior_score", "verdict"])282 st.markdown(283 "<div class='pg-callout'>This tab is context for the benchmark, not the core paper claim. The paper's "284 "main contribution is the pooling rule plus mechanistic validation, not a new encoder sweep.</div>",285 unsafe_allow_html=True,286 )287 288 289if __name__ == "__main__":290 main()291 