如何提取2D热力图指定区域计数并绘制1D柱状图？

问题描述 #

我正在Python中尝试绘制热力图指定区域内计数的柱状图。已完成热力图的编写，也实现了为指定区域绘制边界的功能，但卡在提取指定区域计数并绘制柱状图的环节。该热力图是一个2D直方图，我希望将其选中区域转换为1D直方图，x轴与热力图x轴一致，y轴为计数（即热力图中颜色代表的数值）。

现有代码（含关键修正） #

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LogNorm
from matplotlib.patches import Polygon

df = pd.read_csv('adc_tdc_7_disk26.txt', sep=',', header=None)

x, y = df[0].to_numpy(), df[1].to_numpy()

fig, ax = plt.subplots()
counts, xedges, yedges = np.histogram2d(x, y, bins=1000)
counts = counts.T 
# 关键修正：添加extent参数，让轴坐标匹配原始数据范围，否则区域点判断会出错
im = ax.imshow(counts, cmap='magma_r', origin='lower', norm=LogNorm(),
               extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]])
fig.colorbar(im)
click_points = []

def onselect(event):
    click_points.append((event.xdata, event.ydata))
    # 20 clicks draws a boundary
    if len(click_points) >= 20:
        polygon_verts = click_points

        polygon_patch = plt.Polygon(polygon_verts, facecolor='none', edgecolor='r')

        ax.add_patch(polygon_patch)

        fig.canvas.draw()

        hist_points(polygon_patch, counts, xedges, yedges, x, y)
        click_points.clear()

解决方案 #

核心问题说明 #

原代码未设置imshow的extent参数，导致点击的坐标（bin索引）与原始数据值不在同一坐标系，区域内点的判断完全失效。修正后即可正常提取区域内数据，以下提供两种实现方式：

方法一：基于原始选中点统计（直观易实现）

直接提取区域内的原始x数据，用与原2D直方图相同的x bins重新统计计数，无需复用已有counts数组：

def hist_points(polygon_patch, counts, xedges, yedges, x, y):
    # 获取区域内点的掩码
    selected_mask = polygon_patch.contains_points(np.column_stack((x, y)))
    # 提取选中的x数据
    selected_x = x[selected_mask]
    # 用原2D直方图的x bins统计1D计数
    hist_counts, _ = np.histogram(selected_x, bins=xedges)
    
    # 绘制1D柱状图
    fig2, ax2 = plt.subplots()
    ax2.bar(xedges[:-1], hist_counts, width=np.diff(xedges), edgecolor='black', alpha=0.7)
    ax2.set_xlabel('X (same as heatmap)')
    ax2.set_ylabel('Count')
    ax2.set_title('1D Histogram of Selected Region')
    plt.show()

方法二：基于2D直方图的bin统计（高效适合大数据）

直接复用已生成的counts数组，判断每个bin是否在选中区域内，按x bin累加计数，避免重新遍历原始数据：

def hist_points(polygon_patch, counts, xedges, yedges, x, y):
    # 生成每个bin的中心坐标（用于判断是否在多边形内）
    x_centers = (xedges[:-1] + xedges[1:]) / 2
    y_centers = (yedges[:-1] + yedges[1:]) / 2
    xx, yy = np.meshgrid(x_centers, y_centers)
    bin_centers = np.column_stack((xx.ravel(), yy.ravel()))
    
    # 判断每个bin是否在选中区域内
    bin_in_polygon = polygon_patch.contains_points(bin_centers)
    bin_in_polygon = bin_in_polygon.reshape(counts.shape)
    
    # 按x bin累加区域内的计数
    hist_counts = counts[bin_in_polygon].reshape(-1, counts.shape[1]).sum(axis=0)
    
    # 绘制1D柱状图
    fig2, ax2 = plt.subplots()
    ax2.bar(xedges[:-1], hist_counts, width=np.diff(xedges), edgecolor='black', alpha=0.7)
    ax2.set_xlabel('X (same as heatmap)')
    ax2.set_ylabel('Count')
    ax2.set_title('1D Histogram of Selected Region')
    plt.show()

使用说明 #

1. 运行修正后的代码，在热力图上点击20次绘制区域边界
2. 完成点击后会自动弹出新窗口，显示对应区域的1D柱状图

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