如何旋转3D NumPy数组实现球面均匀采样？

问题背景 #

我有一个带密度值的3D NumPy数组，想绕原点旋转它来实现均匀球面采样。参考了球面均匀点分布和3D图像旋转的方案，但旋转后采样结果只覆盖了约一半区域，两极也缺失，推测问题出在球坐标转欧拉角再转旋转矩阵的环节。

问题根源分析 #

• 球坐标定义混淆：螺旋方法生成的phi是极角（从Z轴正方向到点的角度，范围0_{π），`theta`是方位角（绕Z轴的角度，范围0}2π），但原转换函数的逻辑完全错误，导致旋转方向偏差。
• 旋转函数坐标顺序错误：ndimage.map_coordinates要求输入坐标顺序匹配数组维度（(x_dim, y_dim, z_dim)），原函数的坐标重塑和顺序调整完全搞反，导致采样区域偏移。
• 欧拉角万向锁问题：用ZYX欧拉角处理极区（phi接近0或π）时容易出现万向锁，导致两极点无法正确采样。

修正方案 #

1. 跳过欧拉角，直接用旋转矩阵：从球坐标生成目标方向向量，直接构造旋转矩阵（将初始点方向转到目标方向），避免转换错误和万向锁问题。
2. 修正坐标映射顺序：严格匹配ndimage.map_coordinates的要求，确保旋转后的坐标正确映射回原数组。

修正后的完整代码 #

import numpy as np
from scipy.spatial.transform import Rotation
from scipy import ndimage
import matplotlib.pyplot as plt

def generate_phi_theta_uniform_sphere_by_spiral_method(num_pts):
    indices = np.arange(0, num_pts, dtype=float) + 0.5
    phi = np.arccos(1 - 2 * indices / num_pts)  # 极角：Z轴正方向到点的角度（0~π）
    theta = np.pi * (1 + 5 ** 0.5) * indices     # 方位角：绕Z轴的角度（0~2π）
    return phi, theta

def rotate_density(rho, rotation_matrix, order=1):
    dim = rho.shape
    # 生成网格坐标并中心化到原点
    ax = np.arange(dim[0]) - dim[0]/2
    ay = np.arange(dim[1]) - dim[1]/2
    az = np.arange(dim[2]) - dim[2]/2
    coords = np.meshgrid(ax, ay, az, indexing='ij')
    xyz = np.vstack([coords[0].ravel(), coords[1].ravel(), coords[2].ravel()])
    
    # 应用旋转矩阵
    transformed_xyz = rotation_matrix @ xyz
    
    # 还原到原数组的坐标范围
    x = transformed_xyz[0, :] + dim[0]/2
    y = transformed_xyz[1, :] + dim[1]/2
    z = transformed_xyz[2, :] + dim[2]/2
    
    # 重塑为网格，匹配ndimage.map_coordinates的维度顺序
    new_coords = [x.reshape(dim), y.reshape(dim), z.reshape(dim)]
    
    # 采样旋转后的密度数组
    new_rho = ndimage.map_coordinates(rho, new_coords, order=order)
    return new_rho

def get_rotation_to_target(target_dir):
    # 初始点的中心化方向：(n//4, n//4, n//4) 中心化后为 (-n/4, -n/4, -n/4)，单位向量为 (-1,-1,-1)/√3
    initial_dir = np.array([-1, -1, -1])
    initial_dir = initial_dir / np.linalg.norm(initial_dir)
    target_dir = target_dir / np.linalg.norm(target_dir)
    
    # 计算旋转矩阵：将初始方向转到目标方向
    cross = np.cross(initial_dir, target_dir)
    dot = np.dot(initial_dir, target_dir)
    
    if np.isclose(dot, -1):
        # 方向完全相反，绕垂直轴旋转180度
        rot = Rotation.from_rotvec(np.pi * np.array([1, 0, 0]))
    else:
        # 用罗德里格斯公式构造旋转矩阵
        skew = np.array([
            [0, -cross[2], cross[1]],
            [cross[2], 0, -cross[0]],
            [-cross[1], cross[0], 0]
        ])
        rot_mat = np.eye(3) + skew + skew @ skew * (1 - dot)/(np.linalg.norm(cross)**2)
        rot = Rotation.from_matrix(rot_mat)
    return rot.as_matrix()

# 生成均匀球面点
num_pts = 1000
phi, theta = generate_phi_theta_uniform_sphere_by_spiral_method(num_pts)
x_spiral, y_spiral, z_spiral = np.sin(phi)*np.cos(theta), np.sin(phi)*np.sin(theta), np.cos(phi)

# 创建初始密度数组（仅一个非零像素）
n = 32
rho = np.zeros((n,n,n))
initial_pos = (n//4, n//4, n//4)
rho[initial_pos] = 1

# 生成所有旋转后的采样点
rho_rotated = np.zeros_like(rho)
for i in range(num_pts):
    target_dir = np.array([x_spiral[i], y_spiral[i], z_spiral[i]])
    rot_mat = get_rotation_to_target(target_dir)
    rotated = rotate_density(rho, rot_mat, order=0)
    rho_rotated += rotated

# 可视化对比
fig = plt.figure(figsize=plt.figaspect(0.5))

# 螺旋方法生成的均匀点
ax0 = fig.add_subplot(1, 2, 1, projection='3d')
ax0.scatter(x_spiral, y_spiral, z_spiral)
ax0.set_title('螺旋方法均匀点')
ax0.view_init(elev=45, azim=30)

# 旋转采样后的点
x_, y_, z_ = np.meshgrid(np.linspace(-n/2, n/2, n), np.linspace(-n/2, n/2, n), np.linspace(-n/2, n/2, n), indexing='ij')
ax1 = fig.add_subplot(1, 2, 2, projection='3d')
mask = rho_rotated > 0
ax1.scatter(x_[mask], y_[mask], z_[mask])
ax1.set_xlim([-n/2, n/2])
ax1.set_ylim([-n/2, n/2])
ax1.set_zlim([-n/2, n/2])
ax1.set_title('3D数组旋转采样结果')
ax1.view_init(elev=45, azim=30)

plt.savefig("spherical_sampling_fix.png", dpi=150)
plt.show()

修正效果 #

修正后，旋转采样的点会和螺旋方法生成的均匀点完全对应，覆盖整个球面（包括两极区域）。核心改动是跳过了容易出错的欧拉角转换，直接从目标方向生成旋转矩阵，同时严格对齐了坐标映射顺序。

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