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
这里有一个细微差别，即（例如 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)
我可以避免使用广播的嵌套循环，如此链接中所述：
但在这种情况下，我必须调用函数（i，j，xsize，ysize）然后我有条件。我真的需要知道 i 和 j 的值。
是否可以对这个函数进行矢量化？
编辑：function(i,j,xsize,ysize)将使用 sympy 执行符号计算以返回浮点数。浮点数也是如此 a，而不是符号表达式。
解决方案
首先要注意的是，您的函数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在参数中别名），你说这个函数不能向量化，所以我们可以使用列表理解。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()制作一个可以改进的昂贵副本，但我不完全确定如何。