How to efficiently prepare matrices (2-d array) for multiple arguments?

Question

If you want to evaluate a 1-d array for multiple arguments efficiently i.e. without for-loop, you can do this:

x = array([1, 2, 3])
def gen_1d_arr(x):
    arr = array([2 + x, 2 - x,])
    return arr

gen_1d_arr(x).T

and you get:

array([[ 3,  1],
       [ 4,  0],
       [ 5, -1]])

Okay, but how do you do this for 2-d array like below:

def gen_2d_arr(x):
    arr = array([[2 + x, 2 - x,],
                 [2 * x, 2 / x]])
    return arr

and obtain this?:

array([[[ 3.        ,  1.        ],
        [ 2.        ,  2.        ]],

       [[ 4.        ,  0.        ],
        [ 4.        ,  1.        ]],

       [[ 5.        , -1.        ],
        [ 6.        ,  0.66666667]]])

Also, is this generally possible for n-d arrays?

Answer 1

Look at what you get with your function

In [274]: arr = np.array([[2 + x, 2 - x,],
                 [2 * x, 2 / x]])

In [275]: arr
Out[275]: 
array([[[ 3.        ,  4.        ,  5.        ],
        [ 1.        ,  0.        , -1.        ]],

       [[ 2.        ,  4.        ,  6.        ],
        [ 2.        ,  1.        ,  0.66666667]]])

In [276]: arr.shape
Out[276]: (2, 2, 3)

The 3 comes from x. The middle 2 comes from [2+x, 2-x] pairs, and the 1st 2 from the outer list.

Looks like what you want is a (3,2,2) array. One option is to apply a transpose or axis swap to arr.

arr.transpose([2,0,1])

The basic operation of np.array([arr1,arr2]) is to construct a new array with a new dimension in front, i.e. with shape (2, *arr1(shape)).

There are other operations that combine arrays. np.concatenate and its variants hstack, vstack, dstack, column_stack, join arrays. .reshape() and [None,...], atleast_nd etc add dimensions. Look at the code of the stack functions to get some ideas on how to combine arrays using these tools.

On the question of efficiency, my time tests show that concatenate operations are generally faster than np.array. Often np.array converts its inputs to lists, and reparses the values. This gives it more power in cooercing arrays to specific dtypes, but at the expense of time. But I'd only worry about this with large arrays where construction time is significant.