首先要注意的是，您的函数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() 制作了一个可以改进但我不完全确定如何改进的昂贵副本。