CODEWITHSUNDEEP
pythonnumpy
MrBudgens
MrBudgens

Reputation: 33

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

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.

Upvotes: 3

Views: 221

Answers (3)

Epsi95
Epsi95

Reputation: 9047

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))

Upvotes: 1

JaHoff
JaHoff

Reputation: 206

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.

Upvotes: 2

Mad Physicist
Mad Physicist

Reputation: 114230

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.

Upvotes: 4

Related Questions