Redimension numpy array dynamically

Question

I have an array, lets say:

A = np.array([1., 2.])

I can add a new dimension by simply doing something like:

def redim(arr):
   return arr[:, None]

Now, I want to add extra dimensions dynamically in this manner:

def redim(arr, dims):
   return arr[dims]

So I can do things like:

redim(A, [:, None])
redim(A, [None, :])
redim(A, [none, :, None])

Does anybody know if this is even possible?

What kind of second argument should it be so I can modify it dynamically? (string, obj, list)? How do I connect both arguments to get the expanded array?

I know I can go the expand_dims route, but if there is a way to define all the expanded dimensions in one swipe, my code will be much clearer.

Thanks in advance!

Answer 1

I think reshape is what you are looking for.

If string theory taught us one thing, it is that everything requires 11 dimensions:

>>> A = np.array([1,2])
>>> A = A.reshape((2,1,1,1,1,1,1,1,1,1,1))
>>> A
array([[[[[[[[[[[1]]]]]]]]]],






   [[[[[[[[[[2]]]]]]]]]]])

And now you can get slices along any dimension:

>>> A[1,:,:,:,:,0]
array([[[[[[[[[2]]]]]]]]])

The only requirement when using reshape is that your total number of elements is unchanged (i.e. the product of dimensions must be unchanged between the first and second shape)

Edit

If you want to use slices anyway, you have to pass a tuple of arguments, like this:

def redim(a, dims):
    return a[dims]

>>> redim(A, (slice(None), None, slice(None), None))

Numpy arrays have a special behaviour when passed tuple of arguments (which is exactly what multi-indexing like A[1,2,:] does under the cover), and if you want to wrap calls to the [] operator, you actually have to comply with the np.ndarray.__getitem__ function signature.

To get multidimensional indices, it expects a tuple of ints, None or slice objects (which are the objects created when using slice notations, e.g. A[1:2] <=> A.__getitem__(slice(1,2))). It is quite cumbersome to work with, but you can wrap all numpy operations if you stick to the numpy conventions for arrays (summarized here)