如何在Python中为10条等高线添加WD1-WD10自定义标签

问题描述 #

需要给等高线图中的10条等高线分别标注WD1至WD10，目前仅实现了两条等高线的标签添加，现寻求批量实现的方法。现有示例代码如下：

import numpy as np
import matplotlib.pyplot as plt

def C2thin1(mx_mesh, sigma_mesh):
    return Cthin1(np.column_stack((mx_mesh.ravel(), sigma_mesh.ravel()))).reshape(mx_mesh.shape)

def C2thin2(mx_mesh, sigma_mesh):
    return Cthin2(np.column_stack((mx_mesh.ravel(), sigma_mesh.ravel()))).reshape(mx_mesh.shape)

def cage1_contour(C2thin_func, mx_mesh, sigma_mesh, Rstar, dstar):
    return e2dPhidE(C2thin_func, mx_mesh, sigma_mesh, Rstar, dstar)

# 定义mx和sigma的取值范围
mx_values = np.logspace(np.log10(0.01), np.log10(1000), 100)
sigma_values = np.logspace(np.log10(1e-30), np.log10(1e-50), 100)

# 创建绘图用的网格
mx_mesh, sigma_mesh = np.meshgrid(mx_values, sigma_values)

# 计算两条等高线的数值
contour_values1 = cage1_contour(C2thin1, mx_mesh, sigma_mesh, Rstar1, dstar1)
contour_values2 = cage1_contour(C2thin2, mx_mesh, sigma_mesh, Rstar2, dstar2)
# 生成等高线的判定条件
condition_values1 = sen9int(mx_mesh) - contour_values1
condition_values2 = sen9int(mx_mesh) - contour_values2

# 绘制等高线
plt.figure(figsize=(10, 10))
contour1 = plt.contour(mx_mesh, sigma_mesh, condition_values1, levels=[0], colors='purple', linewidths=3, label='WD1')
contour2 = plt.contour(mx_mesh, sigma_mesh, condition_values2, levels=[0], colors='blue', linewidths=3, label='WD2')

plt.xscale('log')
plt.yscale('log')
plt.xlabel(r'$m_{\chi}$ (GeV)', fontsize=24, fontweight='bold')
plt.ylabel(r'$\sigma_{\chi n}$ (cm$^2$)', fontsize=24, fontweight='bold')
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.grid(True, which='both', linestyle='--', linewidth=2, color='black')
plt.axhline(y=1e-41, color='black', linestyle='--', linewidth=2)
plt.axvline(x=100.4, color='black', linestyle='--', linewidth=2)
plt.title('Contour Plot', fontsize=24, fontweight='bold')
plt.legend(fontsize=18)
plt.show()

解决方案 #

通过循环批量处理可避免重复编写冗余代码，同时确保每条等高线对应正确的标签，以下是两种常用实现方式：

方式1：通过图例展示标签 #

将所有需要用到的函数、参数整理为列表，循环遍历生成等高线并设置对应标签：

import numpy as np
import matplotlib.pyplot as plt

# 整理10组对应的函数和参数（需确保已提前定义好C2thin1-C2thin10、Rstar1-Rstar10、dstar1-dstar10）
c2thin_funcs = [C2thin1, C2thin2, C2thin3, C2thin4, C2thin5,
                C2thin6, C2thin7, C2thin8, C2thin9, C2thin10]
rstar_list = [Rstar1, Rstar2, Rstar3, Rstar4, Rstar5,
              Rstar6, Rstar7, Rstar8, Rstar9, Rstar10]
dstar_list = [dstar1, dstar2, dstar3, dstar4, dstar5,
              dstar6, dstar7, dstar8, dstar9, dstar10]
# 自定义10种区分度高的颜色
colors = ['purple', 'blue', 'green', 'red', 'orange',
          'cyan', 'magenta', 'brown', 'gray', 'olive']

