brisc 0.2.0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (61) hide show
  1. brisc-0.2.2/.git_archival.txt +2 -0
  2. brisc-0.2.2/.gitattributes +1 -0
  3. brisc-0.2.2/.gitignore +21 -0
  4. {brisc-0.2.0 → brisc-0.2.2}/MANIFEST.in +2 -0
  5. {brisc-0.2.0/brisc.egg-info → brisc-0.2.2}/PKG-INFO +1 -1
  6. brisc-0.2.2/brisc/cyutils.pxd +576 -0
  7. {brisc-0.2.0 → brisc-0.2.2}/brisc/cyutils.pyx +1 -361
  8. {brisc-0.2.0 → brisc-0.2.2}/brisc/embed.pyx +368 -307
  9. {brisc-0.2.0 → brisc-0.2.2}/brisc/harmonize.pyx +45 -2
  10. {brisc-0.2.0 → brisc-0.2.2}/brisc/hvg.pyx +101 -125
  11. {brisc-0.2.0 → brisc-0.2.2}/brisc/knn.pyx +93 -80
  12. {brisc-0.2.0 → brisc-0.2.2}/brisc/norm_factors.pyx +6 -3
  13. {brisc-0.2.0 → brisc-0.2.2}/brisc/pseudobulk.py +3 -1
  14. {brisc-0.2.0 → brisc-0.2.2}/brisc/single_cell.py +46 -24
  15. {brisc-0.2.0 → brisc-0.2.2/brisc.egg-info}/PKG-INFO +1 -1
  16. {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/SOURCES.txt +7 -0
  17. brisc-0.2.2/brisc.egg-info/scm_file_list.json +110 -0
  18. brisc-0.2.2/brisc.egg-info/scm_version.json +8 -0
  19. {brisc-0.2.0 → brisc-0.2.2}/pyproject.toml +5 -3
  20. brisc-0.2.2/tests/conftest.py +30 -0
  21. brisc-0.2.2/tests/make_test_data.py +105 -0
  22. brisc-0.2.0/brisc/cyutils.pxd +0 -347
  23. {brisc-0.2.0 → brisc-0.2.2}/LICENSE +0 -0
  24. {brisc-0.2.0 → brisc-0.2.2}/README.md +0 -0
  25. {brisc-0.2.0 → brisc-0.2.2}/brisc/__init__.py +0 -0
  26. {brisc-0.2.0 → brisc-0.2.2}/brisc/colorspacious.py +0 -0
  27. {brisc-0.2.0 → brisc-0.2.2}/brisc/concatenate.py +0 -0
  28. {brisc-0.2.0 → brisc-0.2.2}/brisc/cysparse.pyx +0 -0
  29. {brisc-0.2.0 → brisc-0.2.2}/brisc/de.py +0 -0
  30. {brisc-0.2.0 → brisc-0.2.2}/brisc/doublets.pyx +0 -0
  31. {brisc-0.2.0 → brisc-0.2.2}/brisc/hdf5.pyx +0 -0
  32. {brisc-0.2.0 → brisc-0.2.2}/brisc/kmeans.pyx +0 -0
  33. {brisc-0.2.0 → brisc-0.2.2}/brisc/labels.pyx +0 -0
  34. {brisc-0.2.0 → brisc-0.2.2}/brisc/leiden.pyx +0 -0
  35. {brisc-0.2.0 → brisc-0.2.2}/brisc/loess.pyx +0 -0
  36. {brisc-0.2.0 → brisc-0.2.2}/brisc/normalize.pyx +0 -0
  37. {brisc-0.2.0 → brisc-0.2.2}/brisc/pca.pyx +0 -0
  38. {brisc-0.2.0 → brisc-0.2.2}/brisc/pseudobulk_and_markers.pyx +0 -0
  39. {brisc-0.2.0 → brisc-0.2.2}/brisc/qc.pyx +0 -0
  40. {brisc-0.2.0 → brisc-0.2.2}/brisc/snn.pyx +0 -0
  41. {brisc-0.2.0 → brisc-0.2.2}/brisc/sparse.py +0 -0
  42. {brisc-0.2.0 → brisc-0.2.2}/brisc/threadpoolctl.py +0 -0
  43. {brisc-0.2.0 → brisc-0.2.2}/brisc/type_aliases.py +0 -0
  44. {brisc-0.2.0 → brisc-0.2.2}/brisc/utils.py +0 -0
  45. {brisc-0.2.0 → brisc-0.2.2}/brisc/validated_dict.py +0 -0
  46. {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/dependency_links.txt +0 -0
  47. {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/requires.txt +0 -0
  48. {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/top_level.txt +0 -0
  49. {brisc-0.2.0 → brisc-0.2.2}/setup.cfg +0 -0
  50. {brisc-0.2.0 → brisc-0.2.2}/setup.py +0 -0
  51. {brisc-0.2.0 → brisc-0.2.2}/tests/test_cluster.py +0 -0
  52. {brisc-0.2.0 → brisc-0.2.2}/tests/test_de.py +0 -0
  53. {brisc-0.2.0 → brisc-0.2.2}/tests/test_embedding.py +0 -0
  54. {brisc-0.2.0 → brisc-0.2.2}/tests/test_hvg.py +0 -0
  55. {brisc-0.2.0 → brisc-0.2.2}/tests/test_integration.py +0 -0
  56. {brisc-0.2.0 → brisc-0.2.2}/tests/test_markers.py +0 -0
  57. {brisc-0.2.0 → brisc-0.2.2}/tests/test_neighbors.py +0 -0
  58. {brisc-0.2.0 → brisc-0.2.2}/tests/test_normalize.py +0 -0
  59. {brisc-0.2.0 → brisc-0.2.2}/tests/test_pca.py +0 -0
  60. {brisc-0.2.0 → brisc-0.2.2}/tests/test_pseudobulk.py +0 -0
  61. {brisc-0.2.0 → brisc-0.2.2}/tests/test_qc.py +0 -0
@@ -0,0 +1,2 @@
1
+ node: $Format:%H$
2
+ describe-name: $Format:%(describe:tags=true,match=*[0-9]*)$
@@ -0,0 +1 @@
1
+ .git_archival.txt export-subst
brisc-0.2.2/.gitignore ADDED
@@ -0,0 +1,21 @@
1
+ # Python and build artifacts
2
+ brisc/__pycache__/
3
+ build/
4
+ dist/
5
+ *.egg-info/
6
+ wheelhouse/
7
+ conda-bld/
8
+ brisc/_v2/
9
+ brisc/_v3/
10
+ brisc/_v4/
11
+ brisc/_sve/
12
+
13
+ # Cython artifacts (generated C/C++ files and compiled libraries)
14
+ *.c
15
+ *.cpp
16
+ *.so
17
+ *.pyd
18
+
19
+ # Miscellaneous
20
+ *.txt
21
+ !.git_archival.txt
@@ -1,4 +1,6 @@
1
1
  recursive-include brisc *.pyx *.pxd
2
+ prune docs
3
+ prune .github
2
4
  prune brisc/x86_64_v2
3
5
  prune brisc/x86_64_v3
4
6
  prune brisc/x86_64_v4
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: brisc
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Summary: Fast single-cell toolkit
5
5
  Author-email: Michael Wainberg <m.wainberg@utoronto.ca>
6
6
  License-Expression: MIT
@@ -0,0 +1,576 @@
1
+ from cython.parallel cimport parallel, threadid
2
+ from libc.limits cimport UINT_MAX
3
+ from libc.math cimport M_PI
4
+ from libcpp.algorithm cimport lower_bound
5
+ from libcpp.cmath cimport cos, log, sqrt
6
+ from libcpp.pair cimport pair
7
+ from libcpp.vector cimport vector
8
+ from scipy.linalg.cython_blas cimport sgemm as sgemm_, sgemv as sgemv_
9
+
10
+
11
+ ctypedef fused bit_width:
12
+ unsigned
13
+ unsigned long long
14
+
15
+
16
+ ctypedef fused integer:
17
+ int
18
+ unsigned
19
+ long long
20
+ unsigned long long
21
+
22
+
23
+ ctypedef fused numeric:
24
+ int
25
+ unsigned
26
+ long long
27
+ unsigned long long
28
+ float
29
+ double
30
+
31
+
32
+ ctypedef fused signed_integer:
33
+ int
34
+ long long
35
+
36
+
37
+ cdef extern from * nogil:
38
+ """
39
+ #if defined(_WIN32)
40
+ static inline int dummy_recv(int, void *, size_t, int) {
41
+ return -1;
42
+ }
43
+ #define CYTHON_RECV dummy_recv
44
+ #else
45
+ #include <sys/socket.h>
46
+ #define CYTHON_RECV recv
47
+ #endif
48
+ """
49
+ # Cython sees "recv", but writes "CYTHON_RECV" into the generated C++ file
50
+ int recv "CYTHON_RECV"(int sockfd, void *buf, size_t len, int flags)
51
+
52
+
53
+ cdef extern from * nogil:
54
+ """
55
+ #define atomic_and(x, y) _Pragma("omp atomic") x &= y
56
+ #define atomic_or(x, y) _Pragma("omp atomic") x |= y
57
+ """
58
+ void atomic_and(bint &x, bint y)
59
+ void atomic_or(bint &x, bint y)
60
+
61
+
62
+ cdef extern from * nogil:
63
+ """
64
+ #if defined(__GNUC__) || defined(__clang__)
65
+ #define PREFETCH __builtin_prefetch
66
+ #else
67
+ #define PREFETCH
68
+ #endif
69
+ """
70
+ void PREFETCH(const void *addr)
71
+
72
+
73
+ # A drop-in replacement for Cython's `vector` that does not automatically
74
+ # zero-initialize
75
+
76
+ cdef extern from * nogil:
77
+ """
78
+ // A simplified standalone version of boost::noinit_adaptor
79
+ template<class A>
80
+ struct noinit_adaptor : A {
81
+ template<class U>
82
+ struct rebind {
83
+ typedef noinit_adaptor<typename std::allocator_traits<A>::template
84
+ rebind_alloc<U>> other;
85
+ };
86
+
87
+ template<class U>
88
+ void construct(U* p) {
89
+ ::new(p) U;
90
+ }
91
+
92
+ template<class U, class V>
93
+ void construct(U* p, const V& v) {
94
+ ::new(p) U(v);
95
+ }
96
+
97
+ template<class U>
98
+ void destroy(U* p) {
99
+ p->~U();
100
+ }
101
+ };
102
+
103
+ #include <vector>
104
+ template<class T>
105
+ using uninitialized_vector = \
106
+ std::vector<T, noinit_adaptor<std::allocator<T>>>;
107
+ """
108
+ cdef cppclass uninitialized_vector[T]:
109
+ ctypedef T value_type
110
+ ctypedef size_t size_type
111
+ ctypedef ptrdiff_t difference_type
112
+
113
+ cppclass const_iterator
114
+ cppclass iterator:
115
+ iterator() except +
116
+ iterator(iterator&) except +
117
+ T& operator*()
118
+ iterator operator++()
119
+ iterator operator--()
120
+ iterator operator++(int)
121
+ iterator operator--(int)
122
+ iterator operator+(size_type)
123
+ iterator operator-(size_type)
124
+ difference_type operator-(iterator)
125
+ difference_type operator-(const_iterator)
126
+ bint operator==(iterator)
127
+ bint operator==(const_iterator)
128
+ bint operator!=(iterator)
129
+ bint operator!=(const_iterator)
130
+ bint operator<(iterator)
131
+ bint operator<(const_iterator)
132
+ bint operator>(iterator)
133
+ bint operator>(const_iterator)
134
+ bint operator<=(iterator)
135
+ bint operator<=(const_iterator)
136
+ bint operator>=(iterator)
137
+ bint operator>=(const_iterator)
138
+ cppclass const_iterator:
139
+ const_iterator() except +
140
+ const_iterator(iterator&) except +
141
+ const_iterator(const_iterator&) except +
142
+ operator=(iterator&) except +
143
+ const T& operator*()
144
+ const_iterator operator++()
145
+ const_iterator operator--()
146
+ const_iterator operator++(int)
147
+ const_iterator operator--(int)
148
+ const_iterator operator+(size_type)
149
+ const_iterator operator-(size_type)
150
+ difference_type operator-(iterator)
151
+ difference_type operator-(const_iterator)
152
+ bint operator==(iterator)
153
+ bint operator==(const_iterator)
154
+ bint operator!=(iterator)
155
+ bint operator!=(const_iterator)
156
+ bint operator<(iterator)
157
+ bint operator<(const_iterator)
158
+ bint operator>(iterator)
159
+ bint operator>(const_iterator)
160
+ bint operator<=(iterator)
161
+ bint operator<=(const_iterator)
162
+ bint operator>=(iterator)
163
+ bint operator>=(const_iterator)
164
+
165
+ cppclass const_reverse_iterator
166
+ cppclass reverse_iterator:
167
+ reverse_iterator() except +
168
+ reverse_iterator(reverse_iterator&) except +
169
+ T& operator*()
170
+ reverse_iterator operator++()
171
+ reverse_iterator operator--()
172
+ reverse_iterator operator++(int)
173
+ reverse_iterator operator--(int)
174
+ reverse_iterator operator+(size_type)
175
+ reverse_iterator operator-(size_type)
176
+ difference_type operator-(iterator)
177
+ difference_type operator-(const_iterator)
178
+ bint operator==(reverse_iterator)
179
+ bint operator==(const_reverse_iterator)
180
+ bint operator!=(reverse_iterator)
181
+ bint operator!=(const_reverse_iterator)
182
+ bint operator<(reverse_iterator)
183
+ bint operator<(const_reverse_iterator)
184
+ bint operator>(reverse_iterator)
185
+ bint operator>(const_reverse_iterator)
186
+ bint operator<=(reverse_iterator)
187
+ bint operator<=(const_reverse_iterator)
188
+ bint operator>=(reverse_iterator)
189
+ bint operator>=(const_reverse_iterator)
190
+ cppclass const_reverse_iterator:
191
+ const_reverse_iterator() except +
192
+ const_reverse_iterator(reverse_iterator&) except +
193
+ operator=(reverse_iterator&) except +
194
+ const T& operator*()
195
+ const_reverse_iterator operator++()
196
+ const_reverse_iterator operator--()
197
+ const_reverse_iterator operator++(int)
198
+ const_reverse_iterator operator--(int)
199
+ const_reverse_iterator operator+(size_type)
200
+ const_reverse_iterator operator-(size_type)
201
+ difference_type operator-(iterator)
202
+ difference_type operator-(const_iterator)
203
+ bint operator==(reverse_iterator)
204
+ bint operator==(const_reverse_iterator)
205
+ bint operator!=(reverse_iterator)
206
+ bint operator!=(const_reverse_iterator)
207
+ bint operator<(reverse_iterator)
208
+ bint operator<(const_reverse_iterator)
209
+ bint operator>(reverse_iterator)
210
+ bint operator>(const_reverse_iterator)
211
+ bint operator<=(reverse_iterator)
212
+ bint operator<=(const_reverse_iterator)
213
+ bint operator>=(reverse_iterator)
214
+ bint operator>=(const_reverse_iterator)
215
+
216
+ uninitialized_vector() except +
217
+ uninitialized_vector(uninitialized_vector&) except +
218
+ uninitialized_vector(size_type) except +
219
+ uninitialized_vector(size_type, T&) except +
220
+ T& operator[](size_type)
221
+ bint operator==(uninitialized_vector&, uninitialized_vector&)
222
+ bint operator!=(uninitialized_vector&, uninitialized_vector&)
223
+ bint operator<(uninitialized_vector&, uninitialized_vector&)
224
+ bint operator>(uninitialized_vector&, uninitialized_vector&)
225
+ bint operator<=(uninitialized_vector&, uninitialized_vector&)
226
+ bint operator>=(uninitialized_vector&, uninitialized_vector&)
227
+ void assign(size_type, const T&)
228
+ void assign[InputIt](InputIt, InputIt) except +
229
+ T& at(size_type) except +
230
+ T& back()
231
+ iterator begin()
232
+ const_iterator const_begin "begin"()
233
+ const_iterator cbegin()
234
+ size_type capacity()
235
+ void clear()
236
+ bint empty()
237
+ iterator end()
238
+ const_iterator const_end "end"()
239
+ const_iterator cend()
240
+ iterator erase(iterator)
241
+ iterator erase(iterator, iterator)
242
+ T& front()
243
+ iterator insert(iterator, const T&) except +
244
+ iterator insert(iterator, size_type, const T&) except +
245
+ iterator insert[InputIt](iterator, InputIt, InputIt) except +
246
+ size_type max_size()
247
+ void pop_back()
248
+ void push_back(T&) except +
249
+ reverse_iterator rbegin()
250
+ const_reverse_iterator const_rbegin "rbegin"()
251
+ const_reverse_iterator crbegin()
252
+ reverse_iterator rend()
253
+ const_reverse_iterator const_rend "rend"()
254
+ const_reverse_iterator crend()
255
+ void reserve(size_type) except +
256
+ void resize(size_type) except +
257
+ void resize(size_type, T&) except +
258
+ size_type size()
259
+ void swap(uninitialized_vector&)
260
+ T* data()
261
+ const T* const_data "data"()
262
+ void shrink_to_fit() except +
263
+ iterator emplace(const_iterator, ...) except +
264
+ T& emplace_back(...) except +
265
+
266
+
267
+ ctypedef unsigned unsigned_ # hack to get templates to work with unsigned
268
+
269
+
270
+ cdef inline pair[unsigned, unsigned] get_thread_offset(
271
+ const signed_integer[::1] indptr,
272
+ const unsigned thread_index,
273
+ const unsigned num_threads) noexcept nogil:
274
+ # A helper function for load-balancing parallel iteration over sparse
275
+ # arrays. Gets the start and end row (for CSR) or column (for CSC) index
276
+ # that each thread will work on, to ensure each thread works on about the
277
+ # same number of sparse matrix elements.
278
+
279
+ cdef unsigned start, end, num_major = indptr.shape[0] - 1
280
+ cdef unsigned long long num_elements_per_thread = \
281
+ indptr[num_major] / num_threads, \
282
+ num_elements_start = thread_index * num_elements_per_thread, \
283
+ num_elements_end = (thread_index + 1) * num_elements_per_thread
284
+
285
+ start = 0 if thread_index == 0 else min(
286
+ <unsigned>(lower_bound(&indptr[0], &indptr[0] + indptr.shape[0],
287
+ num_elements_start) - &indptr[0]),
288
+ num_major)
289
+ end = num_major if thread_index == num_threads - 1 else min(
290
+ <unsigned>(lower_bound(&indptr[0], &indptr[0] + indptr.shape[0],
291
+ num_elements_end) - &indptr[0]),
292
+ num_major)
293
+ return pair[unsigned_, unsigned_](start, end)
294
+
295
+
296
+ cdef inline void get_thread_offsets(const signed_integer[::1] indptr,
297
+ unsigned* thread_offsets,
298
+ const unsigned num_threads) noexcept nogil:
299
+ # A variant of `get_thread_offset()` that gets the start and end row (for
300
+ # CSR) or column (for CSC) index that each thread will work on, for all
301
+ # threads simultaneously. For each thread, `thread_offsets[thread_index]`
302
+ # will contain the start index, and `thread_offsets[thread_index + 1]` will
303
+ # contain the end index.
304
+
305
+ cdef unsigned thread_index, num_major = indptr.shape[0] - 1
306
+ cdef unsigned long long num_elements_start, \
307
+ num_elements_per_thread = indptr[num_major] / num_threads
308
+
309
+ thread_offsets[0] = 0
310
+ for thread_index in range(1, num_threads):
311
+ num_elements_start = thread_index * num_elements_per_thread
312
+ thread_offsets[thread_index] = min(
313
+ <unsigned>(lower_bound(&indptr[0], &indptr[0] + indptr.shape[0],
314
+ num_elements_start) - &indptr[0]),
315
+ num_major)
316
+ thread_offsets[num_threads] = num_major
317
+
318
+
319
+ cdef inline void max_heap_pop(unsigned* labels_i,
320
+ float* distances_i,
321
+ const unsigned k) noexcept nogil:
322
+ # Pops the top element from the max-heap defined by `distances_i[0..k-1]`
323
+ # and `labels_i[0..k-1]`. On output the `k-1`th element is undefined.
324
+
325
+ cdef unsigned label, j = 1, child
326
+ cdef float distance
327
+ distances_i -= 1 # use 1-based indexing for easier node->child translation
328
+ labels_i -= 1
329
+ distance = distances_i[k]
330
+ label = labels_i[k]
331
+ while True:
332
+ child = j << 1
333
+ if child > k:
334
+ break
335
+ child += child < k and distances_i[child] <= distances_i[child + 1]
336
+ if distance > distances_i[child]:
337
+ break
338
+ distances_i[j] = distances_i[child]
339
+ labels_i[j] = labels_i[child]
340
+ j = child
341
+ distances_i[j] = distance
342
+ labels_i[j] = label
343
+
344
+
345
+ cdef inline void max_heap_replace_top(unsigned* labels_i,
346
+ float* distances_i,
347
+ const unsigned label,
348
+ const float distance,
349
+ const unsigned k) noexcept nogil:
350
+ # Replaces the top element from the max-heap defined by
351
+ # `distances_i[0..k-1]` and `labels_i[0..k-1]`. Equivalent to
352
+ # `std::pop_heap` followed by `std::push_heap`, but done more efficiently
353
+ # as a single operation.
354
+
355
+ cdef unsigned j = 1, child
356
+ distances_i -= 1 # use 1-based indexing for easier node->child translation
357
+ labels_i -= 1
358
+ while True:
359
+ child = j << 1
360
+ if child > k:
361
+ break
362
+ child += child < k and distances_i[child] <= distances_i[child + 1]
363
+ if distance > distances_i[child]:
364
+ break
365
+ distances_i[j] = distances_i[child]
366
+ labels_i[j] = labels_i[child]
367
+ j = child
368
+ distances_i[j] = distance
369
+ labels_i[j] = label
370
+
371
+
372
+ cdef inline void max_heap_sort(unsigned* labels_i,
373
+ float* distances_i,
374
+ const unsigned k) noexcept nogil:
375
+ cdef unsigned j, label
376
+ cdef float distance
377
+ for j in range(k):
378
+ # Save the root (maximum element)
379
+ distance = distances_i[0]
380
+ label = labels_i[0]
381
+ # Restore the heap property with reduced size `k - i`
382
+ max_heap_pop(labels_i, distances_i, k - j)
383
+ # Place the maximum element after the end of the heap
384
+ distances_i[k - j - 1] = distance
385
+ labels_i[k - j - 1] = label
386
+
387
+
388
+ cdef inline void min_heap_pop(unsigned* labels_i,
389
+ float* distances_i,
390
+ const unsigned k) noexcept nogil:
391
+ # Pops the top element from the min-heap defined by `distances_i[0..k-1]`
392
+ # and `labels_i[0..k-1]`. On output the `k-1`th element is undefined.
393
+
394
+ cdef unsigned label, j = 1, child
395
+ cdef float distance
396
+ distances_i -= 1 # use 1-based indexing for easier node->child translation
397
+ labels_i -= 1
398
+ distance = distances_i[k]
399
+ label = labels_i[k]
400
+ while True:
401
+ child = j << 1
402
+ if child > k:
403
+ break
404
+ child += child < k and distances_i[child] >= distances_i[child + 1]
405
+ if distance < distances_i[child]:
406
+ break
407
+ distances_i[j] = distances_i[child]
408
+ labels_i[j] = labels_i[child]
409
+ j = child
410
+ distances_i[j] = distance
411
+ labels_i[j] = label
412
+
413
+
414
+ cdef inline void min_heap_replace_top(unsigned* labels_i,
415
+ float* distances_i,
416
+ const unsigned label,
417
+ const float distance,
418
+ const unsigned k) noexcept nogil:
419
+ # Replaces the top element from the min-heap defined by
420
+ # `distances_i[0..k-1]` and `labels_i[0..k-1]`. Equivalent to
421
+ # `std::pop_heap` followed by `std::push_heap`, but done more efficiently
422
+ # as a single operation.
423
+
424
+ cdef unsigned j = 1, child
425
+ distances_i -= 1 # use 1-based indexing for easier node->child translation
426
+ labels_i -= 1
427
+ while True:
428
+ child = j << 1
429
+ if child > k:
430
+ break
431
+ child += child < k and distances_i[child] >= distances_i[child + 1]
432
+ if distance < distances_i[child]:
433
+ break
434
+ distances_i[j] = distances_i[child]
435
+ labels_i[j] = labels_i[child]
436
+ j = child
437
+ distances_i[j] = distance
438
+ labels_i[j] = label
439
+
440
+
441
+ cdef inline void min_heap_sort(unsigned* labels_i,
442
+ float* distances_i,
443
+ const unsigned k) noexcept nogil:
444
+ cdef unsigned j, label
445
+ cdef float distance
446
+ for j in range(k):
447
+ # Save the root (minimum element)
448
+ distance = distances_i[0]
449
+ label = labels_i[0]
450
+ # Restore the heap property with reduced size `k - i`
451
+ min_heap_pop(labels_i, distances_i, k - j)
452
+ # Place the minimum element after the end of the heap
453
+ distances_i[k - j - 1] = distance
454
+ labels_i[k - j - 1] = label
455
+
456
+
457
+ cdef inline float norm(float[::1] array) noexcept nogil:
458
+ cdef unsigned i
459
+ cdef float norm = 0
460
+ for i in range(array.shape[0]):
461
+ norm += array[i] * array[i]
462
+ return sqrt(norm)
463
+
464
+
465
+ cdef inline void partial_distances(const float* A,
466
+ const float* B,
467
+ const float* B_norms,
468
+ float* distances,
469
+ const unsigned num_A,
470
+ const unsigned num_B,
471
+ const unsigned ldc,
472
+ const unsigned num_dimensions) \
473
+ noexcept nogil:
474
+ # Calculate the "partial" distance from each row of A (of shape
475
+ # `num_A × num_dimensions`) to each row of B (of shape
476
+ # `num_B × num_dimensions`). Use the identity:
477
+ # ||A - B||² = ||A||² - 2 * A.dot(B.T) + ||B||²,
478
+ # but skip calculating the ||A||² term since the closest row of `B` for a
479
+ # given row of `A` does not depend on ||A||². This is why we call it a
480
+ # "partial" distance. `ldc` is the row stride of `distances`; pass
481
+ # `ldc == num_B` for the untiled case.
482
+
483
+ cdef char transA = b'T', transB = b'N'
484
+ cdef float alpha = -2, beta = 1
485
+ cdef unsigned i, j
486
+
487
+ for i in range(num_A):
488
+ for j in range(num_B):
489
+ # distances = ||B||²
490
+ distances[i * ldc + j] = B_norms[j]
491
+
492
+ # distances -= 2 * A.dot(B.T)
493
+ sgemm(transA, transB, num_B, num_A, num_dimensions, alpha, B,
494
+ num_dimensions, A, num_dimensions, beta, distances, ldc)
495
+
496
+
497
+ cdef inline unsigned rand(unsigned long long* state) noexcept nogil:
498
+ cdef unsigned long long x = state[0]
499
+ state[0] = x * 6364136223846793005ULL + 1442695040888963407ULL
500
+ cdef unsigned s = (x ^ (x >> 18)) >> 27
501
+ cdef unsigned rot = x >> 59
502
+ return (s >> rot) | (s << ((-rot) & 31))
503
+
504
+
505
+ cdef inline unsigned randint(const unsigned bound, unsigned long long* state) \
506
+ noexcept nogil:
507
+ # Lemire's method
508
+ cdef unsigned long long m = \
509
+ <unsigned long long> rand(state) * <unsigned long long> bound
510
+ cdef unsigned threshold, l = <unsigned> m
511
+ if l < bound:
512
+ threshold = -bound % bound
513
+ while l < threshold:
514
+ m = <unsigned long long> rand(state) * <unsigned long long> bound
515
+ l = <unsigned> m
516
+ return <unsigned>(m >> 32)
517
+
518
+
519
+ cdef inline float random_uniform(unsigned long long* state) noexcept nogil:
520
+ # Returns a random number in U(0, 1)
521
+ return <float> rand(state) / UINT_MAX
522
+
523
+
524
+ cdef inline float random_normal(unsigned long long* state) noexcept nogil:
525
+ # Samples a random number from the standard normal distribution
526
+ # using the Box-Muller transform
527
+ cdef float u1, u2, r, theta
528
+ while True:
529
+ u1 = random_uniform(state)
530
+ if u1:
531
+ break
532
+ u2 = random_uniform(state)
533
+ r = sqrt(<float> -2 * log(u1))
534
+ theta = <float> 2 * <float> M_PI * u2
535
+ return r * cos(theta)
536
+
537
+
538
+ cdef inline void sgemv(const char trans,
539
+ const int m,
540
+ const int n,
541
+ const float alpha,
542
+ const float* a,
543
+ const int lda,
544
+ const float* x,
545
+ const int incx,
546
+ const float beta,
547
+ float* y,
548
+ const int incy) noexcept nogil:
549
+ sgemv_(<char*> &trans, <int*> &m, <int*> &n, <float*> &alpha, <float*> a,
550
+ <int*> &lda, <float*> x, <int*> &incx, <float*> &beta, y,
551
+ <int*> &incy)
552
+
553
+
554
+ cdef inline void sgemm(const char transa,
555
+ const char transb,
556
+ const int m,
557
+ const int n,
558
+ const int k,
559
+ const float alpha,
560
+ const float* a,
561
+ const int lda,
562
+ const float* b,
563
+ const int ldb,
564
+ const float beta,
565
+ float* c,
566
+ const int ldc) noexcept nogil:
567
+ sgemm_(<char*> &transa, <char*> &transb, <int*> &m, <int*> &n, <int*> &k,
568
+ <float*> &alpha, <float*> a, <int*> &lda, <float*> b, <int*> &ldb,
569
+ <float*> &beta, c, <int*> &ldc)
570
+
571
+
572
+ cdef inline unsigned long long srand(const unsigned long long seed) \
573
+ noexcept nogil:
574
+ cdef unsigned long long state = seed + 1442695040888963407ULL
575
+ rand(&state)
576
+ return state