annoy

1
/*
2
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
3
# use this file except in compliance with the License. You may obtain a copy of
4
# the License at
5
#
6
# http://www.apache.org/licenses/LICENSE-2.0
7
#
8
# Unless required by applicable law or agreed to in writing, software
9
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
10
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
11
# License for the specific language governing permissions and limitations under
12
# the License.
13
*/
14

15
package annoyindex_test
16

17
import (
18
       "annoyindex"
19
       "os"
20
       "testing"
21
       "math"
22
       "math/rand"
23
       "github.com/stretchr/testify/assert"
24
       "github.com/stretchr/testify/suite"
25
)
26

27
type AnnoyTestSuite struct {
28
    suite.Suite
29
}
30

31
func Round(f float64) float64 {
32
    return math.Floor(f + 0.5)
33
}
34

35
func RoundPlus(f float64, places int) (float64) {
36
     shift := math.Pow(10, float64(places))
37
     return Round(f * shift) / shift
38
}
39

40
func (suite *AnnoyTestSuite) SetupTest() {
41
}
42

43
func (suite *AnnoyTestSuite) TestFileHandling() {
44
     index := annoyindex.NewAnnoyIndexAngular(3)
45
     index.AddItem(0, []float32{0, 0, 1})
46
     index.AddItem(1, []float32{0, 1, 0})
47
     index.AddItem(2, []float32{1, 0, 0})
48
     index.Build(10)
49

50
     index.Save("go_test.ann")
51

52
     info, err := os.Stat("go_test.ann")
53
     if err != nil {
54
        assert.Fail(suite.T(), "Failed to create file, file not found")
55
     }
56
     if info.Size() == 0 {
57
        assert.Fail(suite.T(), "Failed to create file, file size zero")
58
     }
59

60
     annoyindex.DeleteAnnoyIndexAngular(index)
61

62
     index = annoyindex.NewAnnoyIndexAngular(3)
63
     if ret := index.Load("go_test.ann"); ret == false {
64
        assert.Fail(suite.T(), "Failed to load file")
65
     }
66

67
     os.Remove("go_test.ann")
68
     index.Save("go_test2.ann", false)
69

70
     info, err = os.Stat("go_test2.ann")
71
     if err != nil {
72
        assert.Fail(suite.T(), "Failed to create file without prefault, file not found")
73
     }
74
     if info.Size() == 0 {
75
        assert.Fail(suite.T(), "Failed to create file without prefault, file size zero")
76
     }
77

78
     annoyindex.DeleteAnnoyIndexAngular(index)
79

80
     index = annoyindex.NewAnnoyIndexAngular(3)
81
     if ret := index.Load("go_test2.ann", false); ret == false {
82
        assert.Fail(suite.T(), "Failed to load file without prefault")
83
     }
84

85
     os.Remove("go_test2.ann")
86
     index.Save("go_test3.ann", true)
87

88
     info, err = os.Stat("go_test3.ann")
89
     if err != nil {
90
        assert.Fail(suite.T(), "Failed to create file allowing prefault, file not found")
91
     }
92
     if info.Size() == 0 {
93
        assert.Fail(suite.T(), "Failed to create file allowing prefault, file size zero")
94
     }
95

96
     annoyindex.DeleteAnnoyIndexAngular(index)
97

98
     index = annoyindex.NewAnnoyIndexAngular(3)
99
     if ret := index.Load("go_test3.ann", true); ret == false {
100
        assert.Fail(suite.T(), "Failed to load file allowing prefault")
101
     }
102
     annoyindex.DeleteAnnoyIndexAngular(index)
103

104
     os.Remove("go_test3.ann")
105
}
106

107
func (suite *AnnoyTestSuite) TestOnDiskBuild() {
108
     index := annoyindex.NewAnnoyIndexAngular(3)
109
     index.OnDiskBuild("go_test.ann");
110

111
     info, err := os.Stat("go_test.ann")
112
     if err != nil {
113
        assert.Fail(suite.T(), "Failed to create file, file not found")
114
     }
115
     if info.Size() == 0 {
116
        assert.Fail(suite.T(), "Failed to create file, file size zero")
117
     }
118

119
     index.AddItem(0, []float32{0, 0, 1})
120
     index.AddItem(1, []float32{0, 1, 0})
121
     index.AddItem(2, []float32{1, 0, 0})
122
     index.Build(10)
123

124
     index.Unload();
125
     index.Load("go_test.ann");
126

127
     var result []int
128
     index.GetNnsByVector([]float32{3, 2, 1}, 3, -1, &result)
129
     assert.Equal(suite.T(), []int{2, 1, 0}, result)
130

131
     index.GetNnsByVector([]float32{1, 2, 3}, 3, -1, &result)
132
     assert.Equal(suite.T(), []int{0, 1, 2}, result)
133

134
     index.GetNnsByVector([]float32{2, 0, 1}, 3, -1, &result)
135
     assert.Equal(suite.T(), []int{2, 0, 1}, result)
136

137
     annoyindex.DeleteAnnoyIndexAngular(index)
138

139
     os.Remove("go_test.ann")
140
}
141

