如何计算PyTorch二维点集的求和分布热力图网格值？

我有一个形状为(10000,2)的PyTorch张量，包含10000个二维点，点的取值范围在-1到1之间（即$x ∈ [-1, 1]^2$）。我希望生成热力图展示这些点的分布，已经写出了部分matplotlib代码：

torch.meshgrid(torch.linspace(-1, 1, 1000), torch.linspace(-1, 1, 1000))
ax.contourf(x_grid, y_grid, grid_values)

但不知道如何计算对应的grid_values，请问如何基于该点张量计算用于展示点分布的grid_values？

解决方案 #

grid_values本质是二维直方图数据，即统计[-1,1]×[-1,1]平面上每个网格区域内的点数量（或自定义权重求和），可以用PyTorch内置的torch.histogram2d直接计算，步骤如下：

步骤1：拆分点张量的x、y分量 #

假设你的点张量名为points，先拆分出x和y坐标：

x = points[:, 0]
y = points[:, 1]

步骤2：计算二维直方图（grid_values） #

因为要生成1000×1000的网格，所以设置直方图的bin数量为1000，同时指定坐标范围为[-1,1]：

bin_count = 1000
# 生成bin的边界（n个bin需要n+1个边界点）
x_edges = torch.linspace(-1, 1, bin_count + 1)
y_edges = torch.linspace(-1, 1, bin_count + 1)

# 计算直方图，grid_values即为每个网格的点数量
grid_values, _, _ = torch.histogram2d(
    x, y,
    bins=(bin_count, bin_count),
    range=[[-1, 1], [-1, 1]]
)

步骤3：匹配网格与绘图 #

生成meshgrid时注意使用indexing='xy'（符合matplotlib的坐标逻辑），同时需要将grid_values转置以匹配网格的索引顺序：

x_grid, y_grid = torch.meshgrid(x_edges[:-1], y_edges[:-1], indexing='xy')

# 转置grid_values，确保和x_grid、y_grid的维度对应
ax.contourf(x_grid.numpy(), y_grid.numpy(), grid_values.numpy().T)

完整示例代码 #

import torch
import matplotlib.pyplot as plt

# 模拟输入点张量（替换为你的实际数据）
points = torch.rand(10000, 2) * 2 - 1  # 生成[-1,1]范围内的随机点

# 拆分坐标
x = points[:, 0]
y = points[:, 1]

# 计算直方图数据
bin_count = 1000
x_edges = torch.linspace(-1, 1, bin_count + 1)
y_edges = torch.linspace(-1, 1, bin_count + 1)
grid_values, _, _ = torch.histogram2d(x, y, bins=(bin_count, bin_count), range=[[-1, 1], [-1, 1]])

# 生成网格并绘图
x_grid, y_grid = torch.meshgrid(x_edges[:-1], y_edges[:-1], indexing='xy')
fig, ax = plt.subplots()
ax.contourf(x_grid.numpy(), y_grid.numpy(), grid_values.numpy().T)
plt.show()

扩展：自定义权重求和 #

如果需要展示的是点的某个属性的求和分布（而非点数量），可以传入weights参数：

# 假设每个点有一个权重张量weights，形状(10000,)
weights = torch.rand(10000)
grid_values, _, _ = torch.histogram2d(x, y, bins=(bin_count, bin_count), range=[[-1,1], [-1,1]], weights=weights)

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