Select xarray 中海底的值

Question

我有一个带坐标（时间、深度、纬度、经度）的沿海海洋数据集。深层数据被靠近海岸的 nan 掩盖了。我想通过 select 沿着海底的值创建一个坐标（时间、纬度、经度）的新数据集。

目前我正在使用 dataset.bfill('depth').isel({'depth': 0}) 执行此操作，它使用存在的最深数据回填所有 nan，然后切掉最深层。这有效，但效率低下。 bfill 将更新所有时间步长的每个变量，以向下填充缺失值。

海底深度不会在时间步长或变量之间发生变化。我想利用这个事实使这个操作更有效率。假设我有一个 (lat, lon) 数组，其中包含深度坐标的索引，指示海底在该坐标处的位置。制作这个索引数组相对容易，但我不知道如何使用它来 select 正确的数据。

有没有一种方法可以跨所有变量和时间步长有效地使用这个深度索引数组 select 我感兴趣的深度索引？即类似：

>>> dataset
Dimensions:  (t: 5, z: 5, y: 5, x: 5)
Coordinates:
    time     (t) datetime64 2022-02-08 ...
    lon      (x) int64 0 -1 -2 -3 -4
    lat      (y) int64 0 1 2 3 4
    depth    (z) float64 4.25 3.25 2.25 1.25 0.25
Dimensions without coordinates: z, y, x, t
Data variables:
    temp     (t, z, y, x) float64 0.0 nan nan nan nan nan ... 4.0 4.0 4.0 4.0 4.0

>>> depth_indices = compute_ocean_floor_index(
    dataset, depth_variable='depth', coordinate_variables=['lon', 'lat'])
>>> depth_indices
array([[0, 1, 2, 3, 4],
       [1, 1, 2, 3, 4],
       [2, 2, 2, 3, 4],
       [3, 3, 3, 3, 4],
       [4, 4, 4, 4, 4]])
>>> dataset_floor = dataset.some_selector(depth_indices)
>>> dataset_floor
Dimensions:  (t: 5, y: 5, x: 5)
Coordinates:
    time     (t) datetime64 2022-02-08 ...
    lon      (x) int64 0 -1 -2 -3 -4
    lat      (y) int64 0 1 2 3 4
Data variables:
    temp     (t, y, x) float64 0.0 1.0 2.0 3.0 4.0 1.0 ... 4.0 4.0 4.0 4.0 4.0

当前实现通过了以下测试功能。我之后的新实现将通过相同的测试，而不使用 bfill():

import numpy as np
import pandas as pd
import xarray as xr
from numpy.testing import assert_equal

from cemarray.operations import ocean_floor


def test_ocean_floor():
    # Values will be a 3D cube of values, with a slice along the x-axis like
    #     y
    #   44444
    #   3333.
    # d 222..
    #   11...
    #   0....
    values = np.full((5, 5, 5, 5), fill_value=np.nan)
    for i in range(5):
        values[:, i, :i + 1, :i + 1] = i

    temp = xr.DataArray(
        data=values,
        dims=['t', 'z', 'y', 'x'],
    )
    dataset = xr.Dataset(
        data_vars={"temp": temp},
        coords={
            'time': (['t'], pd.date_range('2022-02-08', periods=5)),
            'lon': (['x'], -np.arange(5)),
            'lat': (['y'], np.arange(5)),
            'depth': (['z'], 4.25 - np.arange(5), {'positive': 'down'}),
        }
    )

    floor_dataset = ocean_floor(dataset, ['depth'])

    assert floor_dataset.dims == {
        't': 5,
        'x': 5,
        'y': 5,
    }
    assert set(floor_dataset.coords.keys()) == {'time', 'lon', 'lat'}
    # We should see values for the deepest layer that has a value there
    expected_values = [
        [0, 1, 2, 3, 4],
        [1, 1, 2, 3, 4],
        [2, 2, 2, 3, 4],
        [3, 3, 3, 3, 4],
        [4, 4, 4, 4, 4],
    ]
    assert_equal(
        floor_dataset['temp'].values,
        np.array([expected_values] * 5)
    )

Answer 1

是的！ xarray 的 Advanced Indexing 也适用于许多维度！

创建您想到的索引器“some_selector”（它确实存在！），方法是使用值等于您想要 select 的坐标的 DataArray 进行索引，然后dimensions/coordinates匹配目标结果。在这种情况下，您需要一个由 x, y:

索引的 z 的 DataArray

>>> depth_indices = compute_ocean_floor_index(
    dataset, depth_variable='depth', coordinate_variables=['lon', 'lat'])
>>> depth_indices
array([[0, 1, 2, 3, 4],
       [1, 1, 2, 3, 4],
       [2, 2, 2, 3, 4],
       [3, 3, 3, 3, 4],
       [4, 4, 4, 4, 4]])
>>> selector = xr.DataArray(depth_indices, dims=('y', 'x'))

这个 selector 现在可以用来提取每个 (x, y) 对的 z 水平，忽略 t:

>>> dataset_floor = dataset.isel(z=selector)
>>> dataset_floor
Dimensions:  (t: 5, y: 5, x: 5)
Coordinates:
    time     (t) datetime64 2022-02-08 ...
    lon      (x) int64 0 -1 -2 -3 -4
    lat      (y) int64 0 1 2 3 4
Data variables:
    temp     (t, y, x) float64 0.0 1.0 2.0 3.0 4.0 1.0 ... 4.0 4.0 4.0 4.0 4.0

Select xarray 中海底的值

Select values along the ocean floor in xarray

python

numpy

python-xarray