Uproot与Dask结合时二维数组轴计算异常问题咨询

我尝试用Uproot的Dask功能将ROOT文件中的分支加载为Dask数组，但执行计算时遇到异常：

示例代码：

import dask.array as da

tree = uproot.dask("file.root:tree", library = 'np')
branch_data = tree["testbranch"]
mean = da.mean(branch_data).compute()

该分支数据为二维数组，我希望沿axis=1（每行）计算均值，但结果却是列均值，与np.mean(branch_data.T, axis = 1)的输出一致。添加axis=1参数时，报错“axis out of bounds for array with dimension 1”，但调用compute()后确认数据确实是二维的，da.sum()等方法也存在相同问题。

补充复现代码：

import numpy as np
import uproot

# 创建含树结构的示例文件
with uproot.recreate("test.root") as file:
    file["test_tree"] = {"test_branch": np.random.random((100,10))}

# 标准Uproot方式（期望Dask方式得到相同输出）
tree = uproot.open("./test.root:test_tree")
branch = tree["test_branch"].array(library = 'np')
mean = np.mean(branch, axis = 1)
print(mean)

# Uproot-Dask方式（计算列均值，预期应为每行的均值，结果异常）
tree = uproot.dask("./test.root:test_tree", library = 'np')
branch = tree["test_branch"]
mean = np.mean(branch).compute()
print(mean)

# 预期对应标准Uproot输出但无法运行，报错axis越界
mean = np.mean(branch, axis = 1).compute()
print(mean)

解决方案 #

原因分析 #

Uproot的Dask接口返回的数组并非真正的二维Dask数组，而是一维Dask数组，每个元素对应ROOT分支中一个条目的二维子数组。也就是说，数组形状是(n_entries,)而非(n_entries, n_cols)，这和标准Uproot直接加载为二维NumPy数组的行为不同，因此直接使用Dask的轴参数会触发维度越界错误，默认计算也会误将每个子数组当作单个元素处理，得到列均值。

方法1：用da.map_blocks逐元素计算行均值 #

针对每个子数组单独计算行均值，通过map_blocks批量处理：

import dask.array as da

tree = uproot.dask("./test.root:test_tree", library='np')
branch = tree["test_branch"]

def row_mean(arr):
    return np.mean(arr, axis=1)

# 指定输出数据类型，确保Dask能正确推断结果形状
mean_dask = da.map_blocks(row_mean, branch, dtype=np.float64)
mean = mean_dask.compute()
print(mean)

方法2：拼接为二维Dask数组后计算 #

如果所有子数组的形状一致，可将一维Dask数组拼接为真正的二维数组，再按轴计算：

import dask.array as da

tree = uproot.dask("./test.root:test_tree", library='np')
branch = tree["test_branch"]

# 将每个子数组堆叠成二维数组
branch_2d = da.vstack(branch)
mean = da.mean(branch_2d, axis=1).compute()
print(mean)

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