Alex
Alex

Reputation: 306

Meshgrid a N-columned matrix in Numpy (or smth else)

Python version: 2.7

I have the following numpy 2d array:

array([[ -5.05000000e+01,  -1.05000000e+01],
   [ -4.04000000e+01,  -8.40000000e+00],
   [ -3.03000000e+01,  -6.30000000e+00],
   [ -2.02000000e+01,  -4.20000000e+00],
   [ -1.01000000e+01,  -2.10000000e+00],
   [  7.10542736e-15,  -1.77635684e-15],
   [  1.01000000e+01,   2.10000000e+00],
   [  2.02000000e+01,   4.20000000e+00],
   [  3.03000000e+01,   6.30000000e+00],
   [  4.04000000e+01,   8.40000000e+00]])

If I wanted to find all the combinations of the first and the second columns, I would use np.array(np.meshgrid(first_column, second_column)).T.reshape(-1,2). As a result, I would get a 100*1 matrix with 10*10 = 100 data points. However, my matrix can have 3, 4, or more columns, so I have a problem of using this numpy function.

Question: how can I make an automatically meshgridded matrix with 3+ columns?

UPD: for example, I have the initial array:

[[-50.5 -10.5]
 [  0.    0. ]]

As a result, I want to have the output array like this:

array([[-10.5, -50.5],
   [-10.5,   0. ],
   [  0. , -50.5],
   [  0. ,   0. ]])

or this:

array([[-50.5, -10.5],
   [-50.5,   0. ],
   [  0. , -10.5],
   [  0. ,   0. ]])

Upvotes: 1

Views: 1044

Answers (1)

Divakar
Divakar

Reputation: 221564

You could use * operator on the transposed array version that unpacks those columns sequentially. Finally, a swap axes operation is needed to merge the output grid arrays as one array.

Thus, one generic solution would be -

np.swapaxes(np.meshgrid(*arr.T),0,2)

Sample run -

In [44]: arr
Out[44]: 
array([[-50.5, -10.5],
       [  0. ,   0. ]])

In [45]: np.swapaxes(np.meshgrid(*arr.T),0,2)
Out[45]: 
array([[[-50.5, -10.5],
        [-50.5,   0. ]],

       [[  0. , -10.5],
        [  0. ,   0. ]]])

Upvotes: 1

Related Questions