1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
module type S = sig
type vec
(** Ball Tree structure. *)
type t
(** [t.%(idx)]
Indexed access to vectors/points used to build the ball tree [t]. *)
val ( .%() ) : t -> int -> vec
(** [points t]
Return a list of the vectors/points used to build the ball tree [t]. *)
val points : t -> vec list
(** [points' t]
Return a copied array of the vectors/points used to build the ball tree
[t]. *)
val points' : t -> vec array
(** [make ?leaf_size pts]
Build a ball tree from the list of vectors [pts]. Recursive construction
of the tree will cease branching when a node holds a number of points less
than or equal to [leaf_size]. *)
val make : ?leaf_size:int -> vec list -> t
(** [make' ?leaf_size pts]
Build a ball tree from the array of vectors [pts]. Recursive construction
of the tree will cease branching when a node holds a number of points less
than or equal to [leaf_size]. *)
val make' : ?leaf_size:int -> vec array -> t
(** [search_idxs ?radius t p]
Search through the ball tree [t] for points that are within a distance
[radius] of the target point [p]. Matches are returned as an arbitrarily
ordered list of indices (into the point array used to construct [t]). *)
val search_idxs : ?radius:float -> t -> vec -> int list
(** [search_points ?radius t p]
Search through the ball tree [t] for points that are within a distance
[radius] of the target point [p]. Matches are returned as an arbitrarily
ordered (not sorted from nearest to furthest) list of points. *)
val search_points : ?radius:float -> t -> vec -> vec list
end
module type Projection = sig
type vec
val proj : vec array -> int array -> float array
end
module Proj2 = struct
let proj points idxs =
let mn = ref points.(idxs.(0))
and mx = ref points.(idxs.(0))
and len = Array.length idxs in
for i = 1 to len - 1 do
let p = points.(idxs.(i)) in
(mn := V2.(v (min (x !mn) (x p)) (min (y !mn) (y p))));
mx := V2.(v (max (x !mx) (x p)) (max (y !mx) (y p)))
done;
let d = V2.sub !mx !mn in
let project = if V2.(x d >= y d) then V2.x else V2.y in
Array.map (fun idx -> project points.(idx)) idxs
end
module Proj3 = struct
let proj points idxs =
let mn = ref points.(idxs.(0))
and mx = ref points.(idxs.(0))
and len = Array.length idxs in
for i = 1 to len - 1 do
let p = points.(idxs.(i)) in
(mn := V3.(v (min (x !mn) (x p)) (min (y !mn) (y p)) (min (z !mn) (z p))));
mx := V3.(v (max (x !mx) (x p)) (max (y !mx) (y p)) (max (z !mx) (z p)))
done;
let d = V3.sub !mx !mn in
let dx = V3.x d
and dy = V3.y d
and dz = V3.z d in
let project =
match Float.(compare dx dy, compare dx dz, compare dy dz) with
| 1, 1, _ -> V3.x
| -1, _, 1 -> V3.y
| _, -1, -1 -> V3.z
| _ -> V3.x
in
Array.map (fun idx -> project points.(idx)) idxs
end
module Make (V : V.S) (P : Projection with type vec := V.t) : S with type vec := V.t =
struct
type tree =
| Leaf of int array
| Node of
{ pivot : int
; radius : float
; left : tree
; right : tree
}
type t =
{ points : V.t array
; tree : tree
}
let ( .%() ) t i = t.points.(i)
let points t = List.init (Array.length t.points) (fun i -> t.points.(i))
let points' t = Array.copy t.points
let make'' ?(leaf_size = 25) points =
let rec aux idxs =
let len = Array.length idxs in
if len <= leaf_size
then Leaf idxs
else (
let projected = P.proj points idxs in
let mean_proj =
Array.fold_left (fun sum p -> sum +. p) 0. projected /. Float.of_int len
in
let local_pivot =
let idx = ref 0
and min = ref Float.max_float in
for i = 0 to len - 1 do
let d = Float.abs (projected.(i) -. mean_proj) in
if d < !min
then (
min := d;
idx := i )
done;
!idx
in
let pivot = idxs.(local_pivot) in
let radius =
let p = points.(pivot) in
let f max idx = Float.max max V.(distance p points.(idx)) in
Array.fold_left f 0. idxs
in
let left, right =
let l = ref []
and r = ref [] in
for i = 0 to len - 1 do
if i <> local_pivot
then (
if projected.(i) <= mean_proj then l := idxs.(i) :: !l;
if projected.(i) > mean_proj then r := idxs.(i) :: !r )
done;
aux (Util.array_of_list_rev !l), aux (Util.array_of_list_rev !r)
in
Node { pivot; radius; left; right } )
in
{ points; tree = aux (Array.init (Array.length points) Fun.id) }
let make' ?leaf_size points = make'' ?leaf_size (Array.copy points)
let make ?leaf_size points = make'' ?leaf_size (Array.of_list points)
let search_idxs ?(radius = Util.epsilon) { points; tree } target =
let rec aux = function
| Leaf idxs ->
let f acc i = if V.approx ~eps:radius points.(i) target then i :: acc else acc in
Array.fold_left f [] idxs
| Node tree ->
if not @@ V.approx ~eps:(tree.radius +. radius) points.(tree.pivot) target
then []
else (
let children = List.rev_append (aux tree.left) (aux tree.right) in
if V.approx ~eps:radius points.(tree.pivot) target
then tree.pivot :: children
else children )
in
aux tree
let search_points ?radius t target =
List.map (fun i -> t.points.(i)) (search_idxs ?radius t target)
end