如何为坐标变换后的角度笛卡尔网格填充对应颜色映射单元格

需要在两个坐标系间完成变换并绘制填充单元格，已成功创建填充矩形网格，但无法绘制角度变换后的网格。需求是：用矩形网格的颜色为笛卡尔坐标系下的角度网格填充颜色，尝试过plt.fill()和plt.contour()均未实现目标。

原始代码：

import numpy as np
import matplotlib.pyplot as plt

z_min = 10
z_max = 100
n_zcells = 10
z_array = np.linspace(z_min, z_max, n_zcells)

u_min = -9.5
u_max = 9.5
n_ucells = 15
u_array = np.linspace(u_min, u_max, n_ucells)

[u_grid, z_grid] = np.meshgrid(u_array, z_array)

# 在笛卡尔坐标系中表示极坐标网格
f = 5   # 像素焦距（单位：像素）
B = 10  # 基线长度（单位：cm）

# u = fX/Z；相机到像素坐标的转换公式
x_grid = u_grid*z_grid/f

# 目标网格点
plt.figure()
plt.subplot(2,1,1)
plt.scatter(x_grid, z_grid)
plt.xlabel('横向距离 (cm)')
plt.ylabel('深度 (cm)')
plt.subplot(2,1,2)
plt.scatter(u_grid, z_grid) 
plt.xlabel('像素列 ')
plt.ylabel('深度 (cm)')
plt.show()

# 填充网格单元格颜色
filled_colours = np.random.rand(n_zcells-1,n_ucells-1) # 单元格数量比网格点少1
plt.imshow(filled_colours, cmap=plt.get_cmap('gray'), origin='lower')

核心思路是：将u-z坐标系中的每个矩形单元格转换为x-z坐标系中的四边形，再用PolyCollection批量绘制这些多边形并赋予对应颜色。

完整实现代码：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection

z_min = 10
z_max = 100
n_zcells = 10
z_array = np.linspace(z_min, z_max, n_zcells)

u_min = -9.5
u_max = 9.5
n_ucells = 15
u_array = np.linspace(u_min, u_max, n_ucells)

[u_grid, z_grid] = np.meshgrid(u_array, z_array)

# 坐标系转换参数
f = 5   # 像素焦距（单位：像素）
B = 10  # 基线长度（单位：cm）
x_grid = u_grid*z_grid/f

# 生成每个单元格的四个顶点坐标（u-z坐标系）
polygons = []
for i in range(n_zcells-1):
    for j in range(n_ucells-1):
        # 提取单元格四个角点的u、z值
        u_points = [u_grid[i,j], u_grid[i,j+1], u_grid[i+1,j+1], u_grid[i+1,j]]
        z_points = [z_grid[i,j], z_grid[i,j+1], z_grid[i+1,j+1], z_grid[i+1,j]]
        # 转换为x-z坐标系的顶点
        x_points = np.array(u_points) * np.array(z_points) / f
        # 整理为多边形格式（每个多边形是N个(xy)点的列表）
        poly = np.column_stack((x_points, z_points))
        polygons.append(poly)

# 生成单元格颜色
filled_colours = np.random.rand(n_zcells-1,n_ucells-1).flatten()

# 创建多边形集合并绘制对比图
fig, (ax1, ax2) = plt.subplots(2,1, figsize=(8,10))

# 绘制x-z坐标系的变换后网格（带填充色）
poly_coll = PolyCollection(polygons, cmap='gray', array=filled_colours)
ax1.add_collection(poly_coll)
ax1.scatter(x_grid, z_grid, s=5, color='black')  # 叠加网格点
ax1.set_xlabel('横向距离 (cm)')
ax1.set_ylabel('深度 (cm)')
ax1.set_xlim(x_grid.min(), x_grid.max())
ax1.set_ylim(z_grid.min(), z_grid.max())
ax1.set_aspect('auto')

# 绘制原始u-z坐标系的填充网格
ax2.imshow(filled_colours.reshape(n_zcells-1,n_ucells-1), 
           cmap='gray', origin='lower', 
           extent=[u_min, u_max, z_min, z_max])
ax2.scatter(u_grid, z_grid, s=5, color='black')
ax2.set_xlabel('像素列')
ax2.set_ylabel('深度 (cm)')

plt.tight_layout()
plt.show()

关键说明 #

• 多边形顶点生成：遍历每个u-z网格的单元格，提取四个角点后通过x = u*z/f转换到笛卡尔坐标系，组成四边形顶点集合。
• PolyCollection：该类可批量处理多个多边形，支持直接绑定颜色数组，完美适配网格单元格的颜色映射需求。
• 坐标范围适配：绘制时需手动设置x-z坐标系的轴范围，确保所有变换后的多边形都能显示在视图内。

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