[FIXED] Extend inner array in 3D numpy array

Issue

I've got two numpy arrays, with shapes something like:

A = (2000, 2000, 3) and B = (2000, 2000, 4)

I also have an array of indices, let's say [1,3]

What I need to do is take the two original arrays, and create a new array that is (2000, 2000, 5) where each inner array is the entire contents of A's inner array, and the specified indices of B's.

So an inner array might be:

(A[0], A[1], A[2], B[1], B[3])

I'm currently doing this directly in Python, like this:

width = 2000
height = 2000
out_data = np.empty((height, width, 5), np.float32)
for y in range(height):
    for x in range(width):
        out_data[y][x][:3] = A[y][x]
        out_data[y][x][3] = B[y][x][1]
        out_data[y][x][4] = B[y][x][3]

This is very slow. What I'm wondering is if there's a better way that I can do this within numpy?

If there's a way to do this by extending A, I am happy with that, as A and B are only used for this purpose.

Solution

ix = [1, 3]
new = np.c_[A, B[:, :, ix]]

Minimal example

width = 5
height = 5

A = np.arange(np.prod((width, height, 3))).reshape(width, height,3)
B = -np.arange(np.prod((width, height, 4))).reshape(width, height,4)

OP's code:

ix = [1, 3]  # or (1, 3)
out_data = np.empty((height, width, 5), A.dtype)
for y in range(height):
    for x in range(width):
        out_data[y][x][:3] = A[y][x]
        out_data[y][x][3] = B[y][x][1]
        out_data[y][x][4] = B[y][x][3]

This solution:

new = np.c_[A, B[:, :, ix]]

>>> np.array_equal(out_data, new)
True

Speed

width = 2000
height = 2000

# ...

# time of OP's: 9.55s

%timeit np.c_[A, B[:, :, ix]]
# 143 ms ± 42.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answered By - Pierre D