使用条件和函数向量化嵌套循环
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）
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)
我可以使用广播来避免嵌套循环，如 link:
中所述
但在这种情况下，我必须调用函数 (i,j,xsize,ysize) 然后我有条件。
我真的需要知道 i 和 j 的值。
是否可以向量化此函数？
编辑：function(i,j,xsize,ysize) 将使用 sympy 对 return 浮点数执行符号计算。所以 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），你说过这个函数不能向量化所以我们可以只使用列表理解。我们想要一个具有相同形状 (xsize,ysize) 的矩阵 a，因此对于嵌套列表的理解有点倒退：
# 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
把它们放在一起就是完全向量化的函数：
# 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() 制作了一个可以改进的昂贵副本，但我不完全确定如何.
我有以下功能：
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)
我可以使用广播来避免嵌套循环，如 link:
中所述但在这种情况下，我必须调用函数 (i,j,xsize,ysize) 然后我有条件。 我真的需要知道 i 和 j 的值。
是否可以向量化此函数？
编辑：function(i,j,xsize,ysize) 将使用 sympy 对 return 浮点数执行符号计算。所以 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），你说过这个函数不能向量化所以我们可以只使用列表理解。我们想要一个具有相同形状 (xsize,ysize) 的矩阵 a，因此对于嵌套列表的理解有点倒退：
# 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
把它们放在一起就是完全向量化的函数：
# 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() 制作了一个可以改进的昂贵副本，但我不完全确定如何.