如何为带更新菜单的Plotly散点图添加OLS趋势线并更新信息

问题根源 #

你遇到的核心问题是：初始代码仅基于一组X/Y变量生成趋势线，下拉切换时仅修改散点的数据源，未重新计算对应变量的OLS拟合结果，导致趋势线数据不匹配、悬停模板未更新、拟合方程固定。

修正代码 #

import pandas as pd
import plotly.graph_objects as go
import numpy as np
from sklearn.linear_model import LinearRegression

# 测试数据（可替换为你的数据）
np.random.seed(42)
df = pd.DataFrame({
    'X1': np.linspace(0, 10, 50),
    'X2': np.linspace(5, 15, 50),
    'Y1': 2*np.linspace(0,10,50) + np.random.normal(0,1,50),
    'Y2': 3*np.linspace(5,15,50) + np.random.normal(0,2,50)
})

# 定义所有需要切换的X/Y变量组合
var_combinations = [
    {'x': 'X1', 'y': 'Y1'},
    {'x': 'X1', 'y': 'Y2'},
    {'x': 'X2', 'y': 'Y1'},
    {'x': 'X2', 'y': 'Y2'}
]

fig = go.Figure()

# 为每个变量组合生成独立的散点和趋势线trace
for idx, combo in enumerate(var_combinations):
    x_col = combo['x']
    y_col = combo['y']
    
    # 添加散点trace
    fig.add_trace(go.Scatter(
        x=df[x_col],
        y=df[y_col],
        mode='markers',
        name=f'{x_col} vs {y_col} (散点)',
        hovertemplate=f'{x_col}: %{{x}}<br>{y_col}: %{{y}}<extra></extra>',
        visible=(idx == 0)  # 初始仅显示第一个组合
    ))
    
    # 计算当前组合的OLS趋势线
    X = df[x_col].values.reshape(-1,1)
    y = df[y_col].values
    model = LinearRegression().fit(X, y)
    y_pred = model.predict(X)
    fit_eq = f'y = {model.coef_[0]:.2f}x + {model.intercept_:.2f}<br>R² = {model.score(X,y):.2f}'
    
    # 添加趋势线trace
    fig.add_trace(go.Scatter(
        x=df[x_col],
        y=y_pred,
        mode='lines',
        name=f'{x_col} vs {y_col} (趋势线)',
        hovertemplate=fit_eq + '<extra></extra>',
        visible=(idx == 0)
    ))

# 配置下拉菜单
dropdown_buttons = []
for idx, combo in enumerate(var_combinations):
    x_col = combo['x']
    y_col = combo['y']
    # 控制对应组合的散点和趋势线可见（每个组合对应2个trace：偶数索引为散点，奇数为趋势线）
    visibility = [False] * len(fig.data)
    visibility[2*idx] = True
    visibility[2*idx + 1] = True
    
    dropdown_buttons.append({
        'label': f'{x_col} vs {y_col}',
        'method': 'update',
        'args': [
            {'visible': visibility},
            {'title': f'散点图与OLS趋势线: {x_col} vs {y_col}'}
        ]
    })

# 布局设置
fig.update_layout(
    updatemenus=[{
        'buttons': dropdown_buttons,
        'direction': 'down',
        'showactive': True,
        'x': 0.1,
        'y': 1.15
    }],
    title='散点图与OLS趋势线',
    xaxis_title='X变量',
    yaxis_title='Y变量'
)

fig.show()

关键改动说明 #

1. 预生成所有变量组合的trace：每个X/Y组合对应独立的散点和趋势线trace，下拉切换时仅控制trace的可见性，避免修改数据源导致的趋势线错位。
2. 动态计算OLS拟合：为每个变量组合单独拟合线性模型，生成对应的趋势线数据和拟合方程，确保切换后趋势线是当前变量的正确拟合结果。
3. 自定义悬停模板：为每个散点trace设置对应变量的悬停文本，趋势线直接显示拟合方程，切换后悬停内容自动更新。
4. 下拉菜单精准控制可见性：每个下拉选项对应激活一组散点+趋势线trace，其余trace隐藏，保证界面仅显示当前选中的变量组合图形。

内容的提问来源于stack exchange，提问作者MartyMcFly