在numpy中实现三维阵列减法的较好方法

Question

我想知道一个更好的替代方法来执行三维数组的减法。目前，我正在使用np.array([x for x in xs])对每个第一轴执行减法，如下所示：

import numpy as np

# =============================================================================
# Preparing the data
coords = np.array([
    # Frame 1
    [
        [0, 1, 2],  # Particle 1
        [2, 4, 5],  # Particle 2
        [6, 7, 8],  # Particle 3
        [9, 9, 9],  # Particle 4
    ],
    # Frame 2
    [
        [-0, -1, -2],  # Particle 1
        [-2, -4, -5],  # Particle 2
        [-6, -7, -8],  # Particle 3
        [-9, -9, -9],  # Particle 4
    ],
])
# shape := (n_frames, n_particles, n_dimensions)
assert coords.shape == (2, 4, 3)

centers = np.array([
  [1, 3, 6],  # Frame 1
  [4, 8, 16],  # Frame 2
])
# shape := (n_frames, n_dimensions)
assert centers.shape == (2, 3)
# =============================================================================


# =============================================================================
# Subtract each centers from each particles

# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# I would like to know a better alternative way to do the below
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
n_frames = coords.shape[0]
coords2 = np.asarray(
    [coords[i] - centers[i] for i in range(n_frames)],
)
# shape := (n_frames, n_particles, n_dimensions)
assert coords2.shape == (2, 4, 3)
# =============================================================================


# =============================================================================
# Test if result coordinates are correct
np.testing.assert_array_equal(coords2[0], [
    [0-1, 1-3, 2-6],  # Particle 1
    [2-1, 4-3, 5-6],  # Particle 2
    [6-1, 7-3, 8-6],  # Particle 3
    [9-1, 9-3, 9-6],  # Particle 4
])
np.testing.assert_array_equal(coords2[1], [
    [-0-4, -1-8, -2-16],  # Particle 1
    [-2-4, -4-8, -5-16],  # Particle 2
    [-6-4, -7-8, -8-16],  # Particle 3
    [-9-4, -9-8, -9-16],  # Particle 4
])
# =============================================================================

但是我觉得如果coords和centers有大量的帧(真实数据的coords将是shape := (>10000, >1000, 3))，那么我的实现就会有一些性能问题。

我知道，通过重写上面的代码，Cython将大大提高性能。因此，我可能使用Cython，但我想知道是否有更好(更快)的方法在上执行这种减法，只使用numpy。

谢谢你的阅读；)

Answer 1

coords2 = coords - centers[:, np.newaxis]

在中间有一个额外的长度-1轴的centers视图，使形状对齐，以供广播。