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

我有以下功能：

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      

# 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

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

我的代码很简单

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

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

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

是否可以将此函数矢量化

---

编辑：

函数（i，j，xsize，ysize）

将使用sympy执行符号计算以返回浮点值。所以

a

是一个浮点数，而不是一个符号表达式

首先要注意的是，对于每个索引，您的函数

F（x）

可以描述为

x（idx）*权重（idx）

，其中权重仅取决于

x

的维度。因此，让我们根据函数

get\u weights\u for_shape

来构造代码，这样

F

就相当简单了。为简便起见，

权重

将是一个

（按大小X大小）

矩阵，但我们也可以让

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

首先我们要为每个元素运行

函数

（我在参数中使用了别名

get\u one\u weight

），您说过这个函数不能矢量化，所以我们可以使用列表理解。我们需要一个具有相同形状的矩阵

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

。我们可以使用

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

将这一切放在一起，可以得到完全矢量化的功能：

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      

# 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]

---

据我所知，这似乎是你将得到的最优化的。唯一的改进是将

函数

矢量化（即使这意味着在多个线程中并行调用它），或者如果

.flatte（）

制作了一个可以改进的昂贵副本，但我不完全确定如何改进。

函数做什么？看看这是否可以矢量化是非常重要的。是否应该将

count+=1

缩进到内部循环？现在它只在外部循环中，这意味着大多数赋值都在重写自身，而不是所有的

ff

都被初始化。函数相对复杂。它涉及点积和矩阵乘法的符号计算。

符号计算

？与sympy一样？如果是这样，请调整标记和示例。每个人，

函数

仅取决于输入的大小。这里只需要缓存给定大小的函数结果，然后将其用于相同大小的

x

的不同值。我无法执行此操作：

（xsize，ysize）=dims

。它给了我一个错误：没有足够的值来解包（预期为2，得到1）您没有显示如何获得

xsize

和

ysize

，因此我假设是直接从

x.shape

获取的，如果您的数据只有一维，那么

ysize

=1？我的解释很差。我的向量x总是一维的。它的大小取决于xsize和ysize，它们只是两个整数。如果ysize=2和xsize=3，则一维数组的大小将为ysize*xsize。这里的问题是，对于x[0]，我们有（i=0，j=0）x[1]（i=0，j=1）。这就是我有嵌套循环的原因。如果我可以直接从c中获得

I

和

j

，我就不需要循环了得到你想要的。我想提醒大家注意我在哪里使用了

np.meshgrid

，因为这给了我一个用于计算的矢量化索引，尽管当你的输入匹配循环的形状时它工作得更好。它工作得很好。它比第一个代码慢，但我认为它以更好的方式进行缩放。对于特定的大小，它可以I’我会比我猜的第一个代码更快。但是你能解释一下重用是什么意思吗？