feat: 修复多个bug

app.py

@@ -1,3 +1,4 @@
+import time
 import streamlit as st
 import numpy as np
 from pathlib import Path
@@ -24,7 +25,7 @@ from kan.utils import create_dataset, ex_round  # type: ignore
 # Set torch dtype
 torch.set_default_dtype(torch.float64)
 
-def show_kan_prediction(model, device, samples, placeholder):
+def show_kan_prediction(model, device, samples, placeholder, phase_name):
     """显示KAN的预测结果"""
     # 生成网格数据
     x = np.linspace(-5, 5, 100)
@@ -84,7 +85,7 @@ def show_kan_prediction(model, device, samples, placeholder):
     
     # 更新布局
     fig_kan.update_layout(
-        title='KAN预测结果',
+        title='KAN预测分布',
         showlegend=True,
         width=1200,
         height=600,
@@ -100,7 +101,10 @@ def show_kan_prediction(model, device, samples, placeholder):
     fig_kan.update_yaxes(title_text='Y', row=1, col=2)
     
     # 使用占位符显示图形
-    placeholder.plotly_chart(fig_kan, use_container_width=True)
+    
+    placeholder.plotly_chart(fig_kan, 
+                             use_container_width=False, 
+                             key=f"kan_plot_{phase_name}_{time.time()}")
 
 def create_gmm_plot(dataset, centers, K, samples=None):
     """创建GMM分布的可视化图形"""
@@ -198,6 +202,7 @@ def train_kan(samples, gmm_dataset, device='cuda'):
         device = torch.device('cuda')
     else:
         device = torch.device('cpu')
+    st.info(f"使用设备: {device} 训练网络")
 
     # 转换采样点为tensor
     samples = torch.from_numpy(samples).to(device)
@@ -216,13 +221,17 @@ def train_kan(samples, gmm_dataset, device='cuda'):
     }
 
     # 创建训练进度显示组件
-    st.write("网络结构：")
-    kan_fig_placeholder = st.empty()
-    st.write("预测结果：")
-    kan_plot_placeholder = st.empty()
+    # st.write("网络预测分布：")
+    
+
+    st.write("网络图形结构：")
+    kan_network_arch_placeholder = st.empty()
+
+
     progress_container = st.container()
 
-    total_steps = 100
+    # total_steps = 100
+    total_steps = 50
     steps_per_update = 10
 
     def calculate_error(model, x, y):
@@ -260,44 +269,69 @@ def train_kan(samples, gmm_dataset, device='cuda'):
                 | 测试误差 | {test_error:.8f} |
                 """)
                 
-                # 更新进度和预测结果
-                show_kan_prediction(model, device, samples, kan_plot_placeholder)
-                
+
                 # 更新可视化（每5步更新一次）
-                if step % (steps_per_update * 5) == 0 or step + steps_per_update >= steps:
-                    # 更新预测结果
-                    show_kan_prediction(model, device, samples, kan_plot_placeholder)
+                # if step % (steps_per_update * 5) == 0 or step + steps_per_update >= steps:
+                #     # 更新预测结果
+                #     show_kan_prediction(model, device, samples, kan_plot_placeholder, phase_name)
                     
                     # 更新网络结构图（可选）
-                    if show_plot:
-                        try:
-                            kan_fig = model.plot()
-                            if isinstance(kan_fig, tuple):
-                                kan_fig = kan_fig[0]  # 如果是元组，取第一个元素
-                            if kan_fig is not None:
-                                kan_fig_placeholder.pyplot(kan_fig)
-                                plt.close('all')  # 确保关闭所有图形
-                        except Exception as e:
-                            if step == 0:  # 只在第一次出错时显示警告
-                                st.warning(f"注意：网络结构图显示失败 ({str(e)})")
+                if show_plot:
+                    try:
+                        model.plot()
+                        kan_fig = plt.gcf()
+                        # if isinstance(kan_fig, tuple):
+                            # kan_fig = kan_fig[0]  # 如果是元组，取第一个元素
+                        # if kan_fig is not None:
+                        kan_network_arch_placeholder.pyplot(kan_fig, use_container_width=False)
+                        # plt.close('all')  # 确保关闭所有图形
+                    except Exception as e:
+                        if step == 0:  # 只在第一次出错时显示警告
+                            st.warning(f"注意：网络结构图显示失败 ({str(e)})")
+
 
+                # 更新进度和预测结果
+                show_kan_prediction(model, device, samples, kan_distribution_plot_placeholder, phase_name)
+                
     with progress_container:
         st.markdown("#### 训练过程")
         error_text = st.empty()
 
     # 第一阶段训练
     # 第一阶段：初始训练
-    with st.spinner("初始训练阶段..."):
-        train_phase("第一阶段", total_steps, lamb=0.001, show_plot=False)  # 第一阶段不显示网络图
+    with st.spinner("参数调整中..."):
+        train_phase("第一阶段: 正则化训练", total_steps, lamb=0.001, show_plot=True)  
     
     # 剪枝阶段
     with st.spinner("正在进行网络剪枝优化..."):
         model = model.prune()
         progress_container.info("网络剪枝完成")
     
-    # 第二阶段：精调
-    with st.spinner("最终调优阶段..."):
-        train_phase("第二阶段", total_steps, show_plot=True)  # 第二阶段显示网络图
+    with st.spinner("参数调整中..."):
+        train_phase("第二阶段: 剪枝适应性训练", total_steps, show_plot=True)  
+    
+    with st.spinner("正在进行网格精细化..."):
+        model = model.refine(10)
+        progress_container.info("网格精细化完成")
+
+    with st.spinner("参数调整中..."):
+        train_phase("第三阶段: 网格适应性训练", total_steps, show_plot=True) 
+
+    with st.spinner("符号简化中..."):
+        # model = model.prune()
+        # progress_container.info("网络剪枝完成")
+        model.auto_symbolic()
+        progress_container.info("符号简化完成")
+
+    with st.spinner("参数调整中..."):
+        train_phase("第四阶段：符号适应性训练", total_steps, show_plot=True)  
+    
+    from kan.utils import ex_round
+    from sympy import latex
+    s= ex_round(model.symbolic_formula()[0][0],4)
+
+    st.write("网络公式：")
+    st.latex(latex(s))
     
     # 显示最终误差
     train_error = calculate_error(model, train_dataset['train_input'], train_dataset['train_label'])
@@ -307,7 +341,7 @@ def train_kan(samples, gmm_dataset, device='cuda'):
     - 训练集误差: {train_error:.6f}
     - 测试集误差: {test_error:.6f}
     """)
