CODEWITHSUNDEEP
pythonarraysnumpy
lambdalisue
lambdalisue

Reputation: 125

The better way to perform subtraction of 3-dimensional array in numpy

I would like to know a better alternative way to perform subtraction on three dimensional array. Currently I'm using np.array([x for x in xs]) to perform subtraction on each first axis like: 

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
])
# =============================================================================

But I felt that my implementation would have some performance problem if the coords and centers have tons of frames (the coords of real data would be shape := (>10000, >1000, 3)).

I know that Cython would improve the performance a lot by re-writing the code above. So I probably use Cython but I wondered if there is any better (faster) way to perform such subtraction on by using only numpy.

Thanks for reading ;-)

Upvotes: 0

Views: 1578

Answers (1)

user2357112
user2357112

Reputation: 280788

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

Take a view of centers with an extra length-1 axis in the middle to make the shapes line up right for broadcasting.

Upvotes: 3

Related Questions