Python中numpy多维数组切片自动生成方案咨询

解决方案 #

numpy数组原生支持将由slice对象、整数组成的元组作为索引传入，完全可以动态生成索引元组适配不同维度，无需手动写二维/三维分支判断。

具体实现逻辑如下：

• 动态生成两组fcDXidx的索引元组，分别对应两种偏移的切片逻辑
• 动态生成ccDXidx和M全选前D维的索引元组
• 直接用生成的索引元组完成赋值，无需做维度判断

优化后的完整代码如下：

def computeMatrixElements(M, fcDXidx, ccDXidx, dXi, d, ccregion):
    dxscale = 1./(dXi*dXi)
    D = len(ccregion.Size())
    N = list(ccregion.Size())
    N[d] += 1

    # 生成fcDXidx对应的两组切片元组
    slices_start = tuple(slice(0, N[dim] - ID[d][dim]) for dim in range(D)) + (0,)
    slices_end = tuple(slice(ID[d][dim], N[dim]) for dim in range(D)) + (0,)
    # 生成M、ccDXidx前D维全选的索引元组
    full_slices = tuple(slice(None) for _ in range(D)) + (0,)

    # 无需维度判断，统一赋值逻辑
    M[full_slices + (0,)] = ccDXidx[full_slices] * fcDXidx[slices_start] * dxscale
    M[full_slices + (1,)] = ccDXidx[full_slices] * fcDXidx[slices_end] * dxscale

补充说明：

• slice(start, end)等价于手动写的start:end切片语法
• slice(None)等价于手动写的:，代表当前维度全部选中
• 元组长度会自动匹配数组维度，2维和3维场景都可以正常运行，后续如果要扩展到4维也无需修改核心逻辑

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