How to return an array of at least 4D: efficient method to simulate numpy.atleast_4d

Question

numpy provides three handy routines to turn an array into at least a 1D, 2D, or 3D array, e.g. through numpy.atleast_3d

I need the equivalent for one more dimension: atleast_4d. I can think of various ways using nested if statements but I was wondering whether there is a more efficient and faster method of returning the array in question. In you answer, I would be interested to see an estimate (O(n)) of the speed of execution if you can.

Answer 1

The np.array method has an optional ndmin keyword argument that:

Specifies the minimum number of dimensions that the resulting array should have. Ones will be pre-pended to the shape as needed to meet this requirement.

If you also set copy=False you should get close to what you are after.

---

As a do-it-yourself alternative, if you want extra dimensions trailing rather than leading:

arr.shape += (1,) * (4 - arr.ndim)

Answer 2

Why couldn't it just be something as simple as this:

import numpy as np
def atleast_4d(x):
    if x.ndim < 4:
        y = np.expand_dims(np.atleast_3d(x), axis=3)
    else:
        y = x

    return y

ie. if the number of dimensions is less than four, call atleast_3d and append an extra dimension on the end, otherwise just return the array unchanged.