如何将分块处理的Xarray数据保存为单个Zarr文件而非多个文件？

处理高维度大气数据（Time:5200000、south_north:434、west_east:653），需为每个Time与south_north/west_east的组合映射功率曲线，采用south_north=20、west_east=20的分块大小。当前代码生成600多个独立Zarr分块文件，需求是将结果保存为单个名为Power_curve的Zarr目录，且保留原数据集的维度顺序（确保后续脚本能将实际坐标映射到south_north/west_east点）。

核心思路是预先创建与原数据集维度匹配的空Zarr数据集并指定分块规则，循环计算每个分块后直接写入目标Zarr的对应切片位置，既保证分块处理的内存效率，又能得到单个Zarr目录。

import xarray as xr
import numpy as np

# 假设ds是已加载的原大气数据集
# ds = xr.open_zarr(...)

# 定义分块大小
chunk_size = {'south_north': 20, 'west_east': 20}
# Time维度分块可按需调整，这里沿用原数据集的分块，或设置固定值如10000
time_chunk = ds.chunks['Time'][0] if 'Time' in ds.chunks else 10000

# 创建目标Zarr数据集的结构模板，保留原维度与坐标
template = xr.DataArray(
    data=None,
    dims=('Time', 'south_north', 'west_east'),
    coords={
        'Time': ds.Time,
        'south_north': ds.south_north,
        'west_east': ds.west_east,
        'lat': ds.lat,
        'lon': ds.lon
    },
    name='Power_Curve'
).chunk(chunks={'Time': time_chunk, **chunk_size})

# 初始化空Zarr目录
target_path = "S:\\uv_ds_120m_with_density\\Power_curve.zarr"
template.to_dataset(name='Power_Curve').to_zarr(target_path, mode='w', compute=False)

# 打开目标数据集用于写入
target_ds = xr.open_zarr(target_path, mode='a')

# 计算空间分块数量
num_chunks_sn = int(np.ceil(len(ds.south_north) / chunk_size['south_north']))
num_chunks_we = int(np.ceil(len(ds.west_east) / chunk_size['west_east']))

for i in range(num_chunks_sn):
    for j in range(num_chunks_we):
        # 计算当前分块的切片范围（处理最后一个分块的边界）
        start_sn = i * chunk_size['south_north']
        end_sn = min((i+1)*chunk_size['south_north'], len(ds.south_north))
        slice_sn = slice(start_sn, end_sn)
        
        start_we = j * chunk_size['west_east']
        end_we = min((j+1)*chunk_size['west_east'], len(ds.west_east))
        slice_we = slice(start_we, end_we)
        
        # 提取当前分块数据
        subset = ds.sel(south_north=slice_sn, west_east=slice_we)
        
        # 功率曲线计算逻辑（与原代码一致）
        ws = np.sqrt(np.square(subset.U) + np.square(subset.V))
        WH = np.ceil(ws * 2) / 2
        WL = np.floor(ws * 2) / 2
        Rho_H = (np.ceil(subset.RHO * 40) / 40)
        Rho_L = (np.floor(subset.RHO * 40) / 40)
        
        WH = WH.where(WH > 3.0, 0)
        WH = WH.where(WH < 24.5, 24.5)
        
        WL = WL.where(WL > 3, 0)
        WL = WL.where(WL < 24.5, 24.5)

        Rho_L = Rho_L.where(Rho_L > 0.95, 0.95)
        Rho_L = Rho_L.where(Rho_L < 1.275, 1.275)
        
        Rho_L = Rho_L.astype(str)

        # 基于查找表计算功率
        power = da.sel(row=WH, column=Rho_L) / 2
        power.name = 'Power_Curve'
        
        # 将结果写入目标Zarr的对应位置
        target_ds['Power_Curve'].loc[{'south_north': slice_sn, 'west_east': slice_we}] = power

# 关闭数据集
target_ds.close()

1. 维度与坐标保留：通过模板数组完全保留原数据集的维度顺序和lat、lon坐标，满足后续脚本的坐标映射需求
2. 分块写入优化：直接对目标Zarr的指定切片赋值，避免生成大量独立文件，同时利用分块控制内存占用
3. 边界处理：计算分块范围时加入min函数处理最后一个分块的边界问题，避免超出数据集维度

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