-
Notifications
You must be signed in to change notification settings - Fork 70
/
summator.pyx
100 lines (78 loc) · 2.43 KB
/
summator.pyx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
# cython: language_level=3, boundscheck=False, wraparound=False, cdivision=True
"""
This is the randomization method summator, implemented in cython.
"""
import numpy as np
from cython.parallel import parallel, prange
try:
cimport openmp
except:
...
cimport numpy as np
from libc.math cimport cos, sin
def set_num_threads(num_threads):
cdef int num_threads_c = 1
if num_threads is None:
try:
num_threads_c = openmp.omp_get_num_procs()
except:
...
else:
num_threads_c = num_threads
return num_threads_c
def summate(
const double[:, :] cov_samples,
const double[:] z_1,
const double[:] z_2,
const double[:, :] pos,
num_threads=None,
):
cdef int i, j, d
cdef double phase
cdef int dim = pos.shape[0]
cdef int X_len = pos.shape[1]
cdef int N = cov_samples.shape[1]
cdef double[:] summed_modes = np.zeros(X_len, dtype=float)
cdef int num_threads_c = set_num_threads(num_threads)
for i in prange(X_len, nogil=True, num_threads=num_threads_c):
for j in range(N):
phase = 0.
for d in range(dim):
phase += cov_samples[d, j] * pos[d, i]
summed_modes[i] += z_1[j] * cos(phase) + z_2[j] * sin(phase)
return np.asarray(summed_modes)
cdef (double) abs_square(const double[:] vec) nogil:
cdef int i
cdef double r = 0.
for i in range(vec.shape[0]):
r += vec[i]**2
return r
def summate_incompr(
const double[:, :] cov_samples,
const double[:] z_1,
const double[:] z_2,
const double[:, :] pos,
num_threads=None,
):
cdef int i, j, d
cdef double phase
cdef double k_2
cdef int dim = pos.shape[0]
cdef double[:] e1 = np.zeros(dim, dtype=float)
e1[0] = 1.
cdef double[:] proj = np.empty(dim)
cdef int X_len = pos.shape[1]
cdef int N = cov_samples.shape[1]
cdef double[:, :] summed_modes = np.zeros((dim, X_len), dtype=float)
for i in range(X_len):
for j in range(N):
k_2 = abs_square(cov_samples[:, j])
phase = 0.
for d in range(dim):
phase += cov_samples[d, j] * pos[d, i]
for d in range(dim):
proj[d] = e1[d] - cov_samples[d, j] * cov_samples[0, j] / k_2
summed_modes[d, i] += (
proj[d] * (z_1[j] * cos(phase) + z_2[j] * sin(phase))
)
return np.asarray(summed_modes)