-
Notifications
You must be signed in to change notification settings - Fork 0
/
loc_var_lk.pyx
44 lines (41 loc) · 1.1 KB
/
loc_var_lk.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
import numpy as np
import cython
from libc.math cimport sqrt, exp
from cython.parallel import prange
cimport openmp
import timeit
@cython.boundscheck(False)
@cython.wraparound(False)
def array_double(long N, long M, int num_threads=0):
cdef:
long i, j, k
double cnt = 0.0, *cnt_ptr = &cnt
double var[3]
double summ
double res = 0.0, *res_ptr = &res
double[:, :] inp = np.ones((N, M))
double[:, :] out = np.empty((N, M))
openmp.omp_lock_t lock
if num_threads != 0:
openmp.omp_set_dynamic(0)
openmp.omp_set_num_threads(num_threads)
start = timeit.default_timer()
with nogil:
openmp.omp_init_lock(&lock)
for i in prange(N, schedule='dynamic'):
for j in range(M):
out[i, j] = sqrt(exp(-sqrt(inp[i, j]*(i+j)))) * sqrt(exp(-sqrt(inp[i, j]*(i+j))))
for k in range(3):
var[k] = M
summ = 0
for k in range(3):
summ = summ + var[k]
openmp.omp_set_lock(&lock)
cnt_ptr[0] = cnt_ptr[0] + 1
res_ptr[0] = res_ptr[0] + summ
openmp.omp_unset_lock(&lock)
openmp.omp_destroy_lock(&lock)
stop = timeit.default_timer()
print(stop - start)
print(cnt)
print(res)