feat: KAN可视化有bug

.gitignore

@@ -1,3 +1,5 @@
+figures/
+model/
 *.npz
 # Byte-compiled / optimized / DLL files
 __pycache__/

app.py

@@ -5,6 +5,311 @@ from experiments.gmm_dataset import GeneralizedGaussianMixture
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 from typing import List, Tuple
+import torch
+import os
+import sys
+import matplotlib.pyplot as plt
+
+# Set torch path
+torch.classes.__path__ = [os.path.join(torch.__path__[0], torch.classes.__file__ or "")]
+
+# Add pykan to path
+pykan_path = Path(__file__).parent.parent / 'third_party' / 'pykan'
+sys.path.append(str(pykan_path))
+
+# Import KAN related modules
+from kan import KAN  # type: ignore
+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):
+    """显示KAN的预测结果"""
+    # 生成网格数据
+    x = np.linspace(-5, 5, 100)
+    y = np.linspace(-5, 5, 100)
+    X, Y = np.meshgrid(x, y)
+    xy = np.column_stack((X.ravel(), Y.ravel()))
+    
+    # 使用KAN预测
+    grid_points = torch.from_numpy(xy).to(device)
+    with torch.no_grad():
+        Z_kan = model(grid_points).cpu().numpy().reshape(X.shape)
+    
+    # 创建预测的概率密度图
+    fig_kan = make_subplots(
+        rows=1, cols=2,
+        specs=[[{'type': 'surface'}, {'type': 'contour'}]],
+        subplot_titles=('KAN预测的3D概率密度曲面', 'KAN预测的等高线图')
+    )
+    
+    # 3D Surface
+    surface_kan = go.Surface(
+        x=X, y=Y, z=Z_kan,
+        colorscale='viridis',
+        showscale=True,
+        colorbar=dict(x=0.45)
+    )
+    fig_kan.add_trace(surface_kan, row=1, col=1)
+    
+    # Contour Plot
+    contour_kan = go.Contour(
+        x=x, y=y, z=Z_kan,
+        colorscale='viridis',
+        showscale=True,
+        colorbar=dict(x=1.0),
+        contours=dict(
+            showlabels=True,
+            labelfont=dict(size=12)
+        )
+    )
+    fig_kan.add_trace(contour_kan, row=1, col=2)
+    
+    # 添加采样点
+    if samples is not None:
+        fig_kan.add_trace(
+            go.Scatter(
+                x=samples[:, 0], y=samples[:, 1],
+                mode='markers',
+                marker=dict(
+                    size=8,
+                    color='yellow',
+                    line=dict(color='black', width=1)
+                ),
+                name='训练点'
+            ),
+            row=1, col=2
+        )
+    
+    # 更新布局
+    fig_kan.update_layout(
+        title='KAN预测结果',
+        showlegend=True,
+        width=1200,
+        height=600,
+        scene=dict(
+            xaxis_title='X',
+            yaxis_title='Y',
+            zaxis_title='密度'
+        )
+    )
+    
+    # 更新2D图的坐标轴
+    fig_kan.update_xaxes(title_text='X', row=1, col=2)
+    fig_kan.update_yaxes(title_text='Y', row=1, col=2)
+    
+    # 使用占位符显示图形
+    placeholder.plotly_chart(fig_kan, use_container_width=True)
+
+def create_gmm_plot(dataset, centers, K, samples=None):
+    """创建GMM分布的可视化图形"""
+    # 生成网格数据
+    x = np.linspace(-5, 5, 100)
+    y = np.linspace(-5, 5, 100)
+    X, Y = np.meshgrid(x, y)
+    xy = np.column_stack((X.ravel(), Y.ravel()))
+
+    # 计算概率密度
+    Z = dataset.pdf(xy).reshape(X.shape)
+
+    # 创建2D和3D可视化
+    fig = make_subplots(
+        rows=1, cols=2,
+        specs=[[{'type': 'surface'}, {'type': 'contour'}]],
+        subplot_titles=('3D概率密度曲面', '等高线图与分量中心')
+    )
+
+    # 3D Surface
+    surface = go.Surface(
+        x=X, y=Y, z=Z,
+        colorscale='viridis',
+        showscale=True,
+        colorbar=dict(x=0.45)
+    )
+    fig.add_trace(surface, row=1, col=1)
+
+    # Contour Plot
+    contour = go.Contour(
+        x=x, y=y, z=Z,
+        colorscale='viridis',
+        showscale=True,
+        colorbar=dict(x=1.0),
+        contours=dict(
+            showlabels=True,
+            labelfont=dict(size=12)
+        )
+    )
+    fig.add_trace(contour, row=1, col=2)
+
+    # 添加分量中心点
+    fig.add_trace(
+        go.Scatter(
+            x=centers[:K, 0], y=centers[:K, 1],
+            mode='markers+text',
+            marker=dict(size=10, color='red'),
+            text=[f'C{i+1}' for i in range(K)],
+            textposition="top center",
+            name='分量中心'
+        ),
+        row=1, col=2
+    )
+
+    # 添加采样点（如果有）
+    if samples is not None:
+        fig.add_trace(
+            go.Scatter(
+                x=samples[:, 0], y=samples[:, 1],
+                mode='markers+text',
+                marker=dict(
+                    size=8,
+                    color='yellow',
+                    line=dict(color='black', width=1)
+                ),
+                text=[f'S{i+1}' for i in range(len(samples))],
+                textposition="bottom center",
+                name='采样点'
+            ),
+            row=1, col=2
+        )
+
+    # 更新布局
+    fig.update_layout(
+        title='广义高斯混合分布',
+        showlegend=True,
+        width=1200,
+        height=600,
+        scene=dict(
+            xaxis_title='X',
+            yaxis_title='Y',
+            zaxis_title='密度'
+        )
+    )
+
+    # 更新2D图的坐标轴
+    fig.update_xaxes(title_text='X', row=1, col=2)
+    fig.update_yaxes(title_text='Y', row=1, col=2)
+
+    return fig
+
+def train_kan(samples, gmm_dataset, device='cuda'):
+    """训练KAN网络"""
+    if torch.cuda.is_available() and device == 'cuda':
+        device = torch.device('cuda')
+    else:
+        device = torch.device('cpu')
+
+    # 转换采样点为tensor
+    samples = torch.from_numpy(samples).to(device)
+    # 计算标签（概率密度值）
+    labels = torch.from_numpy(gmm_dataset.pdf(samples.cpu().numpy())).reshape(-1, 1).to(device)
+
+    # 创建KAN模型
+    model = KAN(width=[2,5,1], grid=3, k=3, seed=42, device=device)
+    # 创建训练和测试数据集
+    train_size = int(0.8 * samples.shape[0])
+    train_dataset = {
+        'train_input': samples[:train_size],
+        'train_label': labels[:train_size],
+        'test_input': samples[train_size:],
+        'test_label': labels[train_size:]
+    }
+
+    # 创建训练进度显示组件
+    st.write("网络结构：")
+    kan_fig_placeholder = st.empty()
+    st.write("预测结果：")
+    kan_plot_placeholder = st.empty()
+    progress_container = st.container()
+
+    total_steps = 100
+    steps_per_update = 10
+
+    def calculate_error(model, x, y):
+        """计算预测误差"""
+        with torch.no_grad():
+            pred = model(x)
+            return torch.mean((pred - y) ** 2).item()
+    
+    def train_phase(phase_name, steps, lamb=None, show_plot=True):
+        with progress_container:
+            progress_bar = st.progress(0)
+            status_text = st.empty()
+            
+            for step in range(0, steps, steps_per_update):
+                # 训练几步
+                if lamb is not None:
+                    model.fit(train_dataset, opt="LBFGS", steps=steps_per_update, lamb=lamb)
+                else:
+                    model.fit(train_dataset, opt="LBFGS", steps=steps_per_update)
+                
+                # 更新进度和误差
+                progress = (step + steps_per_update) / steps
+                progress_bar.progress(progress)
+                
+                # 计算当前误差
+                train_error = calculate_error(model, train_dataset['train_input'], train_dataset['train_label'])
+                test_error = calculate_error(model, train_dataset['test_input'], train_dataset['test_label'])
+                # 使用表格格式显示进度和误差
+                status_text.markdown(f"""
+                ### {phase_name}
+                | 项目 | 值 |
+                |:---|:---|
+                | 进度 | {progress:.0%} |
+                | 训练误差 | {train_error:.8f} |
+                | 测试误差 | {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 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)})")
+
+    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("正在进行网络剪枝优化..."):
+        model = model.prune()
+        progress_container.info("网络剪枝完成")
+    
+    # 第二阶段：精调
+    with st.spinner("最终调优阶段..."):
+        train_phase("第二阶段", total_steps, show_plot=True)  # 第二阶段显示网络图
+    
+    # 显示最终误差
+    train_error = calculate_error(model, train_dataset['train_input'], train_dataset['train_label'])
+    test_error = calculate_error(model, train_dataset['test_input'], train_dataset['test_label'])
+    error_text.markdown(f"""
+    #### 训练结果
+    - 训练集误差: {train_error:.6f}
+    - 测试集误差: {test_error:.6f}
+    """)
+    
+    progress_container.success("🎉 训练完成！")
+    return model
 
 def init_session_state():
     """初始化session state"""
@@ -20,6 +325,8 @@ def init_session_state():
         st.session_state.weights = np.ones(3, dtype=np.float64) / 3
     if 'sample_points' not in st.session_state:
         st.session_state.sample_points = None
+    if 'kan_model' not in st.session_state:
+        st.session_state.kan_model = None
 
 def create_default_parameters(K: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """创建默认参数"""
@@ -45,7 +352,7 @@ def generate_latex_formula(p: float, K: int, centers: np.ndarray,
         c = centers[k]
         s = scales[k]
         w = weights[k]
-        component = f"P_{k+1}(x) = \\frac{{{p:.1f}}}{{2\\alpha_{k+1} \\Gamma(1/{p:.1f})}}\\exp(-|\\frac{{x-({c[0]:.1f}, {c[1]:.1f})}}{{{s[0]:.1f}, {s[1]:.1f}}}|^{{{p:.1f}}}) \\\\"
+        component = f"P_{{\\theta_{k+1}}}(x) = \\frac{{{p:.1f}}}{{2\\alpha_{k+1} \\Gamma(1/{p:.1f})}}\\exp(-|\\frac{{x-({c[0]:.1f}, {c[1]:.1f})}}{{{s[0]:.1f}, {s[1]:.1f}}}|^{{{p:.1f}}}) \\\\"
         formula += component
         formula += f"\\pi_{k+1} = {w:.2f} \\\\"
     
@@ -116,15 +423,19 @@ with st.sidebar:
     st.subheader("采样设置")
     n_samples = st.slider("采样点数", 5, 20, 10)
     if st.button("重新采样"):
-        # 生成随机样本
-        samples = []
-        for _ in range(n_samples):
-            # 选择分量
-            k = np.random.choice(K, p=weights)
-            # 从选定的分量生成样本
-            sample = np.random.normal(centers[k], scales[k], size=2)
-            samples.append(sample)
-        st.session_state.sample_points = np.array(samples)
+        # 创建GMM数据集进行采样
+        gmm = GeneralizedGaussianMixture(
+            D=2,
+            K=K,
+            p=st.session_state.p,
+            centers=centers[:K],
+            scales=scales[:K],
+            weights=weights[:K]
+        )
+        # 使用GMM生成采样点
+        samples, _ = gmm.generate_samples(n_samples)
+        st.session_state.sample_points = samples
+        st.session_state.kan_model = None  # 重置KAN模型
 
 # 创建GMM数据集
 dataset = GeneralizedGaussianMixture(
@@ -196,7 +507,7 @@ if st.session_state.sample_points is not None:
     posteriors = []
     for sample in samples:
         component_probs = [
-            weights[k] * np.exp(-np.sum(((sample - centers[k]) / scales[k])**st.session_state.p)) 
+            weights[k] * np.exp(-np.sum(((sample - centers[k]) / scales[k])**st.session_state.p))
             for k in range(K)
         ]
         total = sum(component_probs)
@@ -218,16 +529,6 @@ if st.session_state.sample_points is not None:
         ),
         row=1, col=2
     )
-    
-    # 显示样本点的概率信息
-    st.subheader("采样点信息")
-    for i, (sample, prob, post) in enumerate(zip(samples, probs, posteriors)):
-        st.write(f"样本点 S{i+1} ({sample[0]:.2f}, {sample[1]:.2f}):")
-        st.write(f"- 概率密度: {prob:.4f}")
-        st.write("- 后验概率:")
-        for k in range(K):
-            st.write(f"  - 分量 {k+1}: {post[k]:.4f}")
-        st.write("---")
 
 # 更新布局
 fig.update_layout(
@@ -246,9 +547,70 @@ fig.update_layout(
 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)
 
+# KAN网络训练和预测部分
+if st.session_state.sample_points is not None:
+    st.markdown("---")
+    st.subheader("KAN网络训练与预测")
+    
+    # 训练控制按钮
+    col1, col2, col3 = st.columns([1, 2, 1])
+    with col1:
+        if st.button("拟合KAN", use_container_width=True):
+            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):
+                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)
+
+    st.markdown("---")
+
+# 显示采样点信息
+if st.session_state.sample_points is not None:
+    # 重新计算采样点的概率密度和后验概率
+    samples = st.session_state.sample_points
+    probs = dataset.pdf(samples)
+    posteriors = []
+    for sample in samples:
+        component_probs = [
+            weights[k] * np.exp(-np.sum(((sample - centers[k]) / scales[k])**st.session_state.p))
+            for k in range(K)
+        ]
+        total = sum(component_probs)
+        posteriors.append([p/total for p in component_probs])
+        
+    with st.expander("采样点信息"):
+        # 创建数据列表
+        point_data = []
+        for i, (sample, prob, post) in enumerate(zip(samples, probs, posteriors)):
+            row = {
+                '采样点': f'S{i+1}',
+                'X坐标': f'{sample[0]:.2f}',
+                'Y坐标': f'{sample[1]:.2f}',
+                '概率密度': f'{prob:.4f}'
+            }
+            # 添加每个分量的后验概率
+            for k in range(K):
+                row[f'分量{k+1}后验概率'] = f'{post[k]:.4f}'
+            point_data.append(row)
+        
+        # 显示dataframe
+        st.dataframe(point_data)
+
 # 添加参数说明
 with st.expander("分布参数说明"):
     st.markdown("""