Question

我具有以下功能：

def F(x):    #F receives a numpy vector (x) with size (xsize*ysize)

    ff = np.zeros(xsize*ysize)

    count=0 

    for i in range(xsize):
       for j in range(ysize):

           a=function(i,j,xsize,ysize)

           if (a>xsize):
               ff[count] = x[count]*a
           else
               ff[count] = x[count]*i*j

           count = count +1

    return ff

这里有一个细微差别，那就是（例如xsize = 4，ysize = 3）

c=count
x[c=0] corresponds to x00 i=0,j=0
x[c=1]   x01   i=0, j=1
x[c=2]   x02   i=0, j=2   (i=0,  j = ysize-1)
x[c=3]   x10   i=1, j=0
...   ...  ...
x[c=n]   x32    i=3 j=2  (i=xsize-1, j=ysize-1)

我的代码就是

ff[c] = F[x[c]*a (condition 1)
ff[c] = F[x[c]*i*j (condition 2)

我可以使用广播来避免嵌套循环，如此链接所述：

Python3: vectorizing nested loops

但是在这种情况下，我必须调用函数（i，j，xsize，ysize），然后才有条件。我真的需要知道i和j的值。

可以矢量化此功能吗？

编辑：function(i,j,xsize,ysize)将使用sympy进行符号计算以返回浮点数。因此a是浮点数，而不是符号表达式。

Answer 1

首先要注意的是，对于每个索引，您的函数F(x)可以描述为x(idx) * weight(idx)，其中权重仅取决于x的维数。因此，让我们根据函数get_weights_for_shape来构造代码，以便F非常简单。为简单起见，weights将是一个(xsize by size)矩阵，但我们也可以让F用于平面输入：

def F(x, xsize=None, ysize=None):
    if len(x.shape) == 2:
        # based on how you have put together your question this seems like the most reasonable representation.
        weights = get_weights_for_shape(*x.shape)
        return x * weights
    elif len(x.shape) == 1 and xsize * ysize == x.shape[0]:
        # single dimensional input with explicit size, use flattened weights.
        weights = get_weights_for_shape(xsize, ysize)
        return x * weights.flatten()
    else:
        raise TypeError("must take 2D input or 1d input with valid xsize and ysize")


# note that get_one_weight=function can be replaced with your actual function.
def get_weights_for_shape(xsize, ysize, get_one_weight=function):
    """returns weights matrix for F for given input shape"""
    # will use (xsize, ysize) shape for these calculations.
    weights = np.zeros((xsize,ysize))
    #TODO: will fill in calculations here
    return weights

因此，首先我们要为每个元素运行function（我在参数中为get_one_weight加上别名），您已经说过该函数不能被矢量化，因此我们可以使用list理解。我们希望矩阵a的形状为(xsize,ysize)，因此对于嵌套列表的理解有点倒退：

# notice that the nested list makes the loops in opposite order:
# [ROW for i in Xs]
#  ROW = [f() for j in Ys]
a = np.array([[get_one_weight(i,j,xsize,ysize)
                    for j in range(ysize)
              ] for i in range(xsize)])

有了这个矩阵a > xsize，将为条件赋值提供一个布尔数组：

case1 = a > xsize
weights[case1] = a[case1]

对于另一种情况，我们使用索引i和j。要向量化2D索引，我们可以使用np.meshgrid

[i,j] = np.meshgrid(range(xsize), range(ysize), indexing='ij')
case2 = ~case1 # could have other cases, in this case it's just the rest.
weights[case2] = i[case2] * j[case2]

return weights #that covers all the calculations

将所有内容放在一起可将其作为完全矢量化的功能：

# note that get_one_weight=function can be replaced with your actual function.
def get_weights_for_shape(xsize, ysize, get_one_weight=function):
    """returns weights matrix for F for given input shape"""
    # will use (xsize, ysize) shape for these calculations.
    weights = np.zeros((xsize,ysize))

    # notice that the nested list makes the loop order confusing:
    # [ROW for i in Xs]
    #  ROW = [f() for j in Ys]
    a = np.array([[get_one_weight(i,j,xsize,ysize)
                        for j in range(ysize)
                  ] for i in range(xsize)])

    case1 = (a > xsize)
    weights[case1] = a[case1]

    # meshgrid lets us use indices i and j as vectorized matrices.
    [i,j] = np.meshgrid(range(xsize), range(ysize), indexing='ij')
    case2 = ~case1
    weights[case2] = i[case2] * j[case2]
    #could have more than 2 cases if applicable.

    return weights

涵盖了大部分内容。对于您的特定情况，由于这种繁重的计算仅依赖于输入的形状，因此如果您希望使用类似大小的输入重复调用此函数，则可以缓存所有先前计算的权重：

def get_weights_for_shape(xsize, ysize, _cached_weights={}):
    if (xsize, ysize) not in _cached_weights:
        #assume we added an underscore to real function written above
        _cached_weights[xsize,ysize] = _get_weights_for_shape(xsize, ysize)
    return _cached_weights[xsize,ysize]

据我所知，这似乎是您将获得的最佳优化。唯一的改进是对function进行矢量化处理（即使这意味着仅在多个线程中并行调用它），或者可能.flatten()制作了可以改进的昂贵副本，但我不确定如何

使用条件和函数向量化嵌套循环

1 个答案: