如何利用Xarray的维度与数据创建新数组并实现多维度拓展

解决思路 #

Xarray的核心优势就是基于维度标签的自动广播，完全不需要手动调用np.outer处理维度拓展，也不需要写循环。你原来的问题出在np.outer会强制将输入的高维数组展平为1维再计算外积，自然得不到预期的三维输出。

修改后的可复用实现 #

import numpy as np
import xarray as xr

# 初始化数据集部分仅修正了参数类型bug
data = xr.Dataset(
    data_vars={
        "abs": (("x"), np.ones((200)))
    },
    coords={
        "x": np.linspace(1.5,1.55,200),
        "y": np.arange(int(5e5)),
        "z": np.arange(125)
    }
)

def intensity_decay(da):
    "Return decay curve (exponential decay)."
    init_power = 1e5
    # 直接利用xarray维度广播，自动生成(x,y)维度的数组
    decay = init_power * np.exp(-1 * da.abs * (da.y * 10 ** (-7)))
    decay = decay.where(decay >= 1e-3, 0)
    return decay

# 直接赋值，不需要手动指定维度，函数返回的DataArray自带维度信息
data["depth_decay"] = intensity_decay(data)

def radial_decay(da):
    "Return radial decay curve."
    # 自动广播(x,y)维度的depth_decay和(z)维度的z数组，生成(x,y,z)三维数组
    return da.depth_decay * np.exp(-da.z)

# 测试输出形状，会得到(3, 500000, 125)，符合预期
print(radial_decay(data.isel(x=[1, 4, 10])).shape)

关键优化点说明 #

• 所有运算直接使用Xarray的DataArray对象，保留维度标签，自动处理广播对齐，不需要手动管理维度顺序
• 用Xarray自带的where方法替代numpy索引赋值，避免丢失维度和坐标信息
• 函数返回的都是带维度信息的DataArray，赋值到Dataset时不需要手动指定维度，可复用性更强
• 支持任意维度的输入拓展，后续如果要加第四个维度，只需要按照同样逻辑直接和对应维度的变量相乘即可，不需要修改广播逻辑

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