annoy
1# Copyright (c) 2013 Spotify AB
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14
15import random16
17import numpy18
19from annoy import AnnoyIndex20
21
22def test_random_holes():23f = 1024index = AnnoyIndex(f, "angular")25valid_indices = random.sample(range(2000), 1000) # leave holes26for i in valid_indices:27v = numpy.random.normal(size=(f,))28index.add_item(i, v)29index.build(10)30for i in valid_indices:31js = index.get_nns_by_item(i, 10000)32for j in js:33assert j in valid_indices34for i in range(1000):35v = numpy.random.normal(size=(f,))36js = index.get_nns_by_vector(v, 10000)37for j in js:38assert j in valid_indices39
40
41def _test_holes_base(n, f=100, base_i=100000):42annoy = AnnoyIndex(f, "angular")43for i in range(n):44annoy.add_item(base_i + i, numpy.random.normal(size=(f,)))45annoy.build(100)46res = annoy.get_nns_by_item(base_i, n)47assert set(res) == set([base_i + i for i in range(n)])48
49
50def test_root_one_child():51# See https://github.com/spotify/annoy/issues/22352_test_holes_base(1)53
54
55def test_root_two_children():56_test_holes_base(2)57
58
59def test_root_some_children():60# See https://github.com/spotify/annoy/issues/29561_test_holes_base(10)62
63
64def test_root_many_children():65_test_holes_base(1000)66