FNNTW：最快的最近邻（在）西部

Fastest Neighbor Nearest Neighbor (in the) West 是一个正在开发的 kD-tree 库，旨在成为现有的性能最好的并行 kNN 库之一。

基本用法

import pyfnntw;
import numpy as np;
from scipy.spatial import cKDTree as Tree
from time import time

ND = 10**5
NQ = 10**6
RUNS = 10
TRIALS = 10
WARMUP = 5

# Get some data and queries
data = np.random.uniform(size=(ND, 3))
query = np.random.uniform(size=(NQ, 3))

# Build and query the pynntw tree
tree1 = pyfnntw.Tree(data, 32, 1)
(_, ids1) = tree1.query(query)

# Build and query the scipy tree
tree2 = pyfnntw.Tree(data, 32, 1)
(_, ids2) = tree2.query(query)

if np.all(ids1 == ids2):
    print("Success")
else:
    print("Failure")

kD-trees 的构建和查询过程有几个关键组件，可以提高其性能。

1. 建造树

一个。使用快速选择

select_nth_unstable_by通过使用in Rust形式的快速选择core::slice而不是平均值甚至其他中值算法，我们可以在相同的 O(N) 算法中免费获得left和right子集，而不必再进行另一个 O(N) 比较来获得和left垃圾箱right。因此，构建仍然是 O(N log(N))，但在每次拆分期间删除了在许多其他库中发现的整个 O(N) 操作。

湾。并行构建

许多 kD-Tree 构建实现不是并行的。kD-tree 的每个子树都是一个独立的 kD-tree。因此，经过一些拆分后，可以将剩余子树的构建扔到不同的线程上。这是在用户指定的这里完成的par_split_level，这是并行开始的树深度。有关推荐值，请参阅针对其他代码的基准测试中的注释。

2.查询树

大多数库返回 kNN 距离或距离和相应数据点的索引。这里我们返回距离、索引和邻居。我们获取索引的方式是新的。

一个。HashMap<&'[NotNan<f64>; D], u64>

我们试图做其他库所做的事情，即在某些结构或元组中携带数据的索引，例如&(usize, [NotNan<f64>; D])，但发现性能很差。通过使用ordered_floatcrate 的NotNan<T>struct，浮点数变得可散列。然后，可以只使用查询和候选位置来执行所有查询操作，然后在最后执行 O(1) 查找以检索最近邻居的索引。这HashMap是在构建时与树本身并行构建的，只增加了线程生成的最小开销。

3. 不安全的访问（也许......可能是次要的，没有基准测试）

（我争论过我是否应该提到这个......）因为我们在编译时知道所有数组的形状（即树的维度），unsafe方法get_unchecked和get_unchecked_mut在整个代码中自由使用。尽管尚未明确测量其影响，但理论上这意味着实际上没有进行边界检查。无论如何，LLVM 可能已经优化了许多边界检查。也有算法保证非空的unwrap_unchecked调用。Option<T>（与最佳实践相反，这两个用例是预先优化的）。

守则现状

目前，该库仅提供 1NN 查询。在不久的将来，此功能将扩展到 kNN (k>1) 查询。此外，还将添加对周期性边界条件的支持。最后，计划添加泛型参数以允许单态化f32和f64（目前仅f64支持）。

针对其他代码的基准测试

该库旨在供作者用于计算宇宙学中的汇总统计数据。因此，所选择的基准参数与用于分析宇宙学模拟输出的参数接近。在这样的应用程序中，通常会使用许多子样本或模拟框。因此，组合的构建 + 查询时间很重要，因为可以在分析中构建许多不同的树。我们用

• 单位立方体中 100,000 个均匀随机点的模拟数据集。
• 单位立方体中 1,000,000 个均匀随机点的查询集。

数据集和查询点的 100 多个实现，在使用 48 个线程的 AMD EPYC 7502P 上获得了平均构建（串行）和 1NN 查询时间的以下结果。结果按组合的构建和查询时间排序。

借助 FNNTW 的并行构建功能，AMD EPYC 7502P 的构建时间可低至 8.7 毫秒（at split_level = 1）。由于并行性和原子操作的开销会在数据点数量较少时减慢构建速度，因此并行构建和非并行构建都可以通过Tree:new(..)和获得Tree::new_parallel(..)。后者采用上述参数par_split_level，即并行开始的分割深度。尽管对于 O(1e5) 点的应用，我们看到了最大的改进par_split_level = 1，但我们怀疑最优值par_split_level会随着数据集的大小而增加。