Numpy 仅计算数学表达式，涉及同一数组上的嵌套函数总和

Question

我需要帮助来仅使用 numpy 运算来计算数学表达式。我要计算的表达式如下：

其中：x 是 (N, S) 数组，f 是 numpy 函数（可以与可广播数组一起使用，例如 np.maximum、np.sum、np.prod...）。如果这很重要，在我的例子中 f 是一个对称函数。

到目前为止，我的代码如下所示：

s = 0
for xp in x: # Loop over N...
    s += np.sum(np.prod(f(xp, x), axis=1))

还有我想摆脱的循环。

通常 N 是“大”（大约 30k）但 S 很小（小于 20），所以如果有人能找到只循环 S 的技巧，这仍然是一个重大改进。

我相信问题很容易通过对数组进行 N 复制，但其中一个大小 (32768、32768、20) 需要 150Go 的 RAM，而我没有。但是， (32768, 32768) 适合内存，但我希望有一个不分配此类数组的解决方案。

也许可以将 np.einsum 与精心挑选的数组一起使用？

感谢您的回复。如果缺少任何信息，请告诉我！

祝你有美好的一天！

编辑 1：

我感兴趣的 f 形式包括（目前）：f(x, y) = |x - y|, f(x, y) = |x - y|^2, f(x, y ) = 2 - max(x, y)。

Answer 1

您的循环非常有效。一些可能的方法是

方法一（在 S 上循环）

import numpy as np
def f(x,y):
    return np.abs(x-y)
N = 200
S = 20
x_data = random.rand(N,S) #(i,s)
y_data = random.rand(N,S) #(i',s)
product = f(broadcast_to(x_data[:,0][...,None],(N,N)) ,broadcast_to(y_data[:,0][...,None],(N,N)).T)
for i in range(1,S):
    product *= f(broadcast_to(x_data[:,i][...,None],(N,N)) ,broadcast_to(y_data[:,i][...,None],(N,N)).T)

sum = np.sum(product)

方法2（调度S个块）

import numpy as np
def f(x,y):
    x1 = np.broadcast_to(x[:,None,...],(x.shape[0],y.shape[0],x.shape[1]))
    y1 = np.broadcast_to(y[None,...],(x.shape[0],y.shape[0],x.shape[1]))
    return np.abs(x1-y1)
def f1(x1,y1):
    return np.abs(x1-y1)
N = 5000
S = 20
x_data = np.random.rand(N,S) #(i,s)
y_data = np.random.rand(N,S) #(i',s)
def fun_new(x_data1,y_data1):
    s = 0
    pp =np.split(x_data1,S,axis=0)
    for xp in pp:
        s += np.sum(np.prod(f(xp, y_data1), axis=2))
    return s
def fun_op(x_data1,y_data1):
    s = 0
    for xp in x_data1: # Loop over N...
        s += np.sum(np.prod(f1(xp, y_data1), axis=1))
    return s

fun_new(x_data,y_data)