CODEWITHSUNDEEP
pythonarraysnumpy
Erik
Erik

Reputation: 7342

numpy add variable number of dimensions to an array

Is there a way to add a variable number of singleton dimensions to a numpy array? I want something like atleast_2d but for an arbitrary number of dimensions. expand_dims comes up a lot, but only adds a single dimension.

The only way I know to do this is to compute the shape first and apply it, i.e.

import numpy as np

def atleast_kd(array, k):
    array = np.asarray(array)
    new_shape = array.shape + (1,) * (k - array.ndim)
    return array.reshape(new_shape)

Upvotes: 3

Views: 1097

Answers (1)

zephyrus
zephyrus

Reputation: 1266

A more elegant way to do this is np.broadcast_to. For example:

a = np.random.rand(2,2)
k = # number_extra dimensions
b = np.broadcast_to(a, (1,) * k + a.shape)

This will add the dimensions to the beginning of b. To get the exact same behavior as the given function, you can use np.moveaxis.

Upvotes: 1

Related Questions