-    
+
     progress_container.success("🎉 训练完成！")
     return model
 
@@ -421,7 +455,7 @@ with st.sidebar:
 
     # 采样设置
     st.subheader("采样设置")
-    n_samples = st.slider("采样点数", 5, 20, 10)
+    n_samples = st.slider("采样点数", 5, 1000, 100)
     if st.button("重新采样"):
         # 创建GMM数据集进行采样
         gmm = GeneralizedGaussianMixture(
@@ -548,34 +582,37 @@ fig.update_xaxes(title_text='X', row=1, col=2)
 fig.update_yaxes(title_text='Y', row=1, col=2)
 
 # 显示GMM主图
-st.plotly_chart(fig, use_container_width=True)
+st.plotly_chart(fig, use_container_width=False)
+
 
 # KAN网络训练和预测部分
 if st.session_state.sample_points is not None:
     st.markdown("---")
     st.subheader("KAN网络训练与预测")
+
+    kan_distribution_plot_placeholder = st.empty()
     
     # 训练控制按钮
     col1, col2, col3 = st.columns([1, 2, 1])
     with col1:
-        if st.button("拟合KAN", use_container_width=True):
+        if st.button("拟合KAN", use_container_width=False):
             with st.spinner('训练KAN网络中...'):
                 st.session_state.kan_model = train_kan(st.session_state.sample_points, dataset)
                 st.balloons()
 
     with col3:
         if st.session_state.kan_model is not None:
-            if st.button("清除KAN结果", use_container_width=True):
+            if st.button("清除KAN结果", use_container_width=False):
                 st.session_state.kan_model = None
                 st.rerun()
     
     # 显示KAN预测结果
-    if st.session_state.kan_model is not None:
-        st.subheader("KAN预测结果")
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        kan_plot_placeholder = st.empty()
-        show_kan_prediction(st.session_state.kan_model, device,
-                          st.session_state.sample_points, kan_plot_placeholder)
+    # if st.session_state.kan_model is not None:
+        # st.subheader("KAN预测结果")
+        # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # kan_plot_placeholder = st.empty()
+        # show_kan_prediction(st.session_state.kan_model, device,
+        #                   st.session_state.sample_points, kan_plot_placeholder, "显示结果")
 
     st.markdown("---")
 
@@ -625,4 +662,5 @@ with st.expander("分布参数说明"):
     """)
 
 # 显示当前参数的数学公式
-st.latex(generate_latex_formula(st.session_state.p, K, centers[:K], scales[:K], weights[:K]))
+with st.expander("分布概率密度函数公式"):
+    st.latex(generate_latex_formula(st.session_state.p, K, centers[:K], scales[:K], weights[:K]))