Python高效计算(x,y)相空间轨道符号变化的优化方案

问题背景 #

我通过时间积分计算得到相空间轨道数据，轨道上每个位置对应坐标(x, y)，轨道示例如下图：

现在需要基于x、y坐标数组构建跨零事件序列，序列定义如下图：

序列取值规则：

• 值为0：x坐标在相邻时刻t_n与t_{n+1}之间发生符号变化（跨零）
• 值为1：y坐标在相邻时刻t_n与t_{n+1}之间发生符号变化（跨零）
  上述示例轨道对应的序列形式为：Pxy[n] = 1 0 1 0 1 0 ...

当前计算存在明显性能瓶颈：单组轨道的x、y坐标数组包含1000万个元素——为保证积分精度，积分算法的时间步长必须设置得非常小，因此需要该规模的数组。目前实现的最快算法单次运行耗时约5.3秒，而该算法需要针对100余组不同轨道重复运行，总耗时难以接受。

现有实现方案 #

目前共有两种实现方式：

方案1：基于NumPy的实现 #

参考NumPy符号变化检测的通用方案编写，代码如下：

def calc_Pxy():
    Pxy = np.full(x.shape, -1)
    Pxy[np.where(np.diff(np.signbit(x)))] = 0
    Pxy[np.where(np.diff(np.signbit(y)))] = 1
    return Pxy[Pxy >= 0]

方案2：纯Python for循环实现 #

代码如下：

def calc_Pxy():
    Pxy = []
    for n, [xn, yn] in enumerate(zip(x[:-1], y[:-1])):
        if xn * x[n + 1] < 0:
            Pxy.append(0)
        if yn * y[n + 1] < 0:
            Pxy.append(1)
    return np.array(Pxy)

现存问题 #

1. 小规模数据集性能反差：最初在仅含20个值的小规模数据集上测试时，NumPy实现比普通for循环慢2.5倍，想确认该现象的原因是否是数组重塑、创建新数组带来的固定时间开销？
2. NumPy实现逻辑不符合需求：当同一时间步x、y同时发生符号变化时，现有NumPy实现仅会记录1次跨零事件；而for循环实现会按先记录x跨零、再记录y跨零的顺序存储两次结果，该逻辑不会丢失数据，是我需要的行为。
3. 大规模场景下的优化需求：后续在大规模数据集上测试发现，数据规模增大到100万元素以上时，NumPy方法的速度会反超for循环，但我仍需要调整NumPy实现使其行为与for循环完全一致，同时欢迎进一步提升运算速度的优化思路。

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