# 生成网格数据
mx_values = np.logspace(np.log10(0.01), np.log10(1000), 100)
sigma_values = np.logspace(np.log10(1e-30), np.log10(1e-50), 100)
mx_mesh, sigma_mesh = np.meshgrid(mx_values, sigma_values)

plt.figure(figsize=(10, 10))

# 循环绘制10条等高线
for idx in range(10):
    # 计算当前等高线的数值
    contour_vals = cage1_contour(c2thin_funcs[idx], mx_mesh, sigma_mesh, rstar_list[idx], dstar_list[idx])
    condition_vals = sen9int(mx_mesh) - contour_vals
    # 绘制等高线并设置标签
    plt.contour(mx_mesh, sigma_mesh, condition_vals, levels=[0], 
                colors=colors[idx], linewidths=3, label=f'WD{idx+1}')

# 保持原有坐标轴、标题等设置
plt.xscale('log')
plt.yscale('log')
plt.xlabel(r'$m_{\chi}$ (GeV)', fontsize=24, fontweight='bold')
plt.ylabel(r'$\sigma_{\chi n}$ (cm$^2$)', fontsize=24, fontweight='bold')
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.grid(True, which='both', linestyle='--', linewidth=2, color='black')
plt.axhline(y=1e-41, color='black', linestyle='--', linewidth=2)
plt.axvline(x=100.4, color='black', linestyle='--', linewidth=2)
plt.title('Contour Plot', fontsize=24, fontweight='bold')
plt.legend(fontsize=18)
plt.show()

方式2：直接在等高线线条上标注标签 #

如果希望把WD1等标签直接显示在对应的等高线上，可使用clabel方法：

import numpy as np
import matplotlib.pyplot as plt

# 整理函数和参数列表
c2thin_funcs = [C2thin1, C2thin2, C2thin3, C2thin4, C2thin5,
                C2thin6, C2thin7, C2thin8, C2thin9, C2thin10]
rstar_list = [Rstar1, Rstar2, Rstar3, Rstar4, Rstar5,
              Rstar6, Rstar7, Rstar8, Rstar9, Rstar10]
dstar_list = [dstar1, dstar2, dstar3, dstar4, dstar5,
              dstar6, dstar7, dstar8, dstar9, dstar10]
colors = ['purple', 'blue', 'green', 'red', 'orange',
          'cyan', 'magenta', 'brown', 'gray', 'olive']

# 生成网格数据
mx_values = np.logspace(np.log10(0.01), np.log10(1000), 100)
sigma_values = np.logspace(np.log10(1e-30), np.log10(1e-50), 100)
mx_mesh, sigma_mesh = np.meshgrid(mx_values, sigma_values)

plt.figure(figsize=(10, 10))

# 循环绘制并标注等高线
for idx in range(10):
    contour_vals = cage1_contour(c2thin_funcs[idx], mx_mesh, sigma_mesh, rstar_list[idx], dstar_list[idx])
    condition_vals = sen9int(mx_mesh) - contour_vals
    # 绘制等高线
    contour = plt.contour(mx_mesh, sigma_mesh, condition_vals, levels=[0], 
                          colors=colors[idx], linewidths=3)
    # 直接在等高线上添加标签
    plt.clabel(contour, inline=True, fontsize=16, fmt=f'WD{idx+1}')

# 其他绘图设置保持不变
plt.xscale('log')
plt.yscale('log')
plt.xlabel(r'$m_{\chi}$ (GeV)', fontsize=24, fontweight='bold')
plt.ylabel(r'$\sigma_{\chi n}$ (cm$^2$)', fontsize=24, fontweight='bold')
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.grid(True, which='both', linestyle='--', linewidth=2, color='black')
plt.axhline(y=1e-41, color='black', linestyle='--', linewidth=2)
plt.axvline(x=100.4, color='black', linestyle='--', linewidth=2)
plt.title('Contour Plot', fontsize=24, fontweight='bold')
plt.show()

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