Indexing array with array on numpy - python

It is similar to some questions around SO, but I don't quite understand the trick to get what I want.
I have two arrays,
arr of shape (x, y, z)
indexes of shape (x, y) which hold indexes of interest for z.
For each value of indexes I want to get the actual value in arr where:
arr.x == indexes.x
arr.y == indexes.y
arr.z == indexes[x,y]
This would give an array of shape(x,y) similar to indexes' shape.
For example:
arr = np.arange(99)
arr = arr.reshape(3,3,11)
indexes = np.asarray([
[0,2,2],
[1,2,3],
[3,2,10]])
# indexes.shape == (3,3)
# Example for the first element to be computed
first_element = arr[0,0,indexes[0,0]]
With the above indexes, the expected arrays would look like:
expected_result = np.asarray([
[0,13,24],
[34,46,58],
[69,79,98]])
I tried elements = np.take(arr, indexes, axis=z)
but it gives an array of shape (x, y, x, y)
I also tried things like elements = arr[indexes, indexes,:] but I don't get what I wish.
I saw a few answers involving transposing indexes and transforming it into tuples but I don't understand how it would help.
Note: I'm a bit new to numpy so I don't fully understand indexing yet.
How would you solve this numpy style ?

This can be done using np.take_along_axis
import numpy as np
#sample data
np.random.seed(0)
arr = np.arange(3*4*2).reshape(3, 4, 2) # 3d array
idx = np.random.randint(0, 2, (3, 4)) # array of indices
out = np.squeeze(np.take_along_axis(arr, idx[..., np.newaxis], axis=-1))
In this code, the array of indices gets added one more axis, so it can be broadcasted to the shape of the array arr from which we are making the selection. Then, since the return value of np.take_along_axis has the same shape as the array of indices, we need to remove this extra dimension using np.squeeze.
Another option is to use np.choose, but in this case the axis along which you are making selections must be moved to be the first axis of the array:
out = np.choose(idx, np.moveaxis(arr, -1, 0))

The solution here should work for you: Indexing 3d numpy array with 2d array
Adapted to your code:
ax_0 = np.arange(arr.shape[0])[:,None]
ax_1 = np.arange(arr.shape[1])[None,:]
new_array = arr[ax_0, ax_1, indexes]

You can perform such an operation with np.take_along_axis, the operation can only be applied along one dimension so you will need to reshape your input and indices.
The operation you are looking to perform is:
out[i, j] = arr[i, j, indices[i, j]]
However, we are forced to reshape both arr and indices, i.e. map (i, j) to k, such that we can apply np.take_along_axis. The following operation will take place:
out[k] = arr[k, indices[k]] # indexing along axis=1
The actual usage here comes down to:
>>> put = np.take_along_axis(arr.reshape(9, 11), indices.reshape(9, 1), axis=1)
array([[ 0],
[13],
[24],
[34],
[46],
[58],
[69],
[79],
[91]])
Then reshape back to the shape of indices:
>>> put.reshape(indices.shape)
array([[ 0, 13, 24],
[34, 46, 58],
[69, 79, 91]])

Related

Alternative to loop for for boolean / nonzero indexing of numpy array

I need to select only the non-zero 3d portions of a 3d binary array (or alternatively the true values of a boolean array). Currently I am able to do so with a series of 'for' loops that use np.any, but this does work but seems awkward and slow, so currently investigating a more direct way to accomplish the task.
I am rather new to numpy, so the approaches that I have tried include a) using
np.nonzero, which returns indices that I am at a loss to understand what to do with for my purposes, b) boolean array indexing, and c) boolean masks. I can generally understand each of those approaches for simple 2d arrays, but am struggling to understand the differences between the approaches, and cannot get them to return the right values for a 3d array.
Here is my current function that returns a 3D array with nonzero values:
def real_size(arr3):
true_0 = []
true_1 = []
true_2 = []
print(f'The input array shape is: {arr3.shape}')
for zero_ in range (0, arr3.shape[0]):
if arr3[zero_].any()==True:
true_0.append(zero_)
for one_ in range (0, arr3.shape[1]):
if arr3[:,one_,:].any()==True:
true_1.append(one_)
for two_ in range (0, arr3.shape[2]):
if arr3[:,:,two_].any()==True:
true_2.append(two_)
arr4 = arr3[min(true_0):max(true_0) + 1, min(true_1):max(true_1) + 1, min(true_2):max(true_2) + 1]
print(f'The nonzero area is: {arr4.shape}')
return arr4
# Then use it on a small test array:
test_array = np.zeros([2, 3, 4], dtype = int)
test_array[0:2, 0:2, 0:2] = 1
#The function call works and prints out as expected:
non_zero = real_size(test_array)
>> The input array shape is: (2, 3, 4)
>> The nonzero area is: (2, 2, 2)
# So, the array is correct, but likely not the best way to get there:
non_zero
>> array([[[1, 1],
[1, 1]],
[[1, 1],
[1, 1]]])
The code works appropriately, but I am using this on much larger and more complex arrays, and don't think this is an appropriate approach. Any thoughts on a more direct method to make this work would be greatly appreciated. I am also concerned about errors and the results if the input array has for example two separate non-zero 3d areas within the original array.
To clarify the problem, I need to return one or more 3D portions as one or more 3d arrays beginning with an original larger array. The returned arrays should not include extraneous zeros (or false values) in any given exterior plane in three dimensional space. Just getting the indices of the nonzero values (or vice versa) doesn't by itself solve the problem.
Assuming you want to eliminate all rows, columns, etc. that contain only zeros, you could do the following:
nz = (test_array != 0)
non_zero = test_array[nz.any(axis=(1, 2))][:, nz.any(axis=(0, 2))][:, :, nz.any(axis=(0, 1))]
An alternative solution using np.nonzero:
i = [np.unique(_) for _ in np.nonzero(test_array)]
non_zero = test_array[i[0]][:, i[1]][:, :, i[2]]
This can also be generalized to arbitrary dimensions, but requires a bit more work (only showing the first approach here):
def real_size(arr):
nz = (arr != 0)
result = arr
axes = np.arange(arr.ndim)
for axis in range(arr.ndim):
zeros = nz.any(axis=tuple(np.delete(axes, axis)))
result = result[(slice(None),)*axis + (zeros,)]
return result
non_zero = real_size(test_array)

Numpy view on 1D array via 2D array as indices range

I have a 2D array which describes index ranges for a 1D array like
z = np.array([[0,4],[4,9]])
The 1D array
a = np.array([1,1,1,1,0,0,0,0,0,1,1,1,1])
I want to have a view on the 1D array with the index range defined by z. So, for only the first range
a[z[0][0]:z[0][1]]
How to get it for all ranges? Is it possible to use as_strided with unequal lengths defined by z as shape? I want to avoid to copy data, actually I only want a different view on a for further computation.
In [66]: a = np.array([1,1,1,1,0,0,0,0,0,1,1,1,1])
In [67]: z = np.array([[0,4],[4,9]])
So generating the slices from the rows of z we get 2 arrays:
In [68]: [a[x[0]:x[1]] for x in z]
Out[68]: [array([1, 1, 1, 1]), array([0, 0, 0, 0, 0])]
Individually those arrays are views. But together they aren't an array. The lengths diff, so they can't be vstacked into a (2,?) array. They can be hstacked but that won't be a view.
The calculation core of np.array_split is:
sub_arys = []
sary = _nx.swapaxes(ary, axis, 0)
for i in range(Nsections):
st = div_points[i]
end = div_points[i + 1]
sub_arys.append(_nx.swapaxes(sary[st:end], axis, 0))
Ignoring the swapaxes bit, this is doing the same thing as my list comprehension.
for x, y in z:
array_view = a[x:y]
# do something with array_view

indexing a numpy array with indices from another array

I have a numpy array "data" with dimensions [t, z, x, y]. The
dimensions represent time (t) and three spatial dimensions (x, y, z).
I have a separate array "idx" of indices with dimensions [t, x, y]
describing vertical coordinates in data: each value in idx describes a
single vertical level in data.
I want to pull out the values from data indexed by idx. I've done it
successfully using loops (below). I've read several SO threads and numpy's indexing docs but I haven't been able to make it more pythonic/vectorized.
Is there an easy way I'm just not getting quite right? Or maybe loops
are a clearer way to do this anyway...
import numpy as np
dim = (4, 4, 4, 4) # dimensions time, Z, X, Y
data = np.random.randint(0, 10, dim)
idx = np.random.randint(0, 3, dim[0:3])
# extract vertical indices in idx from data using loops
foo = np.zeros(dim[0:3])
for this_t in range(dim[0]):
for this_x in range(dim[2]):
for this_y in range(dim[3]):
foo[this_t, this_x, this_y] = data[this_t,
idx[this_t, this_x, this_y],
this_x,
this_y]
# surely there's a better way to do this with fancy indexing
# data[idx] gives me an array with dimensions (4, 4, 4, 4, 4, 4)
# data[idx[:, np.newaxis, ...]] is a little closer
# data[tuple(idx[:, np.newaxis, ...])] doesn't quite get it either
# I tried lots of variations on those ideas but no luck yet
In [7]: I,J,K = np.ogrid[:4,:4,:4]
In [8]: data[I,idx,J,K].shape
Out[8]: (4, 4, 4)
In [9]: np.allclose(foo, data[I,idx,J,K])
Out[9]: True
I,J,K broadcast together to the same shape as idx (4,4,4).
More detail on this kind of indexing at
How to take elements along a given axis, given by their indices?

How to assign range(10) to a ndarray which shape is (10, 1) conveniently in Numpy?

I tried this:
import numpy as np
a = np.empty((1, 10, 1), np.int8)
a[0] = range(10)
It throw error: ValueError: could not broadcast input array from shape (10) into shape (10,1)
Several options that work in this case:
a[0, :, 0] = np.arange(10) # assign to 1D slice
a[0].flat = range(10) # assign to flattened 2D slice
a[0] = np.arange(10).reshape(10, 1) # bring the rigth side into correct shape
a[0] = np.arange(10)[:, np.newaxis] # bring the rigth side into correct shape
Note the use of np.arange instead of range. The former directly creates an ndarray with a sequence of values, while the latter creates an iterable that needs to be converted into an array for the assignment.
In the case of assigning to flat it makes sense to use range because both are iterators.
You can do a[0, :, 0] = range(10).

How to generate multi-dimensional 2D numpy index using a sub-index for one dimension

I want to use numpy.ix_ to generate an multi-dimensional index for a 2D space of values. However, I need to use a subindex to look up the indices for one dimension. For example,
assert subindex.shape == (ny, nx)
data = np.random.random(size=(ny,nx))
# Generator returning the index tuples
def get_idx(ny,nx,subindex):
for y in range(ny):
for x in range(nx):
yi = y # This is easy
xi = subindex[y,x] # Get the second index value from the subindex
yield (yi,xi)
# Generator returning the data values
def get_data_vals(ny,nx,data,subindex):
for y in range(ny):
for x in range(nx):
yi = y # This is easy
xi = subindex[y,x] # Get the second index value from the subindex
yield data[y,subindex[y,x]]
So instead of the for loops above, I'd like to use a multi-dimensional index to index data Using numpy.ix_, I guess I would have something like:
idx = numpy.ix_([np.arange(ny), ?])
data[idx]
but I don't know what the second dimension argument should be. I'm guessing it should be something involving numpy.choose?
What you actually seem to want is:
y_idx = np.arange(ny)[:,np.newaxis]
data[y_idx, subindex]
BTW, you could achieve the same thing with y_idx = np.arange(ny).reshape((-1, 1)).
Let's look at a small example:
import numpy as np
ny, nx = 3, 5
data = np.random.rand(ny, nx)
subindex = np.random.randint(nx, size=(ny, nx))
Now
np.arange(ny)
# array([0, 1, 2])
are just the indices for the "y-axis", the first dimension of data. And
y_idx = np.arange(ny)[:,np.newaxis]
# array([[0],
# [1],
# [2]])
adds a new axis to this array (after the existing axis) and effectively transposes it. When you now use this array in an indexing expression together with the subindex array, the former gets broadcasted to the shape of the latter. So y_idx becomes effectively:
# array([[0, 0, 0, 0, 0],
# [1, 1, 1, 1, 1],
# [2, 2, 2, 2, 2]])
And now for each pair of y_idx and subindex you look up an element in the data array.
Here you can find out more about "fancy indexing"
It sounds like you need to do two things:
Find all indices into the data array and
Translate the column indices according to some other array, subindex.
The code below therefore generates indices for all array positions (using np.indices), and reshapes it to (..., 2) -- a 2-D list of coordinates representing each position in the array. For each coordinate, (i, j), we then translate the column coordinate j using the subindex array provided, and then use that translated index as the new column index.
With numpy, it is not necessary to do that in a for-loop--we can simply pass in all the indices at once:
i, j = np.indices(data.shape).reshape((-1, 2)).T
data[i, subindex[i, j]]

Categories

Resources