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.
- brisc-0.2.2/.git_archival.txt +2 -0
- brisc-0.2.2/.gitattributes +1 -0
- brisc-0.2.2/.gitignore +21 -0
- {brisc-0.2.0 → brisc-0.2.2}/MANIFEST.in +2 -0
- {brisc-0.2.0/brisc.egg-info → brisc-0.2.2}/PKG-INFO +1 -1
- brisc-0.2.2/brisc/cyutils.pxd +576 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/cyutils.pyx +1 -361
- {brisc-0.2.0 → brisc-0.2.2}/brisc/embed.pyx +368 -307
- {brisc-0.2.0 → brisc-0.2.2}/brisc/harmonize.pyx +45 -2
- {brisc-0.2.0 → brisc-0.2.2}/brisc/hvg.pyx +101 -125
- {brisc-0.2.0 → brisc-0.2.2}/brisc/knn.pyx +93 -80
- {brisc-0.2.0 → brisc-0.2.2}/brisc/norm_factors.pyx +6 -3
- {brisc-0.2.0 → brisc-0.2.2}/brisc/pseudobulk.py +3 -1
- {brisc-0.2.0 → brisc-0.2.2}/brisc/single_cell.py +46 -24
- {brisc-0.2.0 → brisc-0.2.2/brisc.egg-info}/PKG-INFO +1 -1
- {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/SOURCES.txt +7 -0
- brisc-0.2.2/brisc.egg-info/scm_file_list.json +110 -0
- brisc-0.2.2/brisc.egg-info/scm_version.json +8 -0
- {brisc-0.2.0 → brisc-0.2.2}/pyproject.toml +5 -3
- brisc-0.2.2/tests/conftest.py +30 -0
- brisc-0.2.2/tests/make_test_data.py +105 -0
- brisc-0.2.0/brisc/cyutils.pxd +0 -347
- {brisc-0.2.0 → brisc-0.2.2}/LICENSE +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/README.md +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/__init__.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/colorspacious.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/concatenate.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/cysparse.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/de.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/doublets.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/hdf5.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/kmeans.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/labels.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/leiden.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/loess.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/normalize.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/pca.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/pseudobulk_and_markers.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/qc.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/snn.pyx +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/sparse.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/threadpoolctl.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/type_aliases.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/utils.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc/validated_dict.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/dependency_links.txt +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/requires.txt +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/brisc.egg-info/top_level.txt +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/setup.cfg +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/setup.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_cluster.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_de.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_embedding.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_hvg.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_integration.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_markers.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_neighbors.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_normalize.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_pca.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_pseudobulk.py +0 -0
- {brisc-0.2.0 → brisc-0.2.2}/tests/test_qc.py +0 -0
|
@@ -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
|
|
@@ -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
|