如何在xarray数据集上逐像素应用决策树并支持Dask分块？

问题背景 #

我有一个由多个对齐栅格层组成的xarray数据集，代表森林属性，需要应用决策树生成分类结果。由于数据覆盖超大区域，包含15+输入层和数十个终端节点，必须用Dask实现并行分块处理。

示例数据集 #

import xarray as xr

peat = [[1, 1], [0, 0]] 
fire_recent = [[1, 0], [1, 0]] 
lon = [[-99.83, -99.32], [-99.79, -99.23]]
lat = [[42.25, 42.21], [42.63, 42.59]]

ds_2x2 = xr.Dataset(
    data_vars=dict(
        peat=(["x", "y"], peat),
        fire_recent=(["x", "y"], fire_recent),
    ),
    coords=dict(
        lon=(["x", "y"], lon),
        lat=(["x", "y"], lat),
    ),
)

示例决策树（原错误实现） #

class LossInPeat():
  def __init__(self, no_peat, peat):
    self.no_peat = no_peat
    self.peat = peat
  
  def classify(self, ds_2x2):
    if ds_2x2.peat == 0:
      return 'NoPeat'
    else:
      return 'Peat'

报错情况 #

运行LossInPeat().classify(ds_2x2)时触发：

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

原因是if语句试图判断整个数组的布尔值，而非逐像素处理。

解决方案 #

核心思路是将决策树逻辑改为矢量化操作，配合xarray的apply_ufunc实现分块并行计算，天然支持Dask。

1. 修改决策树为矢量化逻辑 #

用numpy的矢量化条件判断替代标量if-else，直接处理数组：

import numpy as np

class LossInPeat():
    def __init__(self, no_peat=None, peat=None):
        self.no_peat = no_peat
        self.peat = peat
    
    def classify(self, peat_array):
        # 矢量化条件判断，逐像素生成结果
        return np.where(peat_array == 0, 'NoPeat', 'Peat')

2. 用apply_ufunc实现分块并行 #

通过apply_ufunc调用矢量化函数，同时启用Dask并行：

# 对数据集进行Dask分块（根据内存调整块大小）
ds_dask = ds_2x2.chunk({'x': 1000, 'y': 1000})

# 应用分类函数
result = xr.apply_ufunc(
    LossInPeat().classify,
    ds_dask['peat'],
    input_core_dims=[['x', 'y']],  # 指定要处理的空间维度
    output_dtypes=['U6'],  # 输出字符串类型的长度
    dask='parallelized'  # 启用Dask分块并行
)

# 查看结果
print(result)

3. 复杂多分支决策树扩展 #

如果是包含多输入、多分支的复杂决策树，用np.select实现多条件判断：

class ComplexForestClassifier():
    def classify(self, peat, fire_recent):
        # 定义多分支条件与对应标签
        conditions = [
            (peat == 0) & (fire_recent == 0),
            (peat == 0) & (fire_recent == 1),
            (peat == 1) & (fire_recent == 0),
            (peat == 1) & (fire_recent == 1)
        ]
        choices = [
            'NoPeat_NoFire',
            'NoPeat_RecentFire',
            'Peat_NoFire',
            'Peat_RecentFire'
        ]
        return np.select(conditions, choices, default='Unknown')

# 调用多输入分类函数
result_complex = xr.apply_ufunc(
    ComplexForestClassifier().classify,
    ds_dask['peat'],
    ds_dask['fire_recent'],
    input_core_dims=[['x', 'y'], ['x', 'y']],  # 每个输入变量的核心维度
    output_dtypes=['U20'],
    dask='parallelized'
)

关键注意事项 #

• 所有决策树逻辑必须用矢量化操作实现，禁止使用标量if判断数组
• 分块大小需根据内存和数据规模调整，推荐1000x1000或2000x2000的块尺寸
• 输出字符串标签时，要指定足够长度的output_dtypes，避免截断

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