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视图，使形状排列正确，以便进行广播。

在numpy中执行三维数组减法的更好方法

1 个答案: