Как развернуть 2D координаты lat/lon в 1D сетку в xarray

Пошагово в xarray: преобразуем 2D координаты lat/lon в 1D оси lat и lon, раскладываем значения на регулярную сетку, пропуски заполняем NaN. Подходит для NetCDF.

При работе с данными NetCDF в xarray часто возникает задача унификации: преобразовать многомерные массивы координат в стандартные одномерные координаты. Цель — разложить значения на регулярную широтно-долготную сетку и оставить пустые позиции как NaN, чтобы последующая логика могла опираться на согласованные оси.

Задача

Имея DataArray, у которого координаты заданы двумерными массивами широты/долготы, преобразовать его в DataArray с одномерными координатами lat и lon, переназначив значения соответствующим ячейкам сетки.

import xarray as xr

arr0 = xr.DataArray(
    [[0, 1], [2, 3]],
    coords={
        "lon": (["ny", "nx"], [[30, 40], [40, 50]]),
        "lat": (["ny", "nx"], [[10, 10], [20, 20]]),
    },
    dims=["ny", "nx"],
)

Требуемое представление с 1D-координатами и NaN там, где точка не отображается:

xr.DataArray(
    [[0, 1, np.nan], 
     [np.nan, 2, 3]],
    coords={
        "lat": [10, 20],
        "lon": [30, 40, 50],
    })

Что происходит и почему это так

Данные хранят значения на сетке с индексами ny и nx, а координаты lon и lat даны как двумерные массивы по тем же измерениям. Многие инструменты рассчитывают на 1D-оси lat и lon, чтобы увязать данные с регулярной сеткой. Чтобы привести структуру к стандартной форме, значения нужно перенести из пар 2D‑координат на наборы уникальных широт и долгот. Те позиции, которых не было в исходном сопоставлении, должны остаться NaN.

Решение

Простой подход: свернуть данные в список точек с помощью xarray.DataArray.stack, вычислить уникальные значения координат, создать пустую регулярную сетку и вернуть значения на места по координатному поиску.

import xarray as xr
import numpy as np

arr0 = xr.DataArray(
    [[0, 1], [2, 3]],
    coords={
        "lon": (["ny", "nx"], [[30, 40], [40, 50]]),
        "lat": (["ny", "nx"], [[10, 10], [20, 20]]),
    },
    dims=["ny", "nx"],
)

# Свернуть ny/nx в единый индекс
stacked = arr0.stack(idx=("ny", "nx"))

# Собрать уникальные значения координат
lat_unique = np.unique(arr0.lat.values)
lon_unique = np.unique(arr0.lon.values)

# Подготовить регулярную сетку, заполненную NaN
regular = xr.DataArray(
    np.full((len(lat_unique), len(lon_unique)), np.nan),
    coords={"lat": lat_unique, "lon": lon_unique},
    dims=["lat", "lon"]
)

# Отобразить каждую исходную точку в соответствующее положение 1D lat/lon
for k in range(stacked.size):
    y = float(stacked.lat.values[k])
    x = float(stacked.lon.values[k])
    v = stacked.values[k]
    regular.loc[dict(lat=y, lon=x)] = v

print(regular)

Результат:

<xarray.DataArray (lat: 2, lon: 3)> Size: 48B
array([[ 0.,  1., nan],
       [nan,  2.,  3.]])
Coordinates:
  * lat      (lat) int64 16B 10 20
  * lon      (lon) int64 24B 30 40 50

Почему это важно

Стандартные 1D‑оси lat и lon упрощают объединение наборов данных, которые различаются лишь способом задания координат. После ремаппинга индексация по широте/долготе становится прямой, а отсутствующие комбинации явно остаются как NaN, что делает дальнейшую обработку предсказуемой.

Выводы

Если DataArray хранит координаты как 2D‑массива по пространственным измерениям, сверните исходную сетку в набор точек, извлеките уникальные широты и долготы, выделите регулярную сетку и переназначьте данные по координатному соответствию. Так вы выровняете данные по стандартным осям, сохраните исходные значения там, где координаты совпадают, и оставите несуществующие ячейки как NaN.