Extend inner array in 3D numpy array-CodePudding

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.

CodePudding user response：

You can dstack you arrays along the last axis:

out_data = np.dstack([A, B[:,:,[1,3]]])

CodePudding user response：

You should be able to use np.concatenate like this for your example.

np.concatenate((a, b[:,:,(1,3)]), axis=2)

On 2000x2000 arrays as input, this ran about 61 times as fast as your original code and roughly equal to the direct index.

CodePudding user response：

The following does the same as your code in vectorized manner, without the for loops:

out = np.empty((2000, 2000, 5))
out[:, :, :3] = A
out[:, :, 3:5] = B[:, :, [1,3]]

CodePudding user response：

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)