The code below is supplied with your installation (see Python examples).
Approximate Nearest Neighbors Example
1# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
2#
3# Redistribution and use in source and binary forms, with or without modification,
4# are permitted provided that the following conditions are met:
5# 1. Redistributions of source code must retain the above copyright notice,
6# this list of conditions and the following disclaimer.
7# 2. Redistributions in binary form must reproduce the above copyright notice,
8# this list of conditions and the following disclaimer in the documentation
9# and/or other materials provided with the distribution.
10# 3. Neither the name of the copyright holder nor the names of its contributors
11# may be used to endorse or promote products derived from this software without
12# specific prior written permission.
13#
14# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
15# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
16# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
17# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
18# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
19# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
20# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
21# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
22# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
23# POSSIBILITY OF SUCH DAMAGE.
24#
25
26"""
27Approximate nearest neighbors example Python script
28"""
29
30import sys
31import numpy as np
32from aoclda.neighbors import approximate_neighbors
33
34
35def approximate_neighbors_example():
36 """
37 Approximate Nearest Neighbors Example using IVFFlat algorithm
38 """
39
40 # Define training data (16 samples, 2 features)
41 # Data forms 4 clusters
42 x_train = np.array([[0.0, -0.1],
43 [1.1, 0.0],
44 [0.0, 1.1],
45 [1.0, 1.0],
46 [6.0, 0.0],
47 [7.2, 0.1],
48 [6.1, 1.0],
49 [7.0, 1.1],
50 [0.0, 10.0],
51 [1.0, 10.2],
52 [0.1, 11.0],
53 [1.1, 11.1],
54 [10.0, 10.0],
55 [11.1, 10.0],
56 [10.0, 11.2],
57 [11.0, 11.0]], dtype=np.float64)
58
59 # Define query points
60 x_test = np.array([[3.5, 0.4],
61 [0.4, 5.0],
62 [5.6, 5.1]], dtype=np.float64)
63
64 print("\nApproximate Nearest Neighbors for a small data matrix\n")
65
66 # n_list=4: create 4 clusters (centroids)
67 # n_probe=2: search 2 nearest clusters at query time
68 # seed=123: for reproducibility
69 ann = approximate_neighbors(n_neighbors=3, n_list=4, n_probe=1,
70 metric='sqeuclidean', kmeans_iter=10, seed=123)
71
72 # Train and add data to the index in one step
73 ann.train_and_add(x_train)
74 print("\nIndex trained and populated successfully.")
75
76 # Query the k-nearest neighbors
77 k_dist, k_ind = ann.kneighbors(x_test, return_distance=True)
78
79 # Print results
80 print("\nIndices of the approximate nearest neighbors:\n")
81 print(k_ind)
82 print("\nCorresponding squared Euclidean distances:\n")
83 print(k_dist)
84
85 # Print cluster information
86 print("\nCluster centroids:\n")
87 print(ann.cluster_centroids)
88 print("\nNumber of vectors per cluster:")
89 print(ann.list_sizes)
90 print("\nTotal vectors in index:", ann.n_index)
91
92 # Demonstrate the effect of n_probe on search accuracy
93 # Lower n_probe = faster but potentially less accurate
94 # Higher n_probe = slower but more accurate
95 print("\n--- Comparing n_probe values ---\n")
96
97 # Search with n_probe = 1 (set in constructor)
98 k_dist_approx, k_ind_approx = ann.kneighbors(x_test, return_distance=True)
99 print("n_probe = 1 (approximate):")
100 print(" Indices:")
101 print(k_ind_approx)
102 print(" Distances:")
103 print(k_dist_approx)
104
105 # Search with n_probe = n_list (exact search over all clusters)
106 ann.n_probe = ann.n_list
107 k_dist_exact, k_ind_exact = ann.kneighbors(x_test, return_distance=True)
108 print(f"\nn_probe = {ann.n_list} (exact, searches all clusters):")
109 print(" Indices:")
110 print(k_ind_exact)
111 print(" Distances:")
112 print(k_dist_exact)
113
114 # Demonstrate the train then add workflow (alternative to train_and_add)
115 print("\n--- Demonstrating train/add workflow ---\n")
116
117 ann2 = approximate_neighbors(n_neighbors=3, n_list=4, n_probe=1,
118 metric='sqeuclidean', seed=123)
119
120 # Train only (computes centroids, does not add data)
121 ann2.train(x_train)
122 print("Index trained (centroids computed).")
123
124 # Add data to the index
125 ann2.add(x_train)
126 print("Data added to the index.")
127
128 # Query again
129 k_dist2, k_ind2 = ann2.kneighbors(x_test, return_distance=True)
130
131 # Verify results match the train_and_add workflow
132 if np.array_equal(k_ind2, k_ind_approx):
133 print("\nResults match the train_and_add workflow.")
134 else:
135 print("\nWarning: Results differ from train_and_add workflow.")
136 sys.exit(1)
137
138 print("\nApproximate nearest neighbors successfully computed\n")
139
140
141if __name__ == "__main__":
142 approximate_neighbors_example()
The code below can be found in ann.cpp in the examples folder of your installation.
Approximate Nearest Neighbors Example
1/*
2 * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
3 *
4 * Redistribution and use in source and binary forms, with or without modification,
5 * are permitted provided that the following conditions are met:
6 * 1. Redistributions of source code must retain the above copyright notice,
7 * this list of conditions and the following disclaimer.
8 * 2. Redistributions in binary form must reproduce the above copyright notice,
9 * this list of conditions and the following disclaimer in the documentation
10 * and/or other materials provided with the distribution.
11 * 3. Neither the name of the copyright holder nor the names of its contributors
12 * may be used to endorse or promote products derived from this software without
13 * specific prior written permission.
14 *
15 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
16 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
17 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
18 * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
19 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
20 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
21 * OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
22 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
23 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
24 * POSSIBILITY OF SUCH DAMAGE.
25 *
26 */
27
28#include "aoclda.h"
29#include <iomanip>
30#include <iostream>
31#include <vector>
32
33/*
34 * Approximate Nearest Neighbors example
35 *
36 * This example demonstrates the approximate nearest neighbors workflow
37 * for efficient similarity search.
38 */
39
40int main() {
41 std::cout << "---------------------------------------------------------------"
42 << std::endl;
43 std::cout << "Approximate Nearest Neighbors (double precision)" << std::endl;
44 std::cout << "---------------------------------------------------------------"
45 << std::endl;
46 std::cout << std::fixed;
47 std::cout.precision(5);
48
49 da_status status;
50 bool pass = true;
51
52 // Define problem dimensions
53 da_int n_samples = 16;
54 da_int n_features = 2;
55 da_int n_queries = 3;
56 da_int k = 3; // Number of neighbors to find
57
58 // Training data (column-major order)
59 // 4 clusters of 4 points each
60 std::vector<double> X_train{0.0, 1.1, 0.0, 1.0, 6.0, 7.2, 6.1, 7.0,
61 0.0, 1.0, 0.1, 1.1, 10.0, 11.1, 10.0, 11.0,
62 -0.1, 0.0, 1.1, 1.0, 0.0, 0.1, 1.0, 1.1,
63 10.0, 10.2, 11.0, 11.1, 10.0, 10.0, 11.2, 11.0};
64 da_int ldx_train = n_samples;
65
66 // Query points
67 std::vector<double> X_test{3.5, 0.4, 5.6, 0.4, 5.0, 5.1};
68 da_int ldx_test = n_queries;
69
70 // Initialize the handle
71 da_handle handle = nullptr;
72 pass = da_handle_init_d(&handle, da_handle_approx_nn) == da_status_success;
73
74 // Set options for the approximate nearest neighbors algorithm
75 pass &= da_options_set_string(handle, "algorithm", "ivfflat") == da_status_success;
76 pass &= da_options_set_string(handle, "metric", "sqeuclidean") == da_status_success;
77 pass &= da_options_set_int(handle, "number of neighbors", k) == da_status_success;
78 pass &= da_options_set_int(handle, "n_list", 4) == da_status_success;
79 pass &= da_options_set_int(handle, "n_probe", 1) == da_status_success;
80 pass &= da_options_set_int(handle, "k-means_iter", 10) == da_status_success;
81 pass &= da_options_set_int(handle, "seed", 123) == da_status_success;
82 pass &= da_options_set_real_d(handle, "train fraction", 1.0) == da_status_success;
83
84 if (!pass) {
85 std::cout << "Error setting up handle or options.\n";
86 da_handle_print_error_message(handle);
87 da_handle_destroy(&handle);
88 return 1;
89 }
90
91 // Set training data
92 status = da_approx_nn_set_training_data_d(handle, n_samples, n_features,
93 X_train.data(), ldx_train);
94 if (status != da_status_success) {
95 std::cout << "Error setting training data.\n";
96 da_handle_print_error_message(handle);
97 da_handle_destroy(&handle);
98 return 1;
99 }
100
101 // Train the index
102 status = da_approx_nn_train_d(handle);
103 if (status != da_status_success) {
104 std::cout << "Error training the index.\n";
105 da_handle_print_error_message(handle);
106 da_handle_destroy(&handle);
107 return 1;
108 }
109 std::cout << "Index trained successfully.\n\n";
110
111 // Add data to the index
112 status = da_approx_nn_add_d(handle, n_samples, n_features, X_train.data(), ldx_train);
113 if (status != da_status_success) {
114 std::cout << "Error adding data to the index.\n";
115 da_handle_print_error_message(handle);
116 da_handle_destroy(&handle);
117 return 1;
118 }
119 std::cout << "Data added to the index.\n\n";
120
121 // Allocate output arrays for neighbor indices and distances
122 std::vector<da_int> k_ind(n_queries * k);
123 std::vector<double> k_dist(n_queries * k);
124
125 // Query the k-nearest neighbors
126 status = da_approx_nn_kneighbors_d(handle, n_queries, n_features, X_test.data(),
127 ldx_test, k_ind.data(), k_dist.data(), k, 1);
128 if (status != da_status_success) {
129 std::cout << "Error computing k-nearest neighbors.\n";
130 da_handle_print_error_message(handle);
131 da_handle_destroy(&handle);
132 return 1;
133 }
134
135 // Print neighbor indices
136 std::cout << "Approximate nearest neighbor indices:\n";
137 for (da_int i = 0; i < n_queries; i++) {
138 std::cout << " Query " << i << ": ";
139 for (da_int j = 0; j < k; j++) {
140 std::cout << k_ind[i + j * n_queries] << " ";
141 }
142 std::cout << std::endl;
143 }
144 std::cout << std::endl;
145
146 // Extract cluster centroids from the trained index
147 da_int n_list = 4;
148 da_int centroids_dim = n_list * n_features;
149 std::vector<double> centroids(centroids_dim);
150
151 status = da_handle_get_result_d(handle, da_approx_nn_cluster_centroids,
152 ¢roids_dim, centroids.data());
153 if (status == da_status_success) {
154 std::cout << "Cluster centroids:\n";
155 for (da_int i = 0; i < n_list; i++) {
156 std::cout << " Centroid " << i << ": ";
157 for (da_int j = 0; j < n_features; j++) {
158 std::cout << centroids[i + j * n_list] << " ";
159 }
160 std::cout << std::endl;
161 }
162 std::cout << std::endl;
163 }
164
165 // Extract list sizes (number of vectors in each cluster)
166 std::vector<da_int> list_sizes(n_list);
167 status = da_handle_get_result_int(handle, da_approx_nn_list_sizes, &n_list,
168 list_sizes.data());
169 if (status == da_status_success) {
170 std::cout << "Number of vectors per cluster:\n";
171 for (da_int i = 0; i < n_list; i++) {
172 std::cout << " Cluster " << i << ": " << list_sizes[i] << " vectors\n";
173 }
174 std::cout << std::endl;
175 }
176
177 // Demonstrate train_and_add (alternative workflow)
178 std::cout << "Demonstrating train_and_add workflow...\n";
179
180 // Reset training data
181 status = da_approx_nn_set_training_data_d(handle, n_samples, n_features,
182 X_train.data(), ldx_train);
183 if (status != da_status_success) {
184 std::cout << "Error resetting training data.\n";
185 da_handle_destroy(&handle);
186 return 1;
187 }
188
189 // Train and add in one step
190 status = da_approx_nn_train_and_add_d(handle);
191 if (status != da_status_success) {
192 std::cout << "Error in train_and_add.\n";
193 da_handle_print_error_message(handle);
194 da_handle_destroy(&handle);
195 return 1;
196 }
197 std::cout << "train_and_add completed successfully.\n\n";
198
199 // Demonstrate the effect of n_probe on search accuracy
200 // Lower n_probe = faster but potentially less accurate
201 // Higher n_probe = slower but more accurate (recovers exact search when n_probe = n_list)
202 std::cout << "Comparing n_probe values (approximate vs exact search):\n\n";
203
204 // Search with n_probe = 1
205 status = da_approx_nn_kneighbors_d(handle, n_queries, n_features, X_test.data(),
206 ldx_test, k_ind.data(), k_dist.data(), k, 1);
207 if (status != da_status_success) {
208 std::cout << "Error computing k-nearest neighbors.\n";
209 da_handle_destroy(&handle);
210 return 1;
211 }
212 std::cout << "n_probe = 1 (approximate):\n";
213 std::cout << " Indices:\n";
214 for (da_int i = 0; i < n_queries; i++) {
215 std::cout << " Query " << i << ": ";
216 for (da_int j = 0; j < k; j++)
217 std::cout << k_ind[i + j * n_queries] << " ";
218 std::cout << std::endl;
219 }
220 std::cout << " Distances:\n";
221 for (da_int i = 0; i < n_queries; i++) {
222 std::cout << " Query " << i << ": ";
223 for (da_int j = 0; j < k; j++)
224 std::cout << k_dist[i + j * n_queries] << " ";
225 std::cout << std::endl;
226 }
227
228 // Search with n_probe = n_list (exact search over all clusters)
229 da_options_set_int(handle, "n_probe", n_list);
230 status = da_approx_nn_kneighbors_d(handle, n_queries, n_features, X_test.data(),
231 ldx_test, k_ind.data(), k_dist.data(), k, 1);
232 if (status != da_status_success) {
233 std::cout << "Error computing k-nearest neighbors.\n";
234 da_handle_destroy(&handle);
235 return 1;
236 }
237 std::cout << "\nn_probe = " << n_list << " (exact, searches all clusters):\n";
238 std::cout << " Indices:\n";
239 for (da_int i = 0; i < n_queries; i++) {
240 std::cout << " Query " << i << ": ";
241 for (da_int j = 0; j < k; j++)
242 std::cout << k_ind[i + j * n_queries] << " ";
243 std::cout << std::endl;
244 }
245 std::cout << " Distances:\n";
246 for (da_int i = 0; i < n_queries; i++) {
247 std::cout << " Query " << i << ": ";
248 for (da_int j = 0; j < k; j++)
249 std::cout << k_dist[i + j * n_queries] << " ";
250 std::cout << std::endl;
251 }
252 std::cout << std::endl;
253
254 // Clean up
255 da_handle_destroy(&handle);
256
257 std::cout << "---------------------------------------------------------------"
258 << std::endl;
259
260 return 0;
261}