用python写一个求数组主曲率的脚本

主曲率是曲面上某点处曲率半径的倒数，可以使用以下代码来计算：

import numpy as np

def principal_curvatures(surf):
    """
    计算曲面上每个点的主曲率
    :param surf: 曲面点集，形如 (N, 3) 的矩阵
    :return: 包含每个点的主曲率和法向量的元组
    """
    # 计算曲面法向量
    n = np.zeros(surf.shape)
    q = np.zeros((surf.shape[0], 3, 3))
    for i, x in enumerate(surf):
        ix, iy = np.gradient(surf[:, :2])
        nx = np.cross(ix, iy)
        nx /= np.linalg.norm(nx, axis=1).reshape(-1, 1)
        n[i] = nx[i]
        q[i] = np.outer(nx[i], nx[i])
    # 计算二次型矩阵
    m = np.zeros((surf.shape[0], 2, 2))
    m[:, 0, 0] = np.einsum('ijk,ik->i', q, np.array([1, 0, 0]))
    m[:, 0, 1] = np.einsum('ijk,ik->i', q, np.array([1, 1, 0])) / 2.
    m[:, 1, 0] = m[:, 0, 1]
    m[:, 1, 1] = np.einsum('ijk,ik->i', q, np.array([0, 1, 0]))
    # 计算主曲率
    k1 = (m[:, 0, 0] + m[:, 1, 1] + np.sqrt((m[:, 0, 0] - m[:, 1, 1]) ** 2 + 4 * m[:, 0, 1] ** 2)) / 2.
    k2 = (m[:, 0, 0] + m[:, 1, 1] - np.sqrt((m[:, 0, 0] - m[:, 1, 1]) ** 2 + 4 * m[:, 0, 1] ** 2)) / 2.
    return k1, k2, n

这个函数需要一个形如 (N, 3) 的曲面点集作为输入，返回每个点的主曲率和对应的法向量。主曲率被存储在 `k1` 和 `k2` 中，法向量被存储在 `n` 中。