如何围绕指定直线旋转2D numpy数组，使该直线成为新X轴？

问题背景 #

现有随机生成的2D NumPy数组（可视为图像），已通过高值点拟合出一条直线，需要旋转数组，让这条直线对齐到新的X轴。

步骤1：修正端点提取代码 #

原代码存在变量名笔误，先修正高值点的提取逻辑：

import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

# 生成原始数组
x = np.linspace(-5,5,8)
y = np.linspace(-25,25,7)
z = np.random.randn(7,8)*0.5
z = z/np.max(z)
z[5,6] = 2; z[1,1] = 2

# 提取高z值的点（阈值为均值+3倍标准差）
line_mask = np.where(z > z.mean() + 3*z.std(), 1, 0)
y_idx, x_idx = np.nonzero(line_mask)
# 拟合直线
fitline = stats.linregress(x=x_idx, y=y_idx)

步骤2：计算旋转角度 #

通过拟合直线的斜率，计算其与原始X轴的夹角（弧度制），后续需要反向旋转这个角度，让直线对齐新X轴：

# 计算直线与原始X轴的夹角
theta = np.arctan(fitline.slope)

步骤3：执行数组旋转 #

提供两种实现方式，按需选择：

方式一：使用scipy.ndimage.rotate（简单快捷） #

直接调用现成的旋转函数，指定旋转角度即可：

from scipy.ndimage import rotate

# 将弧度转为角度，反向旋转使拟合直线对齐X轴
rotation_angle = -theta * 180 / np.pi
# 旋转数组，保持原尺寸，采用双线性插值
rotated_z = rotate(z, rotation_angle, reshape=False, order=1)

# 可视化对比
plt.figure(figsize=(10,5))
plt.subplot(121)
plt.pcolormesh(x,y,z); plt.colorbar(); plt.title('原始数组')
plt.subplot(122)
plt.pcolormesh(x,y,rotated_z); plt.colorbar(); plt.title('旋转后数组')
plt.show()

方式二：手动坐标变换（自定义性强） #

如果需要更精细的控制（比如调整旋转中心、输出范围），可以手动构建旋转矩阵并完成坐标映射：

from scipy.ndimage import map_coordinates

# 选择旋转中心（这里用数组几何中心，也可以换成拟合点的质心）
center_y, center_x = np.array(z.shape) // 2
# 或者用拟合点的质心：
# center_x = np.mean(x_idx)
# center_y = np.mean(y_idx)

# 构建旋转矩阵（绕中心旋转-theta，使拟合直线对齐X轴）
rot_matrix = np.array([
    [np.cos(theta), -np.sin(theta)],
    [np.sin(theta), np.cos(theta)]
])

# 生成原始网格坐标
y_grid, x_grid = np.meshgrid(np.arange(z.shape[0]), np.arange(z.shape[1]), indexing='ij')
# 坐标平移→旋转→平移回原坐标系
coords = np.stack([x_grid - center_x, y_grid - center_y])
rotated_coords = rot_matrix @ coords.reshape(2, -1)
rotated_x = rotated_coords[0].reshape(z.shape) + center_x
rotated_y = rotated_coords[1].reshape(z.shape) + center_y

# 双线性插值获取旋转后的值
rotated_z = map_coordinates(z, [rotated_y, rotated_x], order=1)

# 可视化对比
plt.figure(figsize=(10,5))
plt.subplot(121)
plt.pcolormesh(x,y,z); plt.colorbar(); plt.title('原始数组')
plt.subplot(122)
plt.pcolormesh(x,y,rotated_z); plt.colorbar(); plt.title('旋转后数组')
plt.show()

关键细节说明 #

• 旋转中心：默认用数组几何中心，若想让拟合直线的中心保持在旋转后的中心，可替换为拟合点的质心坐标。
• 插值方式：order=1为双线性插值，平衡速度与精度；若需更高精度可选用order=3（双三次插值），但计算耗时会增加。
• 输出尺寸：rotate函数的reshape=False会保持原数组大小，旋转后边缘会出现空白（用0填充）；若想完整显示旋转后的图像，可设置reshape=True，此时输出数组尺寸会自动调整。

内容的提问来源于stack exchange，提问作者quics-ilver