|
| 1 | +import unittest |
| 2 | + |
| 3 | +import numpy as np |
| 4 | + |
| 5 | +import hnswlib |
| 6 | + |
| 7 | +class RandomSelfTestCase(unittest.TestCase): |
| 8 | + def testRandomSelf(self): |
| 9 | + |
| 10 | + data1 = np.asarray([[1, 0, 0], |
| 11 | + [0, 1, 0], |
| 12 | + [0, 0, 1], |
| 13 | + [1, 0, 1], |
| 14 | + [1, 1, 1], |
| 15 | + ]) |
| 16 | + |
| 17 | + for space, expected_distances in [ |
| 18 | + ('l2', [[0., 1., 2., 2., 2.]]), |
| 19 | + ('ip', [[-2., -1., 0., 0., 0.]]), |
| 20 | + ('cosine', [[0, 1.835e-01, 4.23e-01, 4.23e-01, 4.23e-01]])]: |
| 21 | + |
| 22 | + for rightdim in range(1, 128, 3): |
| 23 | + for leftdim in range(1, 32, 5): |
| 24 | + data2 = np.concatenate( |
| 25 | + [np.zeros([data1.shape[0], leftdim]), data1, np.zeros([data1.shape[0], rightdim])], axis=1) |
| 26 | + dim = data2.shape[1] |
| 27 | + p = hnswlib.Index(space=space, dim=dim) |
| 28 | + p.init_index(max_elements=5, ef_construction=100, M=16) |
| 29 | + |
| 30 | + p.set_ef(10) |
| 31 | + |
| 32 | + p.add_items(data2) |
| 33 | + |
| 34 | + # Query the elements for themselves and measure recall: |
| 35 | + labels, distances = p.knn_query(np.asarray(data2[-1:]), k=5) |
| 36 | + |
| 37 | + |
| 38 | + diff=np.mean(np.abs(distances-expected_distances)) |
| 39 | + print(dim,space, diff) |
| 40 | + self.assertAlmostEqual(diff, 0, delta=1e-3) |
0 commit comments