142
func (suite *AnnoyTestSuite) TestGetNnsByVector() {
143
     index := annoyindex.NewAnnoyIndexAngular(3)
144
     index.AddItem(0, []float32{0, 0, 1})
145
     index.AddItem(1, []float32{0, 1, 0})
146
     index.AddItem(2, []float32{1, 0, 0})
147
     index.Build(10)
148

149
     var result []int
150
     index.GetNnsByVector([]float32{3, 2, 1}, 3, -1, &result)
151
     assert.Equal(suite.T(), []int{2, 1, 0}, result)
152

153
     index.GetNnsByVector([]float32{1, 2, 3}, 3, -1, &result)
154
     assert.Equal(suite.T(), []int{0, 1, 2}, result)
155

156
     index.GetNnsByVector([]float32{2, 0, 1}, 3, -1, &result)
157
     assert.Equal(suite.T(), []int{2, 0, 1}, result)
158

159
     annoyindex.DeleteAnnoyIndexAngular(index)
160
}
161

162
func (suite *AnnoyTestSuite) TestGetNnsByItem() {
163
     index := annoyindex.NewAnnoyIndexAngular(3)
164
     index.AddItem(0, []float32{2, 1, 0})
165
     index.AddItem(1, []float32{1, 2, 0})
166
     index.AddItem(2, []float32{0, 0, 1})
167
     index.Build(10)
168

169
     var result []int
170
     index.GetNnsByItem(0, 3, -1, &result)
171
     assert.Equal(suite.T(), []int{0, 1, 2}, result)
172

173
     index.GetNnsByItem(1, 3, -1, &result)
174
     assert.Equal(suite.T(), []int{1, 0, 2}, result)
175

176
     annoyindex.DeleteAnnoyIndexAngular(index)
177
}
178

179
func (suite *AnnoyTestSuite) TestGetItem() {
180
     index := annoyindex.NewAnnoyIndexAngular(3)
181
     index.AddItem(0, []float32{2, 1, 0})
182
     index.AddItem(1, []float32{1, 2, 0})
183
     index.AddItem(2, []float32{0, 0, 1})
184
     index.Build(10)
185

186
     var result []float32
187

188
     index.GetItem(0, &result)
189
     assert.Equal(suite.T(), []float32{2, 1, 0}, result)
190

191
     index.GetItem(1, &result)
192
     assert.Equal(suite.T(), []float32{1, 2, 0}, result)
193

194
     index.GetItem(2, &result)
195
     assert.Equal(suite.T(), []float32{0, 0, 1}, result)
196

197
     annoyindex.DeleteAnnoyIndexAngular(index)
198
}
199

200

201
func (suite *AnnoyTestSuite) TestGetDistance() {
202
     index := annoyindex.NewAnnoyIndexAngular(2)
203
     index.AddItem(0, []float32{0, 1})
204
     index.AddItem(1, []float32{1, 1})
205
     index.Build(10)
206

207
     assert.Equal(suite.T(), RoundPlus(math.Pow(2 * (1.0 - math.Pow(2, -0.5)), 0.5), 3), RoundPlus(float64(index.GetDistance(0, 1)), 3))
208

209
     annoyindex.DeleteAnnoyIndexAngular(index)
210
}
211

212
func (suite *AnnoyTestSuite) TestGetDotProductDistance() {
213
    index := annoyindex.NewAnnoyIndexDotProduct(2)
214
    index.AddItem(0, []float32{0, 1})
215
    index.AddItem(1, []float32{1, 1})
216
    index.Build(10)
217

218
    assert.True(suite.T(),
219
        math.Abs(1.0-float64(index.GetDistance(0, 1))) < 0.00001)
220

221
    annoyindex.DeleteAnnoyIndexDotProduct(index)
222
}
223

224
func (suite *AnnoyTestSuite) TestLargeEuclideanIndex() {
225
     index := annoyindex.NewAnnoyIndexEuclidean(10)
226

227
     for j := 0; j < 10000; j += 2 {
228
         p := make([]float32, 0, 10)
229
         for i := 0; i < 10; i++ {
230
	     p = append(p, rand.Float32())
231
	 }
232
	 x := make([]float32, 0, 10)
233
	 for i := 0; i < 10; i++ {
234
	     x = append(x, 1 + p[i] + rand.Float32() * 1e-2)
235
	 }
236
	 y := make([]float32, 0, 10)
237
	 for i := 0; i < 10; i++ {
238
	     y = append(y, 1 + p[i] + rand.Float32() * 1e-2)
239
	 }
240
	 index.AddItem(j, x)
241
	 index.AddItem(j + 1, y)
242
     }
243
     index.Build(10)
244
     for j := 0; j < 10000; j += 2 {
245
         var result []int
246
	 index.GetNnsByItem(j, 2, -1, &result)
247

248
         assert.Equal(suite.T(), result, []int{j, j + 1})
249

250
	 index.GetNnsByItem(j + 1, 2, -1, &result)
251
	 assert.Equal(suite.T(), result, []int{j + 1, j})
252
     }
253
     annoyindex.DeleteAnnoyIndexEuclidean(index)
254
}
255

256
func TestAnnoyTestSuite(t *testing.T) {
257
    suite.Run(t, new(AnnoyTestSuite))
258
}
259