Update app.py

app.py

@@ -9,6 +9,8 @@ import pickle
 import hashlib
 import time
 import threading
+import re
+from datetime import datetime
 
 # ===== 页面设置 =====
 st.set_page_config(page_title="多肽虚拟筛选工具", layout="wide")
@@ -16,27 +18,58 @@ st.title("多肽虚拟筛选工具")
 
 # ===== 防休眠机制 =====
 def keep_alive():
-    """每隔一段时间在后台执行一次轻量操作，防止应用被云平台冻结"""
     while True:
-        time.sleep(300)  # 每5分钟
+        time.sleep(300)
         try:
             _ = st.session_state.get("_", None)
         except Exception:
             pass
 
-# 启动后台保活线程
 if "alive_thread" not in st.session_state:
     alive_thread = threading.Thread(target=keep_alive, daemon=True)
     alive_thread.start()
     st.session_state.alive_thread = True
 
+# ===== 默认参数 =====
+DEFAULTS = {
+    "drug": "HAEGTFTSDV",
+    "min_len": 10,
+    "max_len": 12,
+    "cut_rule": "胰蛋白酶 (KR)",
+    "top_n": 10,
+    "content_threshold": 10,
+}
+
+# 初始化 session_state
+for key, value in DEFAULTS.items():
+    if key not in st.session_state:
+        st.session_state[key] = value
+
 # ===== 侧边栏参数 =====
 st.sidebar.header("筛选参数")
 
-drug = st.sidebar.text_input("靶点序列（参照肽）", value="HAEGTFTSDV")
+# 重置按钮（放在最上面）
+def reset_params():
+    for key, value in DEFAULTS.items():
+        st.session_state[key] = value
+
+st.sidebar.button("🔄 重置为默认参数", on_click=reset_params, use_container_width=True)
 
-min_len = st.sidebar.number_input("最小长度", value=10, min_value=5, max_value=50)
-max_len = st.sidebar.number_input("最大长度", value=12, min_value=5, max_value=50)
+drug = st.sidebar.text_input(
+    "靶点序列（参照肽）",
+    key="drug",
+    help="请输入20种天然氨基酸字母序列（如：HAEGTFTSDV）"
+)
+
+# 靶点输入实时校验
+if drug:
+    if re.fullmatch(r'[ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwy]+', drug):
+        st.sidebar.success("✅ 序列格式正确")
+    else:
+        st.sidebar.error("❌ 序列包含非法字符，仅限20种天然氨基酸字母")
+
+min_len = st.sidebar.number_input("最小长度", min_value=5, max_value=50, key="min_len")
+max_len = st.sidebar.number_input("最大长度", min_value=5, max_value=50, key="max_len")
 
 cut_rule = st.sidebar.selectbox(
     "酶切规则",
@@ -50,7 +83,7 @@ cut_rule = st.sidebar.selectbox(
         "中性蛋白酶 (FLIV)",
         "碱性蛋白酶 (KRFLIV)",
     ],
-    index=0,
+    key="cut_rule",
 )
 cut_map = {
     "胰蛋白酶 (KR)": "KR",
@@ -64,9 +97,8 @@ cut_map = {
 }
 CUT = cut_map[cut_rule]
 
-top_n = st.sidebar.number_input("输出Top几", value=10, min_value=5, max_value=100)
-
-content_threshold = st.sidebar.number_input("蛋白含量阈值(%)", value=10, min_value=0, max_value=100)
+top_n = st.sidebar.number_input("输出Top几", min_value=5, max_value=100, key="top_n")
+content_threshold = st.sidebar.number_input("蛋白含量阈值(%)", min_value=0, max_value=100, key="content_threshold")
 
 # ===== 生成缓存文件名 =====
 def get_cache_key():
@@ -188,123 +220,130 @@ def get_content(src_name):
 # ===== 主界面 =====
 st.subheader("开始筛选")
 
-if st.button("开始筛选", type="primary"):
-    # 加载模型
-    tok, mod = load_model()
-
-    # 拉取蛋白数据（优先读本地缓存）
-    food_cache_file = "food_cache.pkl"
-    if os.path.exists(food_cache_file):
-        with open(food_cache_file, "rb") as f:
-            FOOD = pickle.load(f)
-        st.success(f"已从缓存加载 {len(FOOD)} 条蛋白序列")
-    else:
-        with st.spinner("正在从UniProt下载食用蛋白序列..."):
-            FOOD = fetch_proteins()
-            with open(food_cache_file, "wb") as f:
-                pickle.dump(FOOD, f)
-        st.success(f"已获取并缓存 {len(FOOD)} 条蛋白序列")
-
-    # 虚拟酶切
-    allp = []
-    for name, seq in FOOD.items():
-        for p in cut(seq, CUT, min_len, max_len):
-            allp.append({"src": name, "p": p})
-
-    # 去重（按物种简称 + 肽段序列）
-    seen = set()
-    unip = []
-    for x in allp:
-        species_short = (
-            x["src"].split("(")[-1].split(")")[0].strip()
-            if "(" in x["src"]
-            else x["src"]
-        )
-        key = (species_short, x["p"])
-        if key not in seen:
-            seen.add(key)
-            unip.append(x)
-    st.info(f"酶切并去重后共 {len(unip)} 条候选肽")
-
-    # AI打分（智能缓存）
-    cache_key = get_cache_key()
-    score_cache_file = f"score_cache_{cache_key}.pkl"
-
-    if os.path.exists(score_cache_file):
-        with open(score_cache_file, "rb") as f:
-            res = pickle.load(f)
-        st.success(f"已从缓存加载 {len(res)} 条打分结果")
-    else:
-        progress_bar = st.progress(0)
-        status_text = st.empty()
-        v_drug = vec(drug, tok, mod)
-        res = []
-        total = len(unip)
-        for idx, x in enumerate(unip):
-            v_p = vec(x["p"], tok, mod)
-            s = sim(v_p, v_drug)
-            res.append({"src": x["src"], "p": x["p"], "score": round(s, 4)})
-            if idx % 50 == 0:
-                progress_bar.progress(min(idx / total, 1.0))
-                status_text.text(f"打分进度: {idx}/{total}")
-        progress_bar.progress(1.0)
-        status_text.text(f"打分完成: {total}/{total}")
-        with open(score_cache_file, "wb") as f:
-            pickle.dump(res, f)
-        st.success(f"已完成打分并缓存")
-
-    # 含量标注
-    for x in res:
-        x["content_pct"] = get_content(x["src"])
-
-    # 过滤
-    filtered_res = [x for x in res if x["content_pct"] >= content_threshold]
-    if filtered_res:
-        res = filtered_res
-
-    # 排序
-    res.sort(key=lambda x: x["score"], reverse=True)
-
-    # 显示Top结果（限制显示行数，避免内存溢出，索引从1开始）
-    st.subheader("结构相似度排名")
-    df_top = pd.DataFrame(res[:top_n])
-    df_top.index = range(1, len(df_top) + 1)
-    show_cols = ["src", "p", "score"]
-    if "content_pct" in df_top.columns:
-        show_cols.append("content_pct")
-    max_display_rows = 100
-    st.caption(f"显示前 {min(top_n, max_display_rows)} 条结果（共 {top_n} 条）")
-    st.dataframe(df_top[show_cols].head(max_display_rows), width='stretch')
-
-    # 下载按钮
-    csv_top = df_top.to_csv(index=False).encode("utf-8-sig")
-    st.download_button("下载Top结果CSV", csv_top, "out_top.csv", "text/csv")
-
-    df_all = pd.DataFrame(res)
-    csv_all = df_all.to_csv(index=False).encode("utf-8-sig")
-    st.download_button("下载全部结果CSV", csv_all, "out_all.csv", "text/csv")
-
-    # 物种统计（全部候选肽，索引从1开始）
-    st.subheader("物种贡献统计（全部候选肽）")
-    species_counts = {}
-    for x in res:
-        s = (
-            x["src"].split("(")[-1].split(")")[0].strip()
-            if "(" in x["src"]
-            else "unknown"
+# 按钮在计算中禁用
+if st.button("🚀 开始筛选", type="primary", disabled=st.session_state.get("running", False)):
+    # 输入校验
+    if not re.fullmatch(r'[ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwy]+', drug):
+        st.error("靶点序列包含非法字符，请输入有效的氨基酸序列（仅限20种天然氨基酸字母）")
+        st.stop()
+    
+    st.session_state.running = True
+    
+    filter_params = f"靶点：{drug.upper()} | 酶：{cut_rule.split('(')[0].strip()} | 长度：{min_len}-{max_len} | 时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+    
+    try:
+        tok, mod = load_model()
+
+        food_cache_file = "food_cache.pkl"
+        if os.path.exists(food_cache_file):
+            with open(food_cache_file, "rb") as f:
+                FOOD = pickle.load(f)
+            st.success(f"已从缓存加载 {len(FOOD)} 条蛋白序列")
+        else:
+            with st.spinner("正在从UniProt下载食用蛋白序列..."):
+                FOOD = fetch_proteins()
+                with open(food_cache_file, "wb") as f:
+                    pickle.dump(FOOD, f)
+            st.success(f"已获取并缓存 {len(FOOD)} 条蛋白序列")
+
+        allp = []
+        for name, seq in FOOD.items():
+            for p in cut(seq, CUT, min_len, max_len):
+                allp.append({"src": name, "p": p})
+
+        seen = set()
+        unip = []
+        for x in allp:
+            species_short = (
+                x["src"].split("(")[-1].split(")")[0].strip()
+                if "(" in x["src"]
+                else x["src"]
+            )
+            key = (species_short, x["p"])
+            if key not in seen:
+                seen.add(key)
+                unip.append(x)
+        st.info(f"酶切并去重后共 {len(unip)} 条候选肽")
+
+        cache_key = get_cache_key()
+        score_cache_file = f"score_cache_{cache_key}.pkl"
+
+        if os.path.exists(score_cache_file):
+            with open(score_cache_file, "rb") as f:
+                res = pickle.load(f)
+            st.success(f"已从缓存加载 {len(res)} 条打分结果")
+        else:
+            progress_bar = st.progress(0)
+            status_text = st.empty()
+            v_drug = vec(drug.upper(), tok, mod)
+            res = []
+            total = len(unip)
+            for idx, x in enumerate(unip):
+                v_p = vec(x["p"], tok, mod)
+                s = sim(v_p, v_drug)
+                res.append({"src": x["src"], "p": x["p"], "score": round(s, 4)})
+                if idx % 50 == 0:
+                    progress_bar.progress(min(idx / total, 1.0))
+                    status_text.text(f"打分进度: {idx}/{total}")
+            progress_bar.progress(1.0)
+            status_text.text(f"打分完成: {total}/{total}")
+            with open(score_cache_file, "wb") as f:
+                pickle.dump(res, f)
+            st.success(f"已完成打分并缓存")
+
+        for x in res:
+            x["content_pct"] = get_content(x["src"])
+
+        filtered_res = [x for x in res if x["content_pct"] >= content_threshold]
+        if filtered_res:
+            res = filtered_res
+
+        res.sort(key=lambda x: x["score"], reverse=True)
+
+        st.caption(f"📋 {filter_params}")
+
+        st.subheader("结构相似度排名")
+        df_top = pd.DataFrame(res[:top_n])
+        df_top.index = range(1, len(df_top) + 1)
+        show_cols = ["src", "p", "score"]
+        if "content_pct" in df_top.columns:
+            show_cols.append("content_pct")
+        max_display_rows = 100
+        st.caption(f"显示前 {min(top_n, max_display_rows)} 条结果（共 {top_n} 条）")
+        st.dataframe(df_top[show_cols].head(max_display_rows), width='stretch')
+
+        csv_top = df_top.to_csv(index=False).encode("utf-8-sig")
+        st.download_button("下载Top结果CSV", csv_top, "out_top.csv", "text/csv")
+
+        df_all = pd.DataFrame(res)
+        csv_all = df_all.to_csv(index=False).encode("utf-8-sig")
+        st.download_button("下载全部结果CSV", csv_all, "out_all.csv", "text/csv")
+
+        st.subheader("物种贡献统计（全部候选肽）")
+        species_counts = {}
+        for x in res:
+            s = (
+                x["src"].split("(")[-1].split(")")[0].strip()
+                if "(" in x["src"]
+                else "unknown"
+            )
+            species_counts[s] = species_counts.get(s, 0) + 1
+        species_df = pd.DataFrame(
+            list(species_counts.items()), columns=["物种", "候选肽数量"]
         )
-        species_counts[s] = species_counts.get(s, 0) + 1
-    species_df = pd.DataFrame(
-        list(species_counts.items()), columns=["物种", "候选肽数量"]
-    )
-    species_df = species_df.sort_values("候选肽数量", ascending=False)
-    species_df = species_df.reset_index(drop=True)
-    species_df.index = range(1, len(species_df) + 1)
-
-    col1, col2 = st.columns(2)
-    with col1:
-        st.dataframe(species_df.head(20), width='stretch')
-    with col2:
-        st.bar_chart(species_df.set_index("物种"))
-
-    st.success("筛选完成！")
+        species_df = species_df.sort_values("候选肽数量", ascending=False)
+        species_df = species_df.reset_index(drop=True)
+        species_df.index = range(1, len(species_df) + 1)
+        total_count = species_df["候选肽数量"].sum()
+        species_df["占比(%)"] = (species_df["候选肽数量"] / total_count * 100).round(1)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.dataframe(species_df[["物种", "候选肽数量", "占比(%)"]].head(20), width='stretch')
+        with col2:
+            st.bar_chart(species_df.set_index("物种")["候选肽数量"])
+
+        st.success("筛选完成！")
+        
+    finally:
+        st.session_state.running = False