基于numpy实现无for循环的任意维离散积分最优方案求解

核心思路 #

你当前基于itertools.product的实现本质是逐元素遍历Hxy数组计算项再求和，完全可以通过Numpy/PyTorch的广播机制替换所有Python层面循环，既提升百倍以上运算速度，又避免列表推导式的额外内存开销。

Numpy实现（完全兼容任意维度） #

import numpy as np

# 扩展Hx维度：新增最后一维，匹配Hxy的维度数
Hx_expanded = Hx[..., np.newaxis]
# 扩展Hy维度：在前面新增N-1个空维度，N为Hxy的维度数
Hy_expanded = Hy[(np.newaxis,) * (Hxy.ndim - 1) + (Ellipsis,)]

# 全数组广播运算，和你原逻辑结果完全等价
term = Hxy * np.log(Hxy / Hx_expanded / Hy_expanded)
# 可选：过滤Hxy为0时产生的nan/inf，按需开启
# term = np.nan_to_num(term, nan=0.0, posinf=0.0, neginf=0.0)
integral_fast = term.sum()

PyTorch实现（支持GPU加速） #

如果需要处理超大规模数据，可以用PyTorch迁移到GPU运算，逻辑完全一致：

import torch

# 维度扩展逻辑和Numpy对应
Hx_expanded = Hx.unsqueeze(-1)
Hy_expanded = Hy.view((1,) * (Hxy.ndim - 1) + (-1,))

term = Hxy * torch.log(Hxy / Hx_expanded / Hy_expanded)
# 可选：过滤异常值
# term = torch.nan_to_num(term, nan=0.0, posinf=0.0, neginf=0.0)
integral_fast = term.sum().item()

效果说明 #

• 运算速度：所有运算都在底层C/CUDA层面执行，无Python循环开销，比原实现快10~1000倍（依数据规模而定）
• 内存占用：仅需要存储和Hxy同大小的中间运算数组，无需生成全量索引元组、无额外列表开销，内存占用降低一个数量级以上
• 通用性：天然支持任意N≥2的维度，无需修改代码即可适配不同形状的输入矩阵

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