python - 如何在python中加速N维区间树？

Question

考虑以下问题：给定一组n 个区间和一组m个浮点数，为每个浮点数确定包含该浮点数的区间子集。

这个问题已经通过构造一个区间树（或称为范围树或段树）来解决。已经针对一维情况进行了实现，例如 python 的intervaltree 包。通常，这些实现会考虑一个或几个浮点数，即上面的小“m”。

在我的问题设置中，n 和 m 都是非常大的数字（来自解决图像处理问题）。此外，我需要考虑 N 维区间（当 N=3 时称为长方体，因为我正在使用有限元方法对人脑进行建模）。我在python中实现了一个简单的N维区间树，但它循环运行，一次只能取一个浮点数。任何人都可以帮助提高效率方面的实施吗？您可以自由更改数据结构。

import sys
import time
import numpy as np

# find the index of a satisfying x > a in one dimension
def find_index_smaller(a, x):
    idx = np.argsort(a)
    ss = np.searchsorted(a, x, sorter=idx)
    res = idx[0:ss]
    return res

# find the index of a satisfying x < a in one dimension
def find_index_larger(a, x):
    return find_index_smaller(-a, -x)

# find the index of a satisfing amin < x < amax in one dimension
def find_intv_at(amin, amax, x):
    idx = find_index_smaller(amin, x)
    idx2 = find_index_larger(amax[idx], x)
    res = idx[idx2]
    return res

# find the index of a satisfying amin < x < amax in N dimensions
def find_intv_at_nd(amin, amax, x):
    dim = amin.shape[0]
    res = np.arange(amin.shape[-1])
    for i in range(dim):
        idx = find_intv_at(amin[i, res], amax[i, res], x[i])
        res = res[idx]
    return res

我还有两个用于完整性检查和性能测试的测试示例：

def demo1():
    print ("By default, we do a correctness test")
    n_intv = 2
    n_point = 2
    # generate the test data
    point = np.random.rand(3, n_point)
    intv_min = np.random.rand(3, n_intv)
    intv_max = intv_min + np.random.rand(3, n_intv)*8
    print ("point ")
    print (point)
    print ("intv_min")
    print (intv_min)
    print ("intv_max")
    print (intv_max)
    print ("===Indexes of intervals that contain the point===")
    for i in range(n_point):
        print (find_intv_at_nd(intv_min,intv_max, point[:, i]))

def demo2():
    print ("Performance:")
    n_points=100
    n_intv = 1000000

    # generate the test data
    points = np.random.rand(n_points, 3)*512
    intv_min = np.random.rand(3, n_intv)*512
    intv_max = intv_min + np.random.rand(3, n_intv)*8
    print ("point.shape = "+str(points.shape))
    print ("intv_min.shape = "+str(intv_min.shape))
    print ("intv_max.shape = "+str(intv_max.shape))

    starttime = time.time()
    for point in points:
        tmp = find_intv_at_nd(intv_min, intv_max, point)
    print("it took this long to run {} points, with {} interva: {}".format(n_points, n_intv, time.time()-starttime))

我的想法是：

1. 从算法中删除 np.argsort()，因为区间树没有改变，所以排序可以在预处理中完成。
2. 向量化 x。算法为每个 x 运行一个循环。如果我们能摆脱 x 上的循环就好了。

任何贡献将不胜感激。

标签： pythonperformancenumpyvectorizationfinite-element-analysis