如何在Dask数组中实现径向求和？

我需要对一个Dask数组执行径向求和操作，要求保留分块数据并按每个半径进行求和，还可能需要按求和的「像素」总数进行归一化。

该Dask数组的结构如下：

• Dask数组，shape=(264, 256, 1500)，dtype=int16，chunksize=(264, 256, 992)，chunktype=numpy.ndarray
• 数据为强度（计数），维度说明：
  • x：距离（nm），尺寸为(264,)
  • y：距离（nm），尺寸为(256,)
  • energy_scale：能量（eV），尺寸为(1500,)

我已编写以下代码实现圆形掩码并创建新数组，效果良好，但还需实现点的径向求和：

import numpy as np

def cartesian_to_polar(x, y, center):
    x_rel = x - center[1]
    y_rel = y - center[0]
    radius = np.sqrt(x_rel**2 + y_rel**2)
    theta = np.arctan2(y_rel, x_rel)
    return radius, theta

def get_polar_coordinates_grid(center, max_radius, dimensions):
    x = np.arange(dimensions[1])
    y = np.arange(dimensions[0])
    xv, yv = np.meshgrid(x, y)
    rv, thetav = cartesian_to_polar(xv, yv, center)
    
    # Mask to select points within the specified radius
    mask = rv <= max_radius
    
    mask_3D = np.repeat(mask[:, :, np.newaxis], dimensions[2], axis=2)
    
    print('mask_3D shape',mask_3D.shape)
    
    return mask_3D, rv, thetav

def polar_profile(dataset, center, max_radius):
    print('center', center[0],center[1])

    mask_3D, rv, thetav = get_polar_coordinates_grid(center, max_radius, dataset.shape[:3])
    radial_dataset_array = np.repeat(rv[:, :, np.newaxis], dataset.shape[2], axis=2)

    masked_dataset = np.where(mask_3D, dataset, 0)

    masked_r_dataset = np.where(mask_3D, radial_dataset_array, 0)

目前我不确定索引需要求和的像素的最佳方式，希望能提供一个简单NumPy数组的示例作为参考。

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