python - 将 argsort 与掩码相结合以在移动窗口中获取最接近的值

Question

我有一些代码可以根据二维圆形窗口中的相邻值计算图像中的缺失值。它还使用来自同一位置的一个或多个时间相邻图像的值(即在 3 维中移动的相同 2D 窗口)。

对于每个缺失的位置，我需要计算的值不一定基于整个窗口中所有可用的值，而是仅基于空间上最近的 n 个具有值的单元格(在两个图像/Z 轴中位置)，其中 n 是小于 2D 窗口中单元格总数的某个值。

目前，计算窗口中的所有内容要快得多，因为我的排序方法是获取最近的 n 个包含数据的单元格，这是函数中最慢的部分，因为即使距离很远，每次都必须重复它在窗口坐标方面不改变。我不确定这是必要的，我觉得我必须能够一次获得排序的距离，然后在仅选择可用单元格的过程中屏蔽那些距离。

这是我的代码，用于选择要在间隙单元位置的窗口内使用的数据:

# radius will in reality be ~100
radius = 2
y,x = np.ogrid[-radius:radius+1, -radius:radius+1]
dist = np.sqrt(x**2 + y**2)
circle_template = dist > radius

# this will in reality be a very large 3 dimensional array
# representing daily images with some gaps, indicated by 0s
dataStack = np.zeros((2,5,5))
dataStack[1] = (np.random.random(25) * 100).reshape(dist.shape)
dataStack[0] = (np.random.random(25) * 100).reshape(dist.shape)

testdata = dataStack[1]
alternatedata = dataStack[0]
random_gap_locations = (np.random.random(25) * 30).reshape(dist.shape) > testdata
testdata[random_gap_locations] = 0
testdata[radius, radius] = 0

# in reality we will go through every gap (zero) location in the data 
# for each image and for each gap use slicing to get a window of 
# size (radius*2+1, radius*2+1) around it from each image, with the 
# gap being at the centre i.e. 
# testgaplocation = [radius, radius]
# and the variables testdata, alternatedata below will refer to these 
# slices

locations_to_exclude = np.logical_or(circle_template, np.logical_or
                                     (testdata==0, alternatedata==0))
# the places that are inside the circular mask and where both images 
# have data 
locations_to_include = ~locations_to_exclude
number_available = np.count_nonzero(locations_to_include)

# we only want to do the interpolation calculations from the nearest n 
# locations that have data available, n will be ~100 in reality
number_required = 3

available_distances = dist[locations_to_include]
available_data = testdata[locations_to_include]
available_alternates = alternatedata[locations_to_include]

if number_available > number_required:
    # In this case we need to find the closest number_required of elements, based
    # on distances recorded in dist, from available_data and available_alternates
    # Having to repeat this argsort for each gap cell calculation is slow and feels 
    # like it should be avoidable
    sortedDistanceIndices = available_distances.argsort(kind = 'mergesort',axis=None)
    requiredIndices = sortedDistanceIndices[0:number_required]
    selected_data = np.take(available_data, requiredIndices)
    selected_alternates = np.take(available_alternates , requiredIndices)
else:
    # we just use available_data and available_alternates as they are...

# now do stuff with the selected data to calculate a value for the gap cell

这行得通，但是函数总时间的一半以上都花在了屏蔽空间距离数据的 argsort 上。 (总共 1.4mS 中的 ~900uS - 这个函数将运行数百亿次，所以这是一个重要的区别!)

我确信我必须能够在函数外部执行此 argsort 一次，当最初设置空间距离窗口时，然后将这些排序索引包含在掩码中，以获得第一个 howManyToCalculate 索引而无需重新排序。答案可能涉及将我们从中提取的各种位放入记录数组中——但如果是这样的话，我不知道怎么做。任何人都可以看到如何使流程的这一部分更有效率吗？

Answer 1

所以你想在循环外进行排序:

sorted_dist_idcs = dist.argsort(kind='mergesort', axis=None)

然后使用原始代码中的一些变量，这就是我能想到的，尽管它仍然感觉像是一个主要的往返..

loc_to_incl_sorted = locations_to_include.take(sorted_dist_idcs)
sorted_dist_idcs_to_incl = sorted_dist_idcs[loc_to_incl_sorted]

required_idcs = sorted_dist_idcs_to_incl[:number_required]
selected_data = testdata.take(required_idcs)
selected_alternates = alternatedata.take(required_idcs)

请注意，required_idcs 指的是 testdata 中的位置，而不是原始代码中的 available_data。而我使用 take 的这个片段是为了方便地索引扁平化数组。