首页 > 解决方案 > Python有效地在部分更新的数组中找到最大值索引

问题描述

问题:如何在 Python中以(相对)少量已知但变化的索引最有效地在循环迭代期间部分更新的实值一维数组中找到(全局)最大值的索引?特别是,如果一维数组的大小增加而更新值的数量保持有限,我如何保持每次迭代的计算时间不变?

背景:我正在以一种有效的变体 ( Clark, 1980 ) 实施 CLEAN 反卷积算法 ( Högbom, 1974 ; Roberts+, 1987 ),以去除功率谱中的频谱旁瓣(别名)作为频率和方位角阶 m 的函数。该算法使用两个嵌套循环,主要(外部)和次要(内部)循环。在次要迭代期间,我需要找到一维功率谱密度 (PSD) 数组(复值傅里叶变换数组的绝对平方)的最大值的索引,该索引在索引处的次要迭代之间部分更新. 指数ind_sel_maxpsd_check_tempfu_check_tempind_within_beamind_within_beam可能在较小的迭代之间在值和元素数量上有所不同,但是在len(psd_check_temp)给定足够多的迭代的情况下,与 相比,元素的最大数量会变小。示例性地,len(ind_check_temp)最大可以是~2500,并且最大len(psd_check_temp)可以增长到>~300000。len(psd_check_temp)仅在主要迭代之间更改,而在次要迭代之间不更改。在给定的迭代期间,我不知道 的值究竟psd_check_temp会如何变化(通常但不总是它们会减小),但它可能会在下一次迭代期间计算出来。但是,ind_sel_max始终包含在迭代中ind_with_beam并将psd_check_temp[ind_sel_max]在迭代期间更新。

循环结构:

i_major = 0
while continue_major_loop:
    i_minor = 0
    # code of outer loop
    ...
    while continue_minor_loop:
        # code of inner loop
        ...
        i_minor += 1
    # code of outer loop
    ...
    i_major += 1

代码示例 1(内部循环,慢):

        # compute the PSD
        if i_minor == 0:
            psd_check_temp = np.real(fu_check_temp)**2 + np.imag(fu_check_temp)**2
        else:
            psd_check_temp[ind_within_beam] = np.real(fu_check_temp[ind_within_beam])**2 + np.imag(fu_check_temp[ind_within_beam])**2
        # find the index of the maximum PSD
        ind_sel_max = np.argmax(psd_check_temp)
        ...
        ind_within_beam = ...
        fu_check_temp[ind_within_beam] -= ...
        ...

问题 1:在每次迭代中找到最大通孔的索引np.argmax()(参见上面的代码)将使用越来越多的计算时间,因为len(psd_check_temp)在每次主要迭代期间变得越来越大,因此在足够多次迭代后变得非常慢。

问题:由于在迭代过程中只有一小部分 PSD 数组被更新,直觉上我猜应该有一种方法可以使每次迭代的计算时间限制在更新的元素数量(它本身是有界的,如上面所写示例到 ~2500),而不是检查整个数组psd_check_temp的最大值。我怎样才能最好地做到这一点?

我的想法是psd_check_temp在第一次小迭代期间排序一次,并且在每次迭代中仅对更新的 PSD 值进行排序,保留未更新的 PSD 值,找到索引以将排序的更新 PSD 值插入到排序的未更新 PSD 值中np.searchsorted()并通过 插入排序后的更新 PSD 值np.insert()。由于psd_check_temp将保持排序,因此最大 PSD 的索引是最后一个。但是,我还需要相应地处理fu_check_temp几个 1D ( inu_check_temp, im_check_temp) 和 2D ( ind_2d) 索引数组。

代码示例 2(甚至更慢):

...
# create a 2D index array, filled with index -1 (marking invalid points)
ind_2d = np.full((nnu_use,nm), fill_value=-1, dtype=int)
...
while continue_major_loop:
    ...       
    # get a 1D array of indices corresponding to the 1D frequency and m index arrays
    ind_1d = np.arange(len(inu_check_temp))
    # fill the index array with the indices corresponding to the 1D frequency and m index arrays
    ind_2d[inu_check_temp, im_check_temp] = ind_1d
    ...
    while continue_minor_loop:
        if i_minor == 0:
            psd_check_temp = np.real(fu_check_temp)**2 + np.imag(fu_check_temp)**2
            # find the indices that sort the PSD array and sort the PSD and Fourier transform arrays as well
            # as the frequency and m index arrays accordingly
            ind_sort_temp = np.argsort(psd_check_temp)
            psd_check_temp = psd_check_temp[ind_sort_temp]
            fu_check_temp = fu_check_temp[ind_sort_temp]
            inu_check_temp = inu_check_temp[ind_sort_temp]
            im_check_temp = im_check_temp[ind_sort_temp]
            # also update the 2D index array accordingly
            ind_2d[inu_check_temp, im_check_temp] = ind_1d
        else:
            psd_updated_temp = np.real(fu_check_temp[ind_within_beam])**2 + np.imag(fu_check_temp[ind_within_beam])**2
            # (1) get the indices to sort the PSD of the updated Fourier transform
            ind_sort_temp = np.argsort(psd_updated_temp)
            # (2) sort the PSD, Fourier transform, and frequency and m index arrays accordingly
            psd_updated_temp = psd_updated_temp[ind_sort_temp]
            fu_updated_temp = fu_check_temp[ind_within_beam][ind_sort_temp]
            inu_updated_temp = inu_check_temp[ind_within_beam][ind_sort_temp]
            im_updated_temp = im_check_temp[ind_within_beam][ind_sort_temp]
            # (3) get the 1D indices of the array elements that have not been updated; np.setdiff1d() only returns
            # unique elements, but that is ok for index arrays where there are no duplicates; also it returns
            # the indices in order
            ind_complement_temp = np.setdiff1d(ind_1d, ind_1d[ind_within_beam], assume_unique=True)
            # (4) keep only the values that were not updated
            psd_cut_temp = psd_check_temp[ind_complement_temp]
            fu_cut_temp = fu_check_temp[ind_complement_temp]
            inu_cut_temp = inu_check_temp[ind_complement_temp]
            im_cut_temp = im_check_temp[ind_complement_temp]
            # (5) find the indices where to insert the updated values in the sorted arrays
            ind_insert_temp = np.searchsorted(psd_cut_temp, psd_updated_temp)
            # (6) insert the updated values in the sorted arrays
            psd_check_temp = np.insert(psd_cut_temp,ind_insert_temp,psd_updated_temp)
            fu_check_temp = np.insert(fu_cut_temp,ind_insert_temp,fu_updated_temp)
            inu_check_temp = np.insert(inu_cut_temp,ind_insert_temp,inu_updated_temp)
            im_check_temp = np.insert(im_cut_temp,ind_insert_temp,im_updated_temp)
            # (7) also update the 2D index array accordingly
            ind_2d[inu_check_temp, im_check_temp] = ind_1d
        # the index of the maximum corresponds to the last array element, since the array is sorted
        ind_sel_max = len(psd_check_temp) - 1
        ...
        ind_within_beam = ...
        fu_check_temp[ind_within_beam] -= ...
        ...

问题 2:此实现比代码示例 1 中的实现还要慢。对于 10000 次迭代,代码示例 1 的算法需要大约 0.2 秒,代码示例 2 需要 12 秒,我需要 10^6 次或更多次迭代。步骤np.insert()(6)、np.setdiff1d()(3)、ind_2d更新 (7)、保留未更新值 (4)、np.searchsorted()(5)、np.argsort(psd_check_temp)(1) 以及对更新值进行排序 (2) 的计算时间分别为 4.5、2.7、分别为 2.0、1.2、0.6、0.5 和 0.2 秒。此实现也多次请求内存。

非常感谢任何建议/帮助!

标签: pythonarraysnumpyindexingmax

解决方案


由于psd_check_temp在次要循环迭代期间数组似乎保持恒定大小,因此您只能在更新索引上计算其最大值:

        # compute the PSD
        if i_minor == 0:
            psd_check_temp = np.real(fu_check_temp)**2 + np.imag(fu_check_temp)**2
            ind_sel_max = np.argmax(psd_check_temp)
        else:
            psd_check_temp[ind_within_beam] = np.real(fu_check_temp[ind_within_beam])**2 + np.imag(fu_check_temp[ind_within_beam])**2
            ind_sel_max_candidate = ind_within_beam[np.argmax(psd_check_temp[ind_within_beam])]
            if psd_check_temp[ind_sel_max_candidate] > psd_check_temp[ind_sel_max]:
                ind_sel_max = ind_sel_max_candidate
        ...
        ind_within_beam = ...
        fu_check_temp[ind_within_beam] -= ...
        ...

推荐阅读