Repository navigation
Expand file tree
/
Copy pathgraph.py
More file actions
460 lines (391 loc) · 20.7 KB
/
Copy pathgraph.py
File metadata and controls
460 lines (391 loc) · 20.7 KB
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
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
import os
import multiprocessing
import warnings
from typing import List, Optional, Tuple, Union
import numpy as np
import deglib_cpp
import deglib_cpp.distances as cpp_distances
import deglib_cpp.search as cpp_search
import pathlib
from .distances import FloatSpace, Metric
from .cpu import InstructionSet
from .search import Filter
from .utils import assure_array, InvalidShapeException
class DynamicExplorationGraph:
"""
Public DynamicExplorationGraph Facade for End Users.
All public methods accept and return external_labels (User Object IDs).
Internal calls map transparently between external_label and internal_index.
This is the single, minimal graph class. Graph construction and mutation
must be performed via GraphBuilder.
"""
def __init__(self, graph_cpp: deglib_cpp.DynamicExplorationGraph):
self.dynamic_exploration_graph_cpp = graph_cpp
def size(self) -> int:
"""
:return: the number of vertices in the graph
"""
return self.dynamic_exploration_graph_cpp.size()
def get_edges_per_vertex(self) -> int:
"""
:return: the number of edges of each vertex
"""
return self.dynamic_exploration_graph_cpp.get_edges_per_vertex()
def get_po(self) -> int:
"""
:return: the traversal feature prefetch offset (candidates prefetched ahead)
"""
return self.dynamic_exploration_graph_cpp.get_po()
def get_pl(self) -> int:
"""
:return: the traversal feature prefetch look-ahead (cache lines per feature)
"""
return self.dynamic_exploration_graph_cpp.get_pl()
def get_nl(self) -> int:
"""
:return: the traversal neighbor prefetch look-ahead
"""
return self.dynamic_exploration_graph_cpp.get_nl()
def get_feature_space(self) -> FloatSpace:
"""
:return: the feature space
"""
return FloatSpace(float_space_cpp=self.dynamic_exploration_graph_cpp.get_feature_space())
def has_vertex(self, external_label: int) -> bool:
"""
:param external_label: The external label to check for existence.
:returns: whether the given external label is present in the graph.
"""
return self.dynamic_exploration_graph_cpp.has_vertex(external_label)
def is_mutable(self) -> bool:
"""
:return: whether the graph is mutable (can be modified via GraphBuilder)
"""
return self.dynamic_exploration_graph_cpp.is_mutable()
def get_neighbors(self, external_label: int) -> List[int]:
"""
Get the neighbor external labels of the given external label.
:param external_label: The external label to get the neighbors of
:returns: The external labels of the neighbors
"""
return self.dynamic_exploration_graph_cpp.get_neighbors(external_label)
def search(
self,
query: np.ndarray,
eps_or_ef: float = 0.0,
k: int = 10,
filter_labels: Union[None, np.ndarray, Filter] = None,
max_distance_computation_count: int = 0,
threads: int = 0,
return_distances: bool = True,
unsorted: bool = False,
) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray]:
"""
Search for nearest neighbors of query vector(s).
:param query: Query feature vector(s) as numpy array.
:param eps_or_ef: Exploration parameter. Values >= 1.0 are treated as fixed pool size (ef),
values < 1.0 are treated as relative distance margin (eps). Default: 0.0.
:param k: The number of results to return per query.
:param filter_labels: Filter for labels to include.
:param max_distance_computation_count: Distance computation budget limit.
:param threads: Number of parallel worker threads.
:param return_distances: If True, returns (indices, distances). If False, returns only indices.
:param unsorted: If True, returns candidates in unsorted order instead of ascending distance order.
:returns: (indices, distances) tuple if return_distances is True, otherwise indices array.
If fewer than k results are found, unfilled elements are padded with
np.iinfo(np.uint32).max (for indices) and np.nan / infinity (for distances).
"""
# handle query shapes
single_query = len(query.shape) == 1
if single_query:
query = query.reshape(1, -1)
if len(query.shape) != 2:
raise InvalidShapeException("invalid query shape: {}".format(query.shape))
if k > self.size():
warnings.warn(
"k={} is smaller than number of vertices in graph={}. Setting k={}".format(k, self.size(), self.size())
)
k = self.size()
valid_dtype = self.get_feature_space().metric().get_dtype()
query = assure_array(query, "query", valid_dtype)
filter_obj = Filter.create_filter(filter_labels, self.size())
threads = get_num_useful_threads(threads, query.shape[0])
indices_or_tuple = self.dynamic_exploration_graph_cpp.search_batch(
query, eps_or_ef, k, filter_obj, max_distance_computation_count, threads, return_distances, unsorted
)
if return_distances:
indices, distances = indices_or_tuple
if single_query:
return indices[0], distances[0]
return indices, distances
else:
indices = indices_or_tuple
if single_query:
return indices[0]
return indices
def explore(
self,
entry_external_label: Union[int, np.ndarray, list],
k: int,
max_distance_computation_count: int = 0,
eps_or_ef: float = 0.0,
include_entry: bool = True,
threads: int = 1,
filter_labels: Union[None, np.ndarray, Filter] = None,
return_distances: bool = True,
unsorted: bool = False,
) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray]:
"""
An exploration for similar elements, limited by max_distance_computation_count.
Handles both single entry labels (int) and batch entry arrays (ndarray/list).
:param entry_external_label: The external label of the vertex to start exploration from.
:param k: The number of similar feature vectors to return.
:param max_distance_computation_count: Limit the number of distance calculations.
:param eps_or_ef: Exploration parameter. Values >= 1.0 are treated as fixed pool size (ef),
values < 1.0 are treated as relative distance margin (eps). Default: 0.0.
:param include_entry: If True, the entry vertex is included in the result set.
:param threads: The number of threads to use for parallel processing.
:param filter_labels: Labels filter.
:param return_distances: If True, returns (indices, distances). If False, returns only indices.
:param unsorted: If True, returns candidates in order instead of ascending distance order.
:returns: (indices, distances) tuple if return_distances is True, otherwise indices array.
If fewer than k results are found, unfilled elements are padded with
np.iinfo(np.uint32).max (for indices) and np.nan / infinity (for distances).
"""
if k > self.size():
warnings.warn(
"k={} is larger than number of vertices in graph={}. Setting k={}".format(k, self.size(), self.size())
)
k = self.size()
if isinstance(entry_external_label, (np.ndarray, list, tuple)):
arr = np.ascontiguousarray(entry_external_label, dtype=np.uint32)
if arr.ndim == 2:
arr = arr[:, 0]
filter_obj = Filter.create_filter(filter_labels, self.size())
threads = get_num_useful_threads(threads, arr.shape[0])
return self.dynamic_exploration_graph_cpp.explore_batch(
arr,
k,
max_distance_computation_count,
eps_or_ef,
include_entry,
filter_obj,
threads,
return_distances,
unsorted,
)
else:
indices, distances = self.dynamic_exploration_graph_cpp.explore(
int(entry_external_label), k, max_distance_computation_count, eps_or_ef, include_entry
)
if return_distances:
return indices, distances
return indices
def has_path(
self,
entry_external_labels: Union[int, np.ndarray, list],
to_external_label: int,
eps: float = 0.0,
k: int = 1,
return_distances: bool = True,
) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray]:
"""
Greedy best-first traversal towards a stored vertex, returning the traversed path.
The target vertex's own feature vector is used as the query, so the traversal follows the
graph's greedy structure from the entry vertices until the target is reached. With
eps = 0.0 and k = 1 this is a pure greedy descent; larger eps widen the exploration radius
and allow escaping local optima at the cost of more distance computations.
:param entry_external_labels: External label(s) of the vertices to start the traversal from.
:param to_external_label: External label of the vertex the path should end at.
:param eps: Exploration factor widening the search radius. Default: 0.0 (greedy descent).
:param k: Size of the result list driving the exploration radius. Default: 1.
:param return_distances: If True, returns (path, distances). If False, returns only the path.
:returns: (path, distances) tuple if return_distances is True, otherwise the path array.
The path is ordered from the target back to the entry vertex. Both arrays are
empty if the target was not reached within the given budget.
"""
if isinstance(entry_external_labels, (int, np.integer)):
entries = [int(entry_external_labels)]
else:
entries = [int(label) for label in entry_external_labels]
return self.dynamic_exploration_graph_cpp.has_path(entries, int(to_external_label), float(eps), int(k), return_distances)
def save_graph(self, path: pathlib.Path | str):
"""
Save graph to specified file. Creates necessary directories.
:param path: The path where to save the file.
"""
self.dynamic_exploration_graph_cpp.save_graph(str(path))
def get_entry_vertex_indices(self) -> list[int]:
"""Returns the list of internal entry vertex indices."""
return self.dynamic_exploration_graph_cpp.get_entry_vertex_indices()
def set_entry_vertex_indices(self, indices: list[int] | np.ndarray):
"""Sets the list of internal entry vertex indices."""
self.dynamic_exploration_graph_cpp.set_entry_vertex_indices(list(indices))
def optimize(self, sample_count: int = 50, k: int = 100, ef: int = 200, seed: int = 7) -> None:
"""Auto-tune the graph's traversal prefetch parameters (po, pl, nl) by timing traversal over sampled vertices.
The graph self-samples its own stored features as queries, so no external query buffer is
needed. Call once after construction to tune traversal for the expected operating point.
:param sample_count: Number of vertices sampled as queries (capped at the graph size).
:param k: Result count per traversal (the expected query operating point).
:param ef: Beam width per traversal; clamped to at least k.
:param seed: Random seed for query sampling.
"""
self.dynamic_exploration_graph_cpp.optimize(int(sample_count), int(k), int(ef), int(seed))
def to_readonly(
self, feature_space: Optional[FloatSpace] = None, custom_features: Optional[np.ndarray] = None
) -> "DynamicExplorationGraph":
"""
Create a read-only graph from the given graph by only keeping information
that is useful for searching.
Optionally overrides feature space and replaces feature vectors with a custom features buffer
(e.g., swapping FP32 (d+1) features for FP16 d features while preserving graph topology).
:param feature_space: Optional FloatSpace specifying the target feature space.
:param custom_features: Optional NumPy array containing replacement feature vectors.
:return: A new read-only DynamicExplorationGraph.
"""
fs_cpp = feature_space.float_space_cpp if feature_space is not None else None
feat_cpp = np.ascontiguousarray(custom_features) if custom_features is not None else None
graph_cpp = deglib_cpp.read_only_graph_from_graph(self.dynamic_exploration_graph_cpp, fs_cpp, feat_cpp)
return DynamicExplorationGraph(graph_cpp)
def to_mutable(
self,
feature_space: Optional[FloatSpace] = None,
custom_features: Optional[np.ndarray] = None,
capacity: int = 0,
) -> "DynamicExplorationGraph":
"""
Create a mutable SizeBoundedGraph from the given graph (ReadOnly or SizeBounded),
optionally overriding the feature space and feature vectors.
All edge weights are recalculated using the new feature space.
:param feature_space: Optional FloatSpace specifying the target feature space.
If None, uses the graph's existing feature space.
:param custom_features: Optional NumPy array containing replacement feature vectors.
:param capacity: If > 0, sets the capacity to max(capacity, graph.size()).
If 0, capacity equals graph.size().
:return: A new mutable DynamicExplorationGraph.
"""
fs_cpp = feature_space.float_space_cpp if feature_space is not None else None
feat_cpp = np.ascontiguousarray(custom_features) if custom_features is not None else None
graph_cpp = deglib_cpp.size_bounded_graph_from_graph(
self.dynamic_exploration_graph_cpp, fs_cpp, feat_cpp, capacity
)
return DynamicExplorationGraph(graph_cpp)
def to_dynamic(
self,
feature_space: Optional[FloatSpace] = None,
custom_features: Optional[np.ndarray] = None,
chunk_size: int = 1024,
) -> "DynamicExplorationGraph":
"""
Create a mutable DynamicGraph with chunk-based dynamic memory allocation
from the given graph (ReadOnly, SizeBounded, or Dynamic),
optionally overriding the feature space and feature vectors.
All edge weights are recalculated using the new feature space.
:param feature_space: Optional FloatSpace specifying the target feature space.
If None, uses the graph's existing feature space.
:param custom_features: Optional NumPy array containing replacement feature vectors.
:param chunk_size: Target number of vertices per memory chunk (default = 1024).
Will be automatically rounded up to the nearest power of 2 (e.g. 600 -> 1024).
:return: A new mutable DynamicExplorationGraph.
"""
fs_cpp = feature_space.float_space_cpp if feature_space is not None else None
feat_cpp = np.ascontiguousarray(custom_features) if custom_features is not None else None
graph_cpp = deglib_cpp.dynamic_graph_from_graph(
self.dynamic_exploration_graph_cpp, fs_cpp, feat_cpp, chunk_size
)
return DynamicExplorationGraph(graph_cpp)
def __repr__(self) -> str:
return (
f"DynamicExplorationGraph(size={self.size()} edges_per_vertex={self.get_edges_per_vertex()} "
f"dim={self.get_feature_space().dim()})"
)
def get_num_useful_threads(requested: int, max_limit: int):
if requested == 0:
requested = multiprocessing.cpu_count()
return min(requested, max_limit) # dont use more threads than queries
__all__ = [
"DynamicExplorationGraph",
"create_empty",
"create_dynamic_empty",
"create_random_graph",
"load_readonly_graph",
"load_dynamic_graph",
"load_mutable_graph",
]
def create_empty(capacity: int, feature_space: FloatSpace, edges_per_vertex: int = 32) -> DynamicExplorationGraph:
"""
Create an empty mutable DynamicExplorationGraph (SizeBoundedGraph).
:param capacity: The maximal number of vertices of this graph.
:param feature_space: A FloatSpace object defining dimensionality, metric, and instruction set.
:param edges_per_vertex: Number of neighbors for each vertex. Defaults to 32.
:return: A new mutable DynamicExplorationGraph.
"""
graph_cpp = deglib_cpp.create_size_bounded_graph(capacity, edges_per_vertex, feature_space.float_space_cpp)
return DynamicExplorationGraph(graph_cpp)
def create_dynamic_empty(
feature_space: FloatSpace, edges_per_vertex: int = 32, chunk_size: int = 1024
) -> DynamicExplorationGraph:
"""
Create an empty mutable DynamicExplorationGraph with dynamic chunk-based memory allocation (DynamicGraph).
:param feature_space: A FloatSpace object defining dimensionality, metric, and instruction set.
:param edges_per_vertex: Number of neighbors for each vertex (must be even). Defaults to 32.
:param chunk_size: Target number of vertices per memory chunk (default = 1024).
Will be automatically rounded up to the nearest power of 2 (e.g. 600 -> 1024).
:return: A new mutable DynamicExplorationGraph.
"""
graph_cpp = deglib_cpp.create_dynamic_graph(edges_per_vertex, feature_space.float_space_cpp, chunk_size)
return DynamicExplorationGraph(graph_cpp)
def create_random_graph(
features: np.ndarray, feature_space: FloatSpace, edges_per_vertex: int = 32, seed: int = 7
) -> DynamicExplorationGraph:
"""
Create a random exploration graph from the given feature data.
The graph is built by first fully connecting the initial (edges_per_vertex + 1)
vertices, then iteratively inserting each remaining vertex by connecting it
to edges_per_vertex neighbors chosen from existing vertices.
:param features: A 2D NumPy array of feature vectors (shape: vertex_count x dim).
:param feature_space: A FloatSpace object defining dimensionality, metric, and instruction set.
:param edges_per_vertex: Number of neighbors for each vertex. Must be even. Defaults to 32.
:param seed: Random seed for deterministic graph construction. Defaults to 7.
:return: A new DynamicExplorationGraph with a random exploration graph.
"""
valid_dtype = feature_space.metric().get_dtype()
features = assure_array(features, "features", valid_dtype)
graph_cpp = deglib_cpp.create_random_graph(features, edges_per_vertex, feature_space.float_space_cpp, seed)
return DynamicExplorationGraph(graph_cpp)
def load_readonly_graph(path: pathlib.Path | str) -> DynamicExplorationGraph:
"""
Read a saved ReadOnlyGraph from given file.
:param path: The path where to look for the file
:raises FileNotFoundError: If the given file does not exist
:return: A read-only DynamicExplorationGraph.
"""
if not os.path.isfile(path):
raise FileNotFoundError('File "{}" could not be found'.format(path))
graph_cpp = deglib_cpp.load_readonly_graph(str(path))
return DynamicExplorationGraph(graph_cpp)
def load_dynamic_graph(path: pathlib.Path | str, chunk_size: int = 1024) -> DynamicExplorationGraph:
"""
Read a saved graph as a mutable DynamicGraph with chunk-based allocation.
:param path: The path where to look for the file
:param chunk_size: Target number of vertices per memory chunk (default = 1024).
Will be automatically rounded up to the nearest power of 2 (e.g. 600 -> 1024).
:raises FileNotFoundError: If the given file does not exist
:return: A mutable DynamicExplorationGraph.
"""
if not os.path.isfile(path):
raise FileNotFoundError('File "{}" could not be found'.format(path))
graph_cpp = deglib_cpp.load_dynamic_graph(str(path), chunk_size)
return DynamicExplorationGraph(graph_cpp)
def load_mutable_graph(path: pathlib.Path | str, capacity: int = 0) -> DynamicExplorationGraph:
"""
Read a saved graph as a mutable SizeBoundedGraph without recomputing edge weights.
:param path: The path where to look for the file
:param capacity: Optional capacity upper bound. If 0 or <= file size, capacity equals graph size in file.
:raises FileNotFoundError: If the given file does not exist
:return: A mutable DynamicExplorationGraph.
"""
if not os.path.isfile(path):
raise FileNotFoundError('File "{}" could not be found'.format(path))
graph_cpp = deglib_cpp.load_sizebounded_graph(str(path), capacity)
return DynamicExplorationGraph(graph_cpp)