PSI provides three parallel spatial partition trees:
| Header | Paper name | What it is |
|---|---|---|
psi/kd_tree.h |
Pkd-tree | parallel binary $k$d-tree |
psi/orth_tree.h |
P-Orth tree | parallel quad/oct-tree, any dimension |
psi/p_tree.h |
SPaC-tree | parallel 1-D tree over a Morton or Hilbert code |
There is also an AI-generated wiki covering how things work internally. Contributions welcome 🤗.
Every code block below is compiled by tests/unit/test_manual.cpp, so if the
API moves and this file is not updated, the build breaks.
PSI is header-only. It needs C++20, parlaylib and libmorton.
# Already have the repository, e.g. as a submodule:
add_subdirectory(SpaceTreeLib)
target_link_libraries(myapp PRIVATE PSI::PSI)# Or install it once (cmake --install) and find it:
find_package(PSI REQUIRED)
target_link_libraries(myapp PRIVATE PSI::PSI)Directly, without CMake:
g++ -std=c++20 -O3 -pthread myapp.cpp \
-I SpaceTreeLib/include \
-I SpaceTreeLib/include/parlaylib/include \
-I SpaceTreeLib/include/libmorton/includeinclude/ is the only PSI include root, and every header is reached as
psi/…. Build your own targets with optimisation: PSI's -O3 applies to
PSI's own targets, not to yours, and an unoptimised spatial index looks broken
rather than slow. Add -march=native if you are measuring.
Build options worth knowing:
| Option | Default | Meaning |
|---|---|---|
PSI_NATIVE_ARCH |
ON |
-march=native -mcx16. Turn off for a binary that runs on another machine. |
PSI_ASSERTS |
OFF |
keep assert() in an optimised build; the configuration to develop against |
PSI_BUILD_TESTS |
on if top-level | the ctest suite |
PSI_BUILD_BENCHMARKS |
on if top-level | the drivers and baselines; also what pulls in Boost |
PSI_BUILD_EXAMPLES |
on if top-level | the three example programs |
DEBUG |
OFF |
unoptimised build with asserts and validate() |
CGAL |
OFF |
adds the kd_ccp / p_ccp correctness oracles |
Taking the Pkd-tree; the other two differ only where noted under Choosing and configuring a tree. The complete programs are in example/.
#include "psi/kd_tree.h"
#include "psi/dependence/splitter.h"
using coord_type = long;
using point = psi::basic_point<coord_type, 2>;coord_type may be an integer or floating-point type. The dimension is a
template argument, so it is fixed at compile time.
A split rule is a dimension rule paired with a partition rule:
using split_rule = psi::orthogonal_split_rule<psi::rotate_dim<point>,
psi::spatial_median<point>>;| Dimension rule | Picks the split dimension by |
|---|---|
psi::rotate_dim<Point> |
cycling through the dimensions by depth |
psi::max_stretch_dim<Point> |
the widest side of the bounding box |
| Partition rule | Splits at |
|---|---|
psi::object_median<Point> |
the median point, so the two sides hold equal counts |
psi::spatial_median<Point> |
the middle of the box, so the two sides are equal in space |
A traits type carries everything a tree is configured with: the point type, the split rule, the augmentations, the skeleton height and the imbalance ratio. Name it once and hand it to any tree that suits those points.
using traits = psi::tree_traits<point, split_rule>;Only the first two have no default. psi/tree_traits.h lists the rest.
psi::kd_tree<traits> tree;
parlay::sequence<point> points(1000);
/* ... fill points ... */
tree.build(parlay::make_slice(points));build copies what it is given, so the input may go out of scope afterwards.
Calling build again on a live tree replaces its contents.
Every point inside a box, and the k nearest neighbours of a point:
typename decltype(tree)::box_type box;
box.first[0] = 400; box.first[1] = 400; /* lower corner */
box.second[0] = 600; box.second[1] = 600; /* upper corner */
parlay::sequence<point> inside = tree.range_query(box);
auto nearest = tree.knn(query_point, 10); /* sorted, nearest first */knn gives back (point, squared distance) pairs by value, so they stay
valid after the tree changes. Both return fewer results than asked for when
the tree holds fewer points, and both return nothing on an empty tree.
That is all most callers need. The forms below let you own the storage instead, which is what the benchmarks use:
auto [count, count_log] = tree.range_count(box);
parlay::sequence<point> found(tree.get_size());
auto [written, query_log] = tree.range_query(box, parlay::make_slice(found));
/* found[0 .. written) are the points inside box */range_query writes into the buffer you supply and returns how many points it
wrote. A buffer smaller than the answer is not an error: the result is
truncated and written tells you so. The second value in each pair is a
counter block for profiling; ignore it unless you are measuring.
using dis_type = typename point::dis_type;
using nn_pair = std::pair<std::reference_wrapper<point>, dis_type>;
size_t const k = 10;
parlay::sequence<nn_pair> result(k, nn_pair(std::ref(points[0]), 0));
psi::bounded_queue<point, nn_pair> queue(parlay::make_slice(result));
tree.knn(query_point, queue);
/* result[0 .. queue.size()) hold the neighbours, unsorted, with squared
* distances. queue.size() is less than k when the tree holds fewer points. */The queue must be pre-filled, because std::reference_wrapper has no default
constructor; any point will do as the filler. Unlike knn(q, k), these
results reference storage owned by the tree, so do not use them after
modifying it.
tree.batch_insert(parlay::make_slice(more_points));
tree.batch_delete(parlay::make_slice(points_to_remove)); /* must be present */
tree.batch_diff(parlay::make_slice(maybe_present)); /* tolerates absent */batch_delete requires every point to be in the tree; batch_diff is the one
to use when that may not hold. The tree rebuilds imbalanced subtrees on its
own as updates accumulate.
parlay::sequence<point> all(tree.get_size());
size_t n = tree.flatten(parlay::make_slice(all));flatten needs a buffer of exactly get_size(); given any other size it
writes nothing and returns 0.
The destructor frees the tree. delete_tree() does it early and may be
called more than once. A tree is move-only: it owns its nodes, so copying it
would free them twice.
Which tree do I want? compares them. Concretely:
/* Pkd-tree: any dimension rule, any partition rule */
psi::kd_tree<psi::tree_traits<point, split_rule>> kd;
/* P-Orth tree: partition rule must be spatial_median, and the skeleton
* height must be a multiple of the dimension. Both are static_asserts;
* orth_tree_traits picks a height that fits. The arity follows the point. */
psi::orth_tree<psi::orth_tree_traits<point, split_rule>> orth;
/* SPaC-tree: the point must carry the curve code, so it is an aug_point */
struct curve_code {
using id_type = int_fast32_t;
using curve_code_type = uint64_t;
curve_code() : code(0), id(0) {}
void set_member(curve_code_type const &v) { code = v; }
bool operator<(curve_code const &r) const
{
return code == r.code ? id < r.id : code < r.code;
}
bool operator==(curve_code const &r) const { return id == r.id; }
friend std::ostream &operator<<(std::ostream &o, curve_code const &r)
{
return o << r.code << " " << r.id;
}
curve_code_type code;
id_type id;
};
using aug_pt = psi::aug_point<coord_type, 2, curve_code>;
using curve = psi::spatial_filling_curve<psi::morton_curve<aug_pt>>;
psi::p_tree<psi::tree_traits<aug_pt, curve>> p;psi::hilbert_curve<Point> is the other curve. For p_tree, two points are
the same point when their ids match, so give every point a distinct one.
Each node carries an augmentation. By default it is a bounding box, and you need to write nothing:
psi::kd_tree<psi::tree_traits<point, split_rule>> tree;
/* the traits default to box_leaf_aug and box_interior_aug */Supply your own only if a bounding box is not what you want to store. The
contracts are the two concepts in psi/dependence/concepts.h:
template <typename A, typename Slice>
concept leaf_augmentation = requires(A a, Slice in) {
A(); A(in); a.update_aug(in); a.reset();
};
template <typename A>
concept interior_augmentation = requires(A a, bool flag, size_t n) {
a.set_parallel_flag(flag);
a.reset_parallel_flag();
{ a.get_parallel_flag_ini_status() } -> std::convertible_to<bool>;
{ a.force_parallel(n) } -> std::convertible_to<bool>;
};Two things the concepts cannot state. First, an interior augmentation also
needs create and update member templates, because they take the very
interior_type that holds the augmentation, which is not nameable where the
concept is checked — copy their shape from psi/dependence/augmentation.h.
Second, the library looks for a get_box() member to decide whether a
node knows its own bounds; an augmentation that stores a box under a different
name silently falls back to the slower hyperplane paths.
psi::box_leaf_aug and psi::box_interior_aug in
psi/dependence/augmentation.h are the worked example.
build, the batch updates and flatten are internally parallel. A single
query is not: knn, range_count and range_query each run on one thread,
so a loop over query points uses one core. Parallelise the batch yourself:
parlay::parallel_for(0, queries.size(), [&](size_t i) {
/* each iteration does its own query */
});The thread count comes from parlaylib: set PARLAY_NUM_THREADS, or build with
-DSERIAL=ON for a single-threaded library. Queries do not mutate the tree, so
running many at once is safe; updates are not concurrent with anything.
Common to all three trees. Range is any random-access range of points.
| Call | Returns | Notes |
|---|---|---|
build(Range&&) |
— | copies the input; replaces any existing contents |
batch_insert(Range&&) |
— | |
batch_delete(Range&&) |
— | every point must be present |
batch_diff(Range&&) |
— | tolerates points that are absent |
flatten(Range&& out) |
size_t |
needs out.size() == get_size(), else writes nothing and returns 0 |
knn(Point const&, size_t k) |
sequence<pair<Point, dis_type>> |
allocates; sorted, nearest first; points by value |
range_query(box_type const&) |
points_type |
allocates and returns what it found |
knn(Point const&, bounded_queue&) |
profiling counters | results land in the queue; queue.size() is how many |
range_count(box_type const&) |
pair<size_t, counters> |
|
range_query(box_type const&, Range&& out) |
pair<size_t, counters> |
truncates to out.size() |
get_size(), size() |
size_t |
number of points; the two are the same call |
empty() |
bool |
|
bounds() |
box_type const& |
a bound on the points -- see below |
count(Point const&) |
size_t |
how many points share those coordinates |
contains(Point const&) |
bool |
count(p) != 0 |
insert(Point const&) |
— | one point; see the warning below |
erase(Point const&) |
size_t |
one point, returns how many went away |
delete_tree(), clear() |
— | idempotent; the two are the same call |
get_tree_name() |
char const* |
parsed by the plotting scripts; do not change |
bounds() is a bound, not always the tightest one: the batch deletes leave it
alone, so it can outlive the points that stretched it. Queries stay correct
either way, and a rebuild tightens it. p_tree is the exception -- its box is
the cpam augmentation at the root, which every update keeps current.
count and contains compare coordinates, not identity. p_tree keys its
points on a curve code and an id, so two points at the same position with
different ids both count.
insert and erase exist because sometimes one point is genuinely all you
have. A loop over them is the wrong shape: these trees are built to absorb
a whole batch in parallel, and one-at-a-time gives up everything that makes
them fast. Collect the points and call batch_insert or batch_diff once.
They also inherit the preconditions of the batch calls they wrap -- notably
orth_tree does not grow its bounding box, so inserting outside bounds()
is an error there.
bgi::rtree |
PSI |
|---|---|
rtree(first, last) |
build(range) |
insert(v) / insert(first, last) |
insert(p) / batch_insert(range) |
remove(v) / remove(first, last) |
erase(p) / batch_diff(range) |
query(bgi::within(box), out) |
range_query(box) or range_query(box, out) |
query(bgi::nearest(p, k), out) |
knn(p, k) or knn(p, queue) |
size(), empty(), clear(), bounds() |
the same names |
count(v) |
count(p), over coordinates |
begin() / end() |
flatten(out) -- there is no iterator |
PSI has no predicate composition (bgi::satisfies, !, &&) and no lazy
query iterators (qbegin/qend). A query answers into a buffer or a returned
sequence, and filtering is the caller's to do afterwards.
orth_tree::build takes an optional bounding box, and set_bounding_box pins
one before building — useful when later inserts fall outside the initial
extent, since an orth tree does not grow its box. Inserting outside it is an
error.
Member types you will use: point_type, box_type, points_type,
slice_type, coord_type, dis_type. All of them come from the traits, so
traits::box_type and tree::box_type are the same type -- name whichever
reads better where you are.
Three extension points, in increasing order of effort.
The cheapest. Write a type with encode() in the shape of
psi::morton_curve / psi::hilbert_curve (psi/dependence/splitter.h) and
pass psi::spatial_filling_curve<your_curve> as p_tree's split rule. The
code is applied during the build by cpam_sample_sort; p_tree itself never
calls the rule.
Derive from psi::base_split_dim_rule or psi::base_split_partition_rule and
implement every pure virtual — including allow_rebuild(), which batch update
asks before rebuilding a subtree.
Then add a branch for it in orthogonal_split_rule::handle_undivided, which
decides what happens when the chosen dimension cannot be split. Without one
you get a static_assert naming the rule; that assert is the whole reason it
is worth mentioning here.
Not "derive and override". A derived tree has to supply, by name:
| What | Why |
|---|---|
leaf_type, interior_type |
every algorithm in base_tree_impl/ takes them as template arguments |
splitter_type |
the build; the split rule itself comes from the traits |
node_arr_type, node_regions, splitter_num |
multi-way trees only; the skeleton builder reads them off the derived class |
a tag method plus a matching is_* concept |
how the shared code tells the families apart |
build_recursive / serial_build_recursive |
the split rule calls back into these, so the signature is fixed by the family |
friend split_rule_type and the two rebuild_* friends |
the callbacks reach private members |
delete_tree() |
pure virtual on the base |
kd_tree is the smaller model to copy; orth_tree shows the multi-way side.
Start by copying one wholesale and deleting, rather than deriving from scratch.
Two things the base cannot check for you: validate() has no branch for a node
type that is neither binary nor multi-way (it says so and continues), and the
augmentation concepts cannot state the create/update requirement.
cmake -S . -B build -DPSI_ASSERTS=ON && cmake --build build -j
ctest --test-dir build # the unit suite, well under a second
./example/run_examples.sh # the examples check their own answersFor the CGAL oracle, configure with -DCGAL=ON and run
script/checkCorrect.sh fast <data-root> build. Generate a data root with the
tool below.
# In the build directory
./data_generator -p <out-dir> -n <points> -d <dim> -file_num 2 -varden 0 -seed 1Files land in <out-dir>/uniform/<n>_<d>/1.in — or uniform_bigint when
-axis_max is above 10^6, which is the default, and ss_varden… with
-varden 1. Only dimensions 2, 3, 5, 7, 9, 12 and 16 are supported; anything else
throws. -seed makes the output reproducible; without it a fixed default
is used and printed.
include/
├── baselines/ competitor implementations, benchmark only
│ ├── cpam_raw/ pristine CPAM and PAM
│ ├── zdtree/
│ └── zdtree_3d/
├── libmorton/ submodule, Morton codes
├── parlaylib/ submodule, parallel primitives
└── psi/
├── base_tree.h shared machinery
├── base_tree_impl/
├── tree_traits.h the configuration a tree takes, and the box arithmetic
├── tree_traits_impl/
├── dependence/ points, splitters, concepts, augmentations, cpam
├── kd_tree.h + kd_tree_impl/
├── orth_tree.h + orth_tree_impl/
└── p_tree.h + p_tree_impl/
example/ three worked programs, all of which check themselves
tests/unit/ the correctness suite
tests/ benchmark drivers
script/, script_ae/ experiment drivers for the two papers
psi/dependence/cpam/ is a fork of CPAM that p_tree is built on; see
THIRD_PARTY.md.