lm_transformation_factory.pyx 1.73 KB
Newer Older
1

csongor's avatar
csongor committed
2 3 4
import numpy as np
cimport numpy as np

csongor's avatar
csongor committed
5 6 7 8 9 10 11 12 13 14 15 16
def buildLm(inp, **kwargs):
    if inp.dtype == np.dtype('float32'):
        return _buildLm_f(inp, **kwargs)
    else:
        return _buildLm(inp, **kwargs)

def buildIdx(inp, **kwargs):
    if inp.dtype == np.dtype('complex64'):
        return _buildIdx_f(inp, **kwargs)
    else:
        return _buildIdx(inp, **kwargs)

17 18
cpdef np.ndarray[np.float32_t, ndim=1]  _buildIdx_f(np.ndarray[np.complex64_t,
ndim=1] nr, np.int_t lmax):
csongor's avatar
csongor committed
19 20 21
    cdef np.int size = (lmax+1)*(lmax+1)

    cdef np.ndarray final=np.zeros([size], dtype=np.float32)
22 23 24
    final[0:lmax+1] = nr[0:lmax+1].real
    final[lmax+1::2] = np.sqrt(2)*nr[lmax+1:].real
    final[lmax+2::2] = np.sqrt(2)*nr[lmax+1:].imag
csongor's avatar
csongor committed
25 26 27
    return final

cpdef np.ndarray[np.float64_t, ndim=1]  _buildIdx(np.ndarray[np.complex128_t,
28
ndim=1] nr, np.int_t lmax):
csongor's avatar
csongor committed
29 30 31
    cdef np.int size = (lmax+1)*(lmax+1)

    cdef np.ndarray final=np.zeros([size], dtype=np.float64)
32 33 34
    final[0:lmax+1] = nr[0:lmax+1].real
    final[lmax+1::2] = np.sqrt(2)*nr[lmax+1:].real
    final[lmax+2::2] = np.sqrt(2)*nr[lmax+1:].imag
35 36
    return final

csongor's avatar
csongor committed
37 38 39 40 41 42
cpdef np.ndarray[np.complex64_t, ndim=1] _buildLm_f(np.ndarray[np.float32_t,
ndim=1] nr, np.int_t lmax):
    cdef np.int size = (len(nr)-lmax-1)/2+lmax+1

    cdef np.ndarray res=np.zeros([size], dtype=np.complex64)
    res[0:lmax+1] = nr[0:lmax+1]
43
    res[lmax+1:] = np.sqrt(0.5)*(nr[lmax+1::2] + 1j*nr[lmax+2::2])
csongor's avatar
csongor committed
44 45 46 47
    return res

cpdef np.ndarray[np.complex128_t, ndim=1] _buildLm(np.ndarray[np.float64_t,
ndim=1] nr, np.int_t lmax):
csongor's avatar
csongor committed
48 49 50 51
    cdef np.int size = (len(nr)-lmax-1)/2+lmax+1

    cdef np.ndarray res=np.zeros([size], dtype=np.complex128)
    res[0:lmax+1] = nr[0:lmax+1]
52
    res[lmax+1:] = np.sqrt(0.5)*(nr[lmax+1::2] + 1j*nr[lmax+2::2])
53
    return res