Numpy: assemble three 1D arrays into 3D (but not exactly a simple coordinate grid)

user_14945334 0
Viewed 119

I have three 1D arrays, where two are the same length but the third is a different length, e.g.

A = np.array([1, 2, 3, 4])
B = np.array([10, 20, 30, 40])
C = np.array([100, 200, 300, 400, 500])

I want to combine them into a grid as follows:

D = np.array([[[1, 10, 100], [1, 10, 200], [1, 10, 300], [1, 10, 400], [1, 10, 500]],
              [[2, 20, 100], [2, 20, 200], [2, 20, 300], [2, 20, 400], [2, 20, 500]],
              [[3, 30, 100], [3, 30, 200], [3, 30, 300], [3, 30, 400], [3, 30, 500]],
              [[4, 40, 100], [4, 40, 200], [4, 40, 300], [4, 40, 400], [4, 40, 500]]])

It's similar to a coordinate grid, but where A and B change together, so I assume they need to be combined first, and the resulting 2D array then combined with C somehow, but I can't find the right combination of stack, concatenate, meshgrid, etc.

The real arrays will be much larger than this example (thousands of values each), and the code will be called very many times, so speed is important. It will live in the main loop of an MCMC sampler, where A and B are parameters that vary on each loop, while C is a constant.

numpy python
3 Answers

Let's call the lengths of A, B, C a, b, c, respectively. You are looking for an output array of shape (a, c, 3) (or (b, c, 3)). Stacking A and B correctly gives you an array of shape (a, 2):

zipped = np.stack((A, B), axis=-1)

To be able to append C to zipped, you will need to broadcast the arrays to the same shape along the first two axes of the result:

Aa = np.broadcast_to(zipped[..., None, :], (A.size, C.size, 2))
Cc = np.broadcast_to(C[..., None], (A.size, C.size, 1))
result = np.append(Aa, Cc, axis=-1)

The whole operation can be written as a one-liner:

result = np.append(np.broadcast_to(np.stack((A, B), axis=-1)[..., None, :], (A.size, C.size, 2)), np.broadcast_to(C[..., None], (A.size, C.size, 1)), axis=-1)

This operation is fairly efficient: np.broadcast_to creates a view without copying any data. Only the append operation itself makes a copy.

user_29887300

You can exploit the multiplication of differently oriented matrixes followed by concatenation to achieve your desired result:

A = np.array([1,2,3,4]).reshape((4,1,1))
B = np.array([10,20,30,40]).reshape((4,1,1))
C = np.array([100,200,300,400,500]).reshape((1,5,1))

res = np.concatenate([A*np.ones((1,5,1)),B*np.ones((1,5,1)),C*np.ones((4,1,1))],axis=2)
print(res)

Yields: res

Out[152]: 
array([[[  1.,  10., 100.],
        [  1.,  10., 200.],
        [  1.,  10., 300.],
        [  1.,  10., 400.],
        [  1.,  10., 500.]],

       [[  2.,  20., 100.],
        [  2.,  20., 200.],
        [  2.,  20., 300.],
        [  2.,  20., 400.],
        [  2.,  20., 500.]],

       [[  3.,  30., 100.],
        [  3.,  30., 200.],
        [  3.,  30., 300.],
        [  3.,  30., 400.],
        [  3.,  30., 500.]],

       [[  4.,  40., 100.],
        [  4.,  40., 200.],
        [  4.,  40., 300.],
        [  4.,  40., 400.],
        [  4.,  40., 500.]]])

This type of operation is pretty fast, but you might run into performance problems if your matrixes are extremely big. This operation does like to use a lot of memory.

user_138538200

column_stack

A = np.array([1,2,3,4])
B = np.array([10,20,30,40])
C = np.array([100,200,300,400, 500])


op1 = np.column_stack((A,B))
tt = []
for i in C:
   tt.append(np.column_stack((op1, np.array([i]*A.shape[0]))))
   

print(np.concatenate(tt).reshape(C.shape[0],A.shape[0],3).transpose(1,0,2))
user_66606380
Related