如何合并无公共坐标的二维xarray DataArray并保留原始坐标

我有多个二维xarray.DataArray，它们来自对大网格数据集的非重叠边界框切片，坐标值不一定重合。想要沿新维度fid（每个边界框的字符串标识）合并这些数组，同时避免原始二维坐标被扩展、填充nan，还要保留每个切片的原始全局坐标。

示例数据与问题重现 #

首先创建测试数据并切片得到两个子数组：

import xarray as xr
import numpy as np

# 创建网格数据
nx = 9
ny = 9
data = np.random.randint(5, size=(nx, ny))
x_coord = np.linspace(0, 1, nx)
y_coord = np.linspace(0, 1, ny)
da = xr.DataArray(
    data,
    dims=("x_coord", "y_coord"),
    coords={"x_coord": x_coord, "y_coord": y_coord}
)

# 切分两个非重叠子区域
a = da.isel(x_coord=[1, 2, 3], y_coord=[2, 3, 4]).expand_dims(fid=["abc123"])
b = da.isel(x_coord=[6, 7, 8], y_coord=[5, 6, 7]).expand_dims(fid=["def456"])

直接使用xr.concat会扩展全局坐标，导致数组填充大量nan：

xr.concat([a, b], dim="fid")

结果会生成形状为(2,6,6)的数组，大部分位置为nan，且全局坐标被合并为所有子数组坐标的并集。

解决方案 #

我们可以通过局部坐标系+多维全局坐标的方式实现需求，步骤如下：

1. 为每个子数组添加局部坐标 #

为每个子数组创建局部坐标xi/yi（从0开始的相对坐标），同时保留原始全局坐标作为临时变量：

# 处理数组a
a_local = a.assign_coords(
    xi=a.x_coord - a.x_coord.min(),
    yi=a.y_coord - a.y_coord.min(),
    # 保留原始全局坐标
    x_original=a.x_coord,
    y_original=a.y_coord
).rename({"x_coord": "xi", "y_coord": "yi"})

# 处理数组b
b_local = b.assign_coords(
    xi=b.x_coord - b.x_coord.min(),
    yi=b.y_coord - b.y_coord.min(),
    x_original=b.x_coord,
    y_original=b.y_coord
).rename({"x_coord": "xi", "y_coord": "yi"})

2. 合并数组并整理坐标 #

合并后，将临时保存的原始坐标转换为(fid, xi)和(fid, yi)维度的多维坐标：

# 沿fid维度合并
combined = xr.concat([a_local, b_local], dim="fid")

# 将原始坐标转换为多维坐标
combined = combined.assign_coords(
    x_coord=(["fid", "xi"], combined.x_original.data),
    y_coord=(["fid", "yi"], combined.y_original.data)
).drop(["x_original", "y_original"])

最终结果 #

此时combined的形状为(2,3,3)，且保留了所需的坐标结构：

print(combined)

输出示例：

<xarray.DataArray (fid: 2, xi: 3, yi: 3)>
array([[[3, 4, 3],
        [3, 1, 0],
        [3, 2, 4]],

       [[4, 3, 0],
        [3, 2, 2],
        [4, 2, 1]]])
Coordinates:
  * fid      (fid) object 'abc123' 'def456'
  * xi       (xi) float64 0.0 0.125 0.25
  * yi       (yi) float64 0.0 0.125 0.25
    x_coord  (fid, xi) float64 0.125 0.25 0.375 0.75 0.875 1.0
    y_coord  (fid, yi) float64 0.25 0.375 0.5 0.625 0.75 0.875

内容的提问来源于stack exchange，提问作者ogb119