diff --git a/nifty/operators/fft_operator/transformations/lm_transformation_factory.py b/nifty/operators/fft_operator/transformations/lm_transformation_factory.pyx similarity index 50% rename from nifty/operators/fft_operator/transformations/lm_transformation_factory.py rename to nifty/operators/fft_operator/transformations/lm_transformation_factory.pyx index ad1f038a7de822c56802fb798235607409a90465..5ea953302fd406b29b2ea8a4df51b4278e5d303a 100644 --- a/nifty/operators/fft_operator/transformations/lm_transformation_factory.py +++ b/nifty/operators/fft_operator/transformations/lm_transformation_factory.pyx @@ -1,4 +1,6 @@ + import numpy as np +cimport numpy as np def buildLm(inp, **kwargs): if inp.dtype == np.dtype('float32'): @@ -12,36 +14,40 @@ def buildIdx(inp, **kwargs): else: return _buildIdx(inp, **kwargs) -def _buildIdx_f(nr, lmax): - size = (lmax+1)*(lmax+1) +cpdef np.ndarray[np.float32_t, ndim=1] _buildIdx_f(np.ndarray[np.complex64_t, +ndim=1] nr, np.int_t lmax): + cdef np.int size = (lmax+1)*(lmax+1) - final=np.zeros([size], dtype=np.float32) + cdef np.ndarray final=np.zeros([size], dtype=np.float32) 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 return final -def _buildIdx(nr, lmax): - size = (lmax+1)*(lmax+1) +cpdef np.ndarray[np.float64_t, ndim=1] _buildIdx(np.ndarray[np.complex128_t, +ndim=1] nr, np.int_t lmax): + cdef np.int size = (lmax+1)*(lmax+1) - final=np.zeros([size], dtype=np.float64) + cdef np.ndarray final=np.zeros([size], dtype=np.float64) 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 return final -def _buildLm_f(nr, lmax): - size = (len(nr)-lmax-1)/2+lmax+1 +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 - res=np.zeros([size], dtype=np.complex64) + cdef np.ndarray res=np.zeros([size], dtype=np.complex64) res[0:lmax+1] = nr[0:lmax+1] res[lmax+1:] = np.sqrt(0.5)*(nr[lmax+1::2] + 1j*nr[lmax+2::2]) return res -def _buildLm(nr, lmax): - size = (len(nr)-lmax-1)/2+lmax+1 +cpdef np.ndarray[np.complex128_t, ndim=1] _buildLm(np.ndarray[np.float64_t, +ndim=1] nr, np.int_t lmax): + cdef np.int size = (len(nr)-lmax-1)/2+lmax+1 - res=np.zeros([size], dtype=np.complex128) + cdef np.ndarray res=np.zeros([size], dtype=np.complex128) res[0:lmax+1] = nr[0:lmax+1] res[lmax+1:] = np.sqrt(0.5)*(nr[lmax+1::2] + 1j*nr[lmax+2::2]) return res diff --git a/nifty/spaces/lm_space/lm_helper.pyx b/nifty/spaces/lm_space/lm_helper.pyx new file mode 100644 index 0000000000000000000000000000000000000000..e2a977e07e1d83aee21d09c4612ca41a6dcc9999 --- /dev/null +++ b/nifty/spaces/lm_space/lm_helper.pyx @@ -0,0 +1,23 @@ +import numpy as np +cimport numpy as np + +cpdef _distance_array_helper(np.ndarray[np.int_t] index_array, np.int_t lmax): + cdef np.int size = (lmax + 1) * (lmax + 1) + + cdef np.ndarray res = np.zeros([len(index_array)], dtype=np.float128) + + for idx, index_full in enumerate(index_array): + if index_full <= lmax: + index_half = index_full + else: + if (index_full - lmax) % 2 == 0: + index_half = (index_full + lmax) / 2; + else: + index_half = (index_full + lmax + 1) / 2; + + 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 + + return res diff --git a/nifty/spaces/lm_space/lm_space.py b/nifty/spaces/lm_space/lm_space.py index c8c098fca8afc1e2af2a8457d785490fecf35bd5..19ec9ab0e23ab229aad444868cf11c2127b48964 100644 --- a/nifty/spaces/lm_space/lm_space.py +++ b/nifty/spaces/lm_space/lm_space.py @@ -25,18 +25,13 @@ from nifty.spaces.space import Space from nifty.config import nifty_configuration as gc,\ dependency_injector as gdi +from lm_helper import _distance_array_helper + from d2o import arange gl = gdi.get('libsharp_wrapper_gl') hp = gdi.get('healpy') -def _distance_array_helper(index_array, lmax): - u=2*lmax+1 - index_half=(index_array+np.minimum(lmax,index_array)+1)//2 - m = (np.ceil((u - np.sqrt(u*u - 8 * (index_half - lmax))) / 2)).astype(int) - res = (index_half - m * (u - m) // 2).astype(np.float64) - return res - class LMSpace(Space): """ diff --git a/setup.py b/setup.py index f0097a315df5f99640a72d1aa8103ad962fd31c4..4d8c1a9c8bc5e147c55b796c21ea991a727eb0c6 100644 --- a/setup.py +++ b/setup.py @@ -35,6 +35,9 @@ setup(name="ift_nifty", package_dir={"nifty": "nifty"}, zip_safe=False, ext_modules=cythonize([ + "nifty/operators/fft_operator/transformations/" + "lm_transformation_factory.pyx", + "nifty/spaces/lm_space/lm_helper.pyx", "nifty/operators/smoothing_operator/smooth_util.pyx" ]), include_dirs=[numpy.get_include()],