Nature期刊图复现 | Python 实现双估计量对比森林图

import osimport numpy as npimport pandas as pdimport matplotlib.pyplot as pltoutput_dir = "图表"os.makedirs(output_dir, exist_ok=True)excel_path = "data.xlsx"if not os.path.exists(excel_path):    raise FileNotFoundError(f"未在当前目录下找到 {excel_path} 文件。")df = pd.read_excel(excel_path)if 'ci_lower' not in df.columns or 'ci_upper' not in df.columns:    if 'ci' in df.columns:        ci_parsed = df['ci'].astype(str).str.replace(r'[\[\]\s]', '', regex=True).str.split(',', expand=True)        df['ci_lower'] = pd.to_numeric(ci_parsed[0])        df['ci_upper'] = pd.to_numeric(ci_parsed[1])    else:        raise ValueError("Excel 数据中必须包含 'ci_lower'/'ci_upper' 列或 'ci' 范围列。")color_themes = {    1: {'PATE': '#0022cc', 'SATE': '#f59e24', 'note': '经典蓝橙：原图配色，对比鲜明，重点突出'},    2: {'PATE': '#1f77b4', 'SATE': '#ff7f0e', 'note': '标准通用：Matplotlib 默认类别色彩，兼容性极高'},    3: {'PATE': '#2b5c8f', 'SATE': '#d95f02', 'note': '学术沉稳：深邃蓝与低饱和度暗橙，适合正式论文'},    4: {'PATE': '#008080', 'SATE': '#cd5c5c', 'note': '自然冷暖：绿松石色配印第安红，沉静而不失活力'},    5: {'PATE': '#4a148c', 'SATE': '#00c853', 'note': '高亮科技：深紫与翠绿搭配，极具未来感与对比度'},    6: {'PATE': '#2c3e50', 'SATE': '#e74c3c', 'note': '平铺商务：经典浅黑灰配高亮红，商业报告常用'},    7: {'PATE': '#0f4c81', 'SATE': '#ff6f61', 'note': '时尚活力：经典蓝与活珊瑚橘，兼具专业与现代感'},    8: {'PATE': '#1b4d3e', 'SATE': '#c19a6b', 'note': '英伦复古：不列颠绿与优雅驼色，整体质感偏向温润'},    9: {'PATE': '#3f51b5', 'SATE': '#ff4081', 'note': '材料设计：标准靛蓝与亮丽粉红，视觉冲击力强'},    10: {'PATE': '#333333', 'SATE': '#888888', 'note': '工业纯灰：纯黑与深灰组合，适合极简或单色印刷'},    11: {'PATE': '#4682b4', 'SATE': '#b22222', 'note': '传统学院：钢青色与耐火砖红，北美高校论文常见'},    12: {'PATE': '#4a69bd', 'SATE': '#e55039', 'note': '温和莫兰迪：低饱和度的灰蓝与烟熏红，柔和不刺眼'},    13: {'PATE': '#1e3799', 'SATE': '#f6b93b', 'note': '皇家典雅：深海皇室蓝配亮金黄，对比尊贵显眼'},    14: {'PATE': '#0a3d62', 'SATE': '#3c6382', 'note': '渐变单色：同轴深蓝色与中蓝色，适合强调层次而非对立'},    15: {'PATE': '#228b22', 'SATE': '#8b0000', 'note': '圣诞复古：森林深绿与经典暗红，色彩饱满沉稳'},    16: {'PATE': '#6c5ce7', 'SATE': '#fd79a8', 'note': '赛博霓虹：亮紫色与极客粉红，偏向年轻化数据视觉'},    17: {'PATE': '#2d3436', 'SATE': '#fdcb6e', 'note': '黑金质感：高纯度深碳灰与黄铜金，低调奢华风格'},    18: {'PATE': '#005b41', 'SATE': '#008170', 'note': '墨绿双色：深邃暗绿与松石绿渐变，具有极佳的结构美感'},    19: {'PATE': '#706fd3', 'SATE': '#ff5252', 'note': '温和冲突：淡紫色与高光浅红，突出局部关键变动'},    20: {'PATE': '#2f3542', 'SATE': '#747d8c', 'note': '现代钢印：硬派深铁灰与亮银灰，适用于中立客观的展示'}}theme_choice = 1selected_colors = color_themes[theme_choice]plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Arial', 'SimHei']plt.rcParams['axes.unicode_minus'] = Falselabels = list(df['dv_label'].unique())[::-1]y_positions = np.arange(len(labels))fig, ax = plt.subplots(figsize=(10, 6.5), dpi=300)offset = 0.15ax.grid(axis='x', color='#e5e5e5', linestyle='-', linewidth=0.8, zorder=0)for i, label in enumerate(labels):    row_pate = df[(df['dv_label'] == label) & (df['est'] == 'PATE')]    if not row_pate.empty:        ate = row_pate['ate'].values[0]        ci_l = row_pate['ci_lower'].values[0]        ci_u = row_pate['ci_upper'].values[0]        xerr = [[ate - ci_l], [ci_u - ate]]        ax.errorbar(ate, i - offset, xerr=xerr, fmt='o', color=selected_colors['PATE'],                    markersize=6, capsize=0, linewidth=1.5, zorder=3)    row_sate = df[(df['dv_label'] == label) & (df['est'] == 'SATE')]    if not row_sate.empty:        ate = row_sate['ate'].values[0]        ci_l = row_sate['ci_lower'].values[0]        ci_u = row_sate['ci_upper'].values[0]        xerr = [[ate - ci_l], [ci_u - ate]]        ax.errorbar(ate, i + offset, xerr=xerr, fmt='o', color=selected_colors['SATE'],                    markersize=6, capsize=0, linewidth=1.5, zorder=3)ax.axvline(x=0, color='black', linestyle='-', linewidth=1.8, zorder=2)ax.set_yticks(y_positions)ax.set_yticklabels(labels, fontsize=11)ax.set_xticks([-0.05, 0.00, 0.05, 0.10])ax.set_xticklabels(['-0.05', '0.00', '0.05', '0.10'], fontsize=11)ax.set_xlabel('Treatment effect', fontsize=13, labelpad=12)ax.set_xlim(-0.06, 0.11)for spine in ['top', 'right', 'left', 'bottom']:    ax.spines[spine].set_visible(False)ax.tick_params(axis='both', which='both', length=0)blank_leg = ax.errorbar([], [], yerr=[], fmt=' ', color='none')pate_leg = ax.errorbar([], [], yerr=[], fmt='o', color=selected_colors['PATE'], markersize=6, linewidth=1.5)sate_leg = ax.errorbar([], [], yerr=[], fmt='o', color=selected_colors['SATE'], markersize=6, linewidth=1.5)legend_handles = [blank_leg, pate_leg, sate_leg]legend_labels = ['Estimand', 'PATE', 'SATE']leg = ax.legend(handles=legend_handles, labels=legend_labels, fontsize=11,                loc='upper center', bbox_to_anchor=(0.65, -0.15), ncol=3, frameon=False,                handlelength=0, handleheight=2.0, numpoints=1, columnspacing=1.8)leg.get_texts()[0].set_weight('bold')leg.legend_handles[0].set_visible(False)plt.tight_layout()output_image_path = os.path.join(output_dir, "forest_plot.png")plt.savefig(output_image_path, bbox_inches='tight')plt.close()print(f"File saved to: {output_image_path}")