使用条件和函数向量化嵌套循环 [英] Vectorizing nested loop with conditionals and functions
问题描述
我有以下功能:
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)
There is one nuance here which is the fact that (example for 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 could avoid the nested loop using broadcasting as explained in this link:
但在这种情况下，我必须调用函数(i,j,xsize,ysize) 然后我有条件.我真的需要知道 i 和 j 的值.
But in this case I have to call the function(i,j,xsize,ysize) and then I have conditionals. I really need to know the value of i and j.
是否可以向量化这个函数?
Is it possible to vectorize this function?
function(i,j,xsize,ysize) 将使用 sympy 执行符号计算以返回浮点数.所以 a 是一个浮点数，而不是一个符号表达式.
the function(i,j,xsize,ysize) will use sympy to perform symbolic calculations to return a float. So a is a float, not a symbolic expression.
推荐答案
首先要注意的是你的函数 F(x) 可以描述为 x(idx) * weight(idx) 对于每个索引，其中权重仅取决于 x 的维度.因此，让我们根据函数 get_weights_for_shape 构建我们的代码，这样 F 就相当简单.为了简单起见，weights 将是一个 (xsize by size) 矩阵，但我们也可以让 F 用于平面输入:
The first thing to note is that your function F(x) can be described as x(idx) * weight(idx) for each index, where weight is only dependent on the dimensions of x. So lets structure our code in terms of a function get_weights_for_shape so that F is fairly simple. For sake of simplicity weights will be a (xsize by size) matrix but we can let F work for flat inputs too:
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) 所以对于嵌套列表的理解有点倒退:
So first we want to run your function (which I have aliased get_one_weight inside the parameters) for each element, you have said that this function can't be vectorized so we can just use list comprehension. We want a matrix a that has the same shape (xsize,ysize) so the comprehension is a bit backwards for a nested list:
# 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 将为条件赋值提供一个布尔数组:
With this matrix a > xsize will give a boolean array for conditional assignment:
case1 = a > xsize
weights[case1] = a[case1]
对于另一种情况，我们使用索引 i 和 j.要矢量化二维索引，我们可以使用 np.meshgrid
For the other case, we use the index i and j. To vectorize a 2D index we can use 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
将它们放在一起给出了完全矢量化的函数:
Putting it all together gives this as the fully vectorized function:
# 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
这涵盖了大部分内容.对于您的特定情况，因为这种繁重的计算仅依赖于输入的形状，如果您希望使用类似大小的输入重复调用此函数，则可以缓存所有先前计算的权重:
And that covers most of it. For your specific case since this heavy calculation only relies on the shape of the input, if you expect to call this function repeatedly with similarly sized input you can cache all previously calculated 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() 制作一个昂贵的副本，可能改进了，但我不完全确定如何.
As far as I can tell this seems like the most optimized you are going to get. The only improvements would be to vectorize function (even if that means just calling it in several threads in parallel) or possibly if .flatten() makes an expensive copy that could be improved but I'm not totally sure how.
这篇关于使用条件和函数向量化嵌套循环的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！
