如何将自定义函数生成的旋转曲面转换为600×600×600的NumPy数组？

你正在开展一个小型项目，目标是获取自定义函数生成的旋转曲面，已实现sin(x)的旋转曲面生成代码，现希望将该曲面转换为600×600×600的NumPy数组用于后续使用。

原曲面生成代码 #

import numpy as np
import matplotlib.pyplot as plt

np.seterr(divide='ignore', invalid='ignore')

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

x = np.linspace(0,np.pi,600) # 设置sin(x)的左右边界
y = np.linspace(0, 2*np.pi,600)
U,V = np.meshgrid(x,y)

def function(x): # 此处可传入自定义函数，当前示例为sin(x)
    return np.sin(x)

a = np.vectorize(function) # 向量化函数，避免长度为1的数组被转为Python标量的错误
b = a(x) # 对x数组逐元素应用函数

X = U
Y = b*np.cos(V)
Z = b*np.sin(V)

ax.set_xlabel('Y axis')
ax.set_ylabel('X axis')
ax.set_zlabel('Z axis')

ax.plot_surface(X, Y, Z)

plt.show()

实现曲面转600×600×600 NumPy数组 #

完全可以实现，核心思路是体素化：将3D空间中的曲面映射到离散的体素网格，数组元素用1标记曲面区域（或内部），0标记非曲面区域。以下是具体实现代码：

import numpy as np
from scipy.spatial import KDTree

# 1. 生成旋转曲面点云
x = np.linspace(0, np.pi, 600)
y = np.linspace(0, 2*np.pi, 600)
U, V = np.meshgrid(x, y)

def function(x):
    return np.sin(x)

# 利用numpy广播替代vectorize，更高效
b = function(x)
X = U
Y = b * np.cos(V)
Z = b * np.sin(V)

# 将曲面的所有点整理为点云数组
surface_points = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T

# 2. 定义目标体素数组尺寸
target_shape = (600, 600, 600)

# 3. 计算坐标映射参数：将曲面的连续坐标转换为体素索引
x_min, x_max = X.min(), X.max()
y_min, y_max = Y.min(), Y.max()
z_min, z_max = Z.min(), Z.max()

# 缩放因子：将原坐标范围映射到0~599的索引范围
x_scale = (target_shape[0]-1) / (x_max - x_min)
y_scale = (target_shape[1]-1) / (y_max - y_min)
z_scale = (target_shape[2]-1) / (z_max - z_min)

# 偏移量：确保最小坐标对应索引0
x_offset = -x_min * x_scale
y_offset = -y_min * y_scale
z_offset = -z_min * z_scale

# 4. 生成所有体素的中心坐标
xv, yv, zv = np.meshgrid(
    np.linspace(x_min, x_max, target_shape[0]),
    np.linspace(y_min, y_max, target_shape[1]),
    np.linspace(z_min, z_max, target_shape[2]),
    indexing='ij'
)
voxel_centers = np.vstack([xv.ravel(), yv.ravel(), zv.ravel()]).T

# 5. 利用KDTree快速判断体素是否属于曲面区域
tree = KDTree(surface_points)
# 计算每个体素中心到曲面的最近距离
distances, _ = tree.query(voxel_centers, k=1)

# 设置距离阈值：自适应体素尺寸的一半，确保曲面区域被准确标记
threshold = (x_max - x_min) / (target_shape[0] * 2)
# 将距离小于阈值的体素标记为1，其余为0
voxel_array = (distances <= threshold).reshape(target_shape).astype(np.int8)

# 验证结果
print("生成的体素数组形状:", voxel_array.shape)
print("曲面占据的体素数量:", np.sum(voxel_array))

补充说明 #

• 代码中去掉了np.vectorize，因为numpy原生支持函数广播，直接调用function(x)即可处理数组，效率更高。
• 若需要填充闭合曲面的内部区域，可替换为射线法判断点是否在曲面内部，或使用open3d库的体素化工具，实现更高效的闭合曲面体素化。

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