Извлечение данных Xarray 4D с использованием 2D массива индексов

Question

У меня есть набор данных xarray с 4 измерениями (time,depth,latitude,longitude). Пример набора данных приведен ниже. Я хочу извлечь данные для каждого временного шага на самом глубоком уровне глубины, где есть действительные данные (дно океана или ближайшая действительная точка сетки). Я нахожу самую глубокую допустимую точку с помощью следующей функции:

    def find_deepest_depth_indexes(ds):
        # First get the vertical True/False of valid values
        idx = ds["thetao"].isel(time=0).isnull()
        idx_vals=idx.values
        # Create the initial final array to store indices (integer type)
        depth_indexes = np.zeros((len(idx.latitude), len(idx.longitude))).astype(int)

        # Now find the deepest depth where values are True and store in indices array
        for i in range(len(ds.longitude.values)):
            for j in range(len(ds.latitude.values)):
                located = np.where(idx_vals[:, j, i] == True)
                depth_indexes[j, i] = int(located[0][0])

        return depth_indexes

Это работает хорошо и быстро, так как мне нужно только один раз найти самые глубокие индексы глубины. Затем я хочу извлечь точки данных в моем 4D-массиве для всех временных шагов с заданными индексами depth. Я использую следующий код:

    def calculate_climatology(varname, ds, depth_indexes):

        # Extract 2D data for each timestep using 2D depth indices on 4D array
        final_array=np.empty((len(ds.time),len(ds.latitude),len(ds.longitude)))
        m,n = depth_indexes.shape
        I,J = np.ogrid[:m,:n]

        for t in range(len(ds.time.values)):
            data_at_timestep=ds[varname][t,:,:,:].squeeze()
            datavalues_at_timestep=data_at_timestep.values

            final_array[t,:,:] = datavalues_at_timestep[depth_indexes, I, J]
        return final_array

Это не работает (работает, но производит только nans в final_array). Я попытался воспроизвести то, что было предложено в качестве решения в предыдущем вопросе SO , но не смог заставить это работать должным образом. Если кто-то видит, что здесь не так, или у него есть альтернативные подходы, я благодарен.

    <xarray.Dataset>
Dimensions:    (depth: 50, latitude: 841, longitude: 961, time: 312)
Coordinates:
  * longitude  (longitude) float32 -180.0 -179.91667 ... -100.083336 -100.0
  * latitude   (latitude) float32 10.0 10.083333 10.166667 ... 79.916664 80.0
  * depth      (depth) float32 0.494025 1.541375 2.645669 ... 5274.784 5727.917
  * time       (time) datetime64[ns] 1993-01-16T12:00:00 ... 2018-12-16T12:00:00
Data variables:
    mlotst     (time, latitude, longitude) float32 dask.array<chunksize=(1, 170, 250), meta=np.ndarray>
    zos        (time, latitude, longitude) float32 dask.array<chunksize=(1, 170, 250), meta=np.ndarray>
    bottomT    (time, latitude, longitude) float32 dask.array<chunksize=(1, 170, 250), meta=np.ndarray>
    sithick    (time, latitude, longitude) float32 dask.array<chunksize=(1, 170, 250), meta=np.ndarray>
    siconc     (time, latitude, longitude) float32 dask.array<chunksize=(1, 170, 250), meta=np.ndarray>
    usi        (time, latitude, longitude) float32 dask.array<chunksize=(1, 170, 250), meta=np.ndarray>
    vsi        (time, latitude, longitude) float32 dask.array<chunksize=(1, 170, 250), meta=np.ndarray>
    thetao     (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 50, 170, 250), meta=np.ndarray>
    so         (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 50, 170, 250), meta=np.ndarray>
    uo         (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 50, 170, 250), meta=np.ndarray>
    vo         (time, depth, latitude, longitude) float32 dask.array<chunksize=(1, 50, 170, 250), meta=np.ndarray>
Attributes:
    title:                         Monthly mean fields for product GLOBAL_REA...
    references:                    http://marine.copernicus.eu
    credit:                        E.U. Copernicus Marine Service Information...
    licence:                       http://marine.copernicus.eu/services-portf...
    contact:                       servicedesk.cmems@mercator-ocean.eu
    producer:                      CMEMS - Global Monitoring and Forecasting ...
    institution:                   Mercator Ocean
    Conventions:                   CF-1.6
    area:                          GLOBAL
    product:                       GLOBAL_REANALYSIS_001_030
    dataset:                       global-reanalysis-001-030-monthly
    source:                        MERCATOR GLORYS12V1
    product_user_manual:           http://marine.copernicus.eu/documents/PUM/...
    quality_information_document:  http://marine.copernicus.eu/documents/QUID...e here