Numpy低内存占用大型矩阵乘法解决方案求助

问题根源 #

原代码中np.exp(1j*np.outer(a1,a2))会生成100000×100000的复数矩阵，单个复数元素占16字节，总内存需求约160GB，这是导致内存溢出的核心原因。完全不需要生成完整的A矩阵，通过数学推导可直接计算结果，从根源降低内存占用。

数学推导 #

设 ( A_{i,j} = e^{i a_1[i] a_2[j]} )，则矩阵乘积 ( (A \cdot b)[i] = \sum_{j=0}^{n-1} e^{i a_1[i] a_2[j]} b[j] )，最终结果 ( d[i] = c[i] \times \sum_{j=0}^{n-1} e^{i a_1[i] a_2[j]} b[j] )。基于此推导可避免生成巨型矩阵。

方法1：Numpy分块迭代（低内存，无需额外库） #

核心思路：分块遍历a1的元素，每次计算一小块的结果，控制临时数据的内存占用，适配绝大多数机器。

import numpy as np

m = 100000
n = 100000
a1 = np.random.rand(m)
a2 = np.random.rand(n)
c = np.random.rand(m) + 1j * np.random.rand(m)
b = np.random.rand(n) + 1j * np.random.rand(n)

# 根据可用内存调整块大小，示例设为1000，内存占用约1.6GB/块
block_size = 1000
d = np.empty(m, dtype=np.complex128)

for i in range(0, m, block_size):
    # 截取当前处理块的a1和c元素
    a1_block = a1[i:i+block_size]
    c_block = c[i:i+block_size]
    # 计算当前块的矩阵乘积，临时生成block_size×n的小矩阵
    dot_block = np.dot(np.exp(1j * np.outer(a1_block, a2)), b)
    # 存入结果数组
    d[i:i+block_size] = c_block * dot_block

方法2：Numba加速（极低内存，适合资源有限的机器） #

针对逐元素计算优化，完全不生成中间矩阵，内存仅用于存储一维数组。之前使用Numba未成功大概率是未采用逐元素计算的逻辑：

import numpy as np
from numba import njit

@njit(parallel=True)  # 开启多线程并行，可根据需求关闭
def compute_d(a1, a2, c, b):
    m = len(a1)
    n = len(a2)
    d = np.empty(m, dtype=np.complex128)
    for i in range(m):
        sum_val = 0j
        ai = a1[i]
        for j in range(n):
            sum_val += np.exp(1j * ai * a2[j]) * b[j]
        d[i] = c[i] * sum_val
    return d

m = 100000
n = 100000
a1 = np.random.rand(m)
a2 = np.random.rand(n)
c = np.random.rand(m) + 1j * np.random.rand(m)
b = np.random.rand(n) + 1j * np.random.rand(n)

d = compute_d(a1, a2, c, b)

关于之前Dask失败的说明 #

直接将Numpy数组转为Dask数组再执行outer和dot，本质仍在尝试生成巨型矩阵，只是拆分了计算步骤，无法解决内存瓶颈。若要用Dask，需基于上述数学推导实现分块计算，复杂度高于前两种方法，因此不优先推荐。

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