Basically I want to simply a function like this:
arr = np.array([[0.2, 0.4], [-0.9, -0.8]])
def take_dimension(arr, dim=None):
"""dim is either None, 0 or 1"""
# return array with specific dimension
if dim is not None:
return arr[dim, :]
# return array unaltered
return arr
print(take_dimension(arr)) # returns full array:
# [[ 0.2 0.4]
# [-0.9 -0.8]]
print(take_dimension(arr, 0)) # returns: [0.2 0.4]
print(take_dimension(arr, 1)) # returns: [-0.9 -0.8]
Tried
Counter to my expectations, return arr[None, :]
does not return arr
as-is, but instead adds a dimension (arr[None, :].shape
outputs (1, 2, 2)
).
Question
Is there a value I can assign to dim
, so that it returns the full array (effectively functioning as :
)? Then I could simplify the code to:
# not working
def take_dimension(arr, dim=':'):
return arr[dim, :]
print(take_dimension(arr))
# IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices