lm_helper.pyx 790 Bytes
Newer Older
Jait Dixit's avatar
Jait Dixit committed
1
2
3
import numpy as np
cimport numpy as np

4
cpdef _distance_array_helper(np.ndarray[np.int_t] index_array, np.int_t lmax):
Jait Dixit's avatar
Jait Dixit committed
5
6
    cdef np.int size = (lmax + 1) * (lmax  + 1)

7
    cdef np.ndarray res = np.zeros([len(index_array)], dtype=np.float128)
Jait Dixit's avatar
Jait Dixit committed
8

9
10
11
12
    for idx, index_full in enumerate(index_array):
        if index_full <= lmax:
            index_half = index_full
        else:
13
14
15
16
17
            if (index_full - lmax) % 2 == 0:
                index_half = (index_full + lmax) / 2;
            else:
                index_half = (index_full + lmax + 1) / 2;

18
19
20
21
        m = (np.ceil(((2 * lmax + 1) -
                      np.sqrt((2 * lmax + 1)**2 -
                              8 * (index_half - lmax))) / 2)).astype(int)
        res[idx] = index_half - m * (2 * lmax + 1 - m) // 2
Jait Dixit's avatar
Jait Dixit committed
22

23
    return res