Как в python xarray создавать и подмножество ленивых переменных, не загружая весь массив данных? ⇐ Python

[Anonymous]

Я пытаюсь создать функцию Python, которая открывает удаленный набор данных (на сервере opendap) с помощью xarray и автоматически лениво создает новые переменные. Вариант использования — вычисление величины и направления, когда доступны компоненты u и v, например:

Код: Выделить всё

import xarray as xr
import dask
import dask.array as da

@dask.delayed
def uv2mag(u, v):
return (u**2 + v**2)**0.5

@dask.delayed
def uv2dir(u, v):
return np.rad2deg(np.arctan2(u, v))

def open_dataset(*args, **kwargs) -> xr.Dataset:

uv = kwargs.pop("uv", None)

ds = xr.open_dataset(*args, **kwargs)

if uv:

uvar, vvar = uv

ds["magnitude"] = (
ds[uvar].dims,
da.from_delayed(uv2mag(ds[uvar], ds[vvar]),
ds[uvar].shape,
dtype=ds[uvar].dtype),
{"long_name": "magnitude"},
)

ds["direction"] = (
ds[uvar].dims,
da.from_delayed(uv2dir(ds[uvar], ds[vvar]),
ds[uvar].shape,
dtype=ds[uvar].dtype),
{"long_name": "direction"},
)

return ds

url = "https://tds.hycom.org/thredds/dodsC/FMRC_ESPC-D-V02_uv3z/FMRC_ESPC-D-V02_uv3z_best.ncd"
uvar = "water_u"
vvar = "water_v"

ds = open_dataset(url, drop_variables="tau", uv=[uvar, vvar])

На первый взгляд кажется, что все работает нормально, и созданы новые переменные величины и направления.

Код: Выделить всё

Dimensions:      (depth: 40, lat: 4251, lon: 4500, time: 121)
Coordinates:
* depth        (depth) float64 0.0 2.0 4.0 6.0 ... 2.5e+03 3e+03 4e+03 5e+03
* lat          (lat) float64 -80.0 -79.96 -79.92 -79.88 ... 89.92 89.96 90.0
* lon          (lon) float64 0.0 0.07996 0.16 0.24 ... 359.7 359.8 359.8 359.9
* time         (time) datetime64[ns] 2024-09-29T12:00:00 ... 2024-10-14T12:...
time_run     (time) datetime64[ns] ...
Data variables:
time_offset  (time) datetime64[ns] ...
water_u      (time, depth, lat, lon) float32 ...
water_v      (time, depth, lat, lon) float32 ...
magnitude    (time, depth, lat, lon) float32 dask.array
direction    (time, depth, lat, lon) float32 dask.array
Attributes: (12/22)
classification_level:
distribution_statement:    Approved for public release; distribution unli...
downgrade_date:            not applicable
classification_authority:  not applicable
institution:               Fleet Numerical Meteorology and Oceanography C...
source:                    HYCOM archive file, GLBz0.04
...
time_origin:               2024-10-05 12:00:00
_CoordSysBuilder:          ucar.nc2.dataset.conv.CF1Convention
cdm_data_type:             GRID
featureType:               GRID
location:                  Proto fmrc:FMRC_ESPC-D-V02_uv3z
history:                   FMRC Best Dataset

Проблема возникает, когда я пытаюсь получить очень небольшую часть данных. «Настоящая» переменная мгновенно возвращает данные, например:

Код: Выделить всё

>>> ds.isel(time=slice(0, 5), depth=0, lat=1000, lon=1000)["water_u"].values
array([0.27400002, 0.23600002, 0.12      , 0.108     , 0.24400002],
dtype=float32)

Но когда я пытаюсь получить подмножество из ленивой переменной, возникает ошибка из-за попытки загрузить в память всю 4D-матрицу, а не только подмножество.

Код: Выделить всё

>>> ds.isel(time=slice(0, 5), depth=0, lat=1000, lon=1000)["magnitude"].values
MemoryError: Unable to allocate 345. GiB for an array with shape (121, 40, 4251, 4500) and data type float32

Есть ли у кого-нибудь предложения, как можно исправить тонкость?
Спасибо.
Мои ожидания заключалось в том, что подмножество, выполняемое isel, также будет применяться к ленивым переменным, и что загрузка в память будет такой же быстрой, как и «реальные» переменные.

Подробнее здесь: https://stackoverflow.com/questions/790 ... oading-the