Index 3D array with index of last axis stored in 2D array

[zxdawn]

This is copied from a question on stackoverflow.

I have a ndarray of shape(z,y,x) containing values. I am trying to index this array with another ndarray of shape(y,x) that contains the z-index of the value I am interested in.
How do I have to index val_arr with z_indices to get the values at the desired z-axis position?

MCVE Code Sample

import numpy as np
val_arr = np.arange(27).reshape(3, 3, 3)

z_indices = np.array([[1, 0, 2],
                      [0, 0, 1],
                      [2, 0, 1]])

index_array = z_indices.choose(val_arr)
print(index_array)

xarray version

import numpy as np
import xarray as xr

val_arr = xr.DataArray(np.arange(27).reshape(3, 3, 3),
                       dims=['z', 'y', 'x'])

z_indices = xr.DataArray(np.array([[1, 0, 2],
                      [0, 0, 1],
                      [2, 0, 1]]),
                      dims=['y', 'x'])

index_array = np.choose(z_indices, val_arr)
print(index_array)

Expected Output

<xarray.DataArray (y: 3, x: 3)>
array([[ 9,  1, 20],
       [ 3,  4, 14],
       [24,  7, 17]])
Dimensions without coordinates: y, x

Problem Description

Is this feature of choose supported in xarray?

[johnomotani]

I think indexed_array = val_arr.isel(z=z_indices) should work, if I'm understanding what you want. I haven't tested though...

[zxdawn]

@johnomotani Thanks, it works. Sorry for this simple question ...