WIP - cythonify

nifty/operators/smoothing_operator/Cython/extended.py

-import numpy as np
-
-def smooth_power_2s(power, k, exclude=1, smooth_length=None):
-
-    if smooth_length == 0:
-        # No smoothing requested, just return the input array.
-        return power
-
-    if (exclude > 0):
-        k = k[exclude:]
-        excluded_power = np.copy(power[:exclude])
-        power = power[exclude:]
-
-    if (smooth_length is None) or (smooth_length < 0):
-        smooth_length = k[1]-k[0]
-
-    nmirror = int(5*smooth_length/(k[1]-k[0]))+2
-	
-    print "nmirror", nmirror
-
-    mpower = np.r_[np.exp(2*np.log(power[0])-np.log(power[1:nmirror][::-1])),power,np.exp(2*np.log(power[-1])-np.log(power[-nmirror:-1][::-1]))]
-
-
-    mk = np.r_[(2*k[0]-k[1:nmirror][::-1]),k,(2*k[-1]-k[-nmirror:-1][::-1])]
-    mdk = np.r_[0.5*(mk[1]-mk[0]),0.5*(mk[2:]-mk[:-2]),0.5*(mk[-1]-mk[-2])]
-    print "mpower", mpower
-    p_smooth = np.empty(mpower.shape)
-    print "p_smooth", p_smooth
-    for i in xrange(len(p_smooth)):
-        l = i-int(2*smooth_length/mdk[i])-1
-        l = max(l,0)
-        u = i+int(2*smooth_length/mdk[i])+2
-        u = min(u,len(p_smooth))
-        print "i", i, "l", l, "u", u
-        C = np.exp(-(mk[l:u]-mk[i])**2/(2.*smooth_length**2))*mdk[l:u]
-        p_smooth[i] = np.sum(C*mpower[l:u])/np.sum(C)
-        # print "p_smooth[",i,"] = ",p_smooth[i]
-
-    print "p_smooth2", " all ", p_smooth
-    p_smooth = p_smooth[nmirror - 1:-nmirror + 1]
-    print "p_smooth3", p_smooth
-    print "p_smooth length ", len(p_smooth)
-
-    if (exclude > 0):
-        p_smooth = np.r_[excluded_power,p_smooth]
-    return p_smooth
-
-def GaussianKernel(mpower, mk, mu,smooth_length):
-    C = np.exp(-(mk - mu) ** 2 / (2. * smooth_length ** 2))
-    return np.sum(C * mpower) / np.sum(C)
-
-
-def smoothie(power, startindex, endindex, k, exclude=1, smooth_length=None):
-
-    if smooth_length == 0:
-        # No smoothing requested, just return the input array.
-        return power
-
-    excluded_power = []
-    if (exclude > 0):
-        k = k[exclude:]
-        excluded_power = np.copy(power[:exclude])
-        power = power[exclude:]
-
-    if (smooth_length is None) or (smooth_length < 0):
-        smooth_length = k[1]-k[0]
-
-    p_smooth = np.empty(endindex-startindex)
-    for i in xrange(startindex, endindex):
-        l = max(i-int(2*smooth_length)-1,0)
-        u = min(i+int(2*smooth_length)+2,len(p_smooth))
-        p_smooth[i-startindex] = GaussianKernel(power[l:u], k[l:u], k[i],
-                                          smooth_length)
-
-    if (exclude > 0):
-        p_smooth = np.r_[excluded_power,p_smooth]
-    return p_smooth
-
-
-def smooth_something(datablock, axis=0, startindex=None, endindex=None,
-                     kernelfunction=lambda x:x, k=None, sigma=None):
-    if startindex == None:
-        startindex=0
-    if endindex == None:
-        endindex=len(datablock)
-    print kernelfunction
-    return np.apply_along_axis(smoothie, axis, datablock,
-                               startindex=startindex, endindex=endindex, k=k,
-                               smooth_length=sigma)
-
-    
\ No newline at end of file

nifty/operators/smoothing_operator/Cython/extended.pyx

+import numpy as np
+cimport numpy as np
+
+ctypedef fused FUSED_TYPES_t:
+    np.int32_t
+    np.int64_t
+    np.float32_t
+    np.float64_t
+    np.complex64_t
+    np.complex128_t
+
+cpdef np.ndarray[FUSED_TYPES_t, ndim=1] GaussianKernel(np.ndarray[FUSED_TYPES_t, ndim=1] mpower,np.ndarray[FUSED_TYPES_t, ndim=1] mk,np.ndarray[FUSED_TYPES_t, ndim=1] mu,np.float smooth_length):
+    cdef np.ndarray[FUSED_TYPES_t, ndim=1] C = np.exp(-(mk - mu) ** 2 / (2. * smooth_length ** 2))
+    return np.sum(C * mpower) / np.sum(C)
+
+
+cpdef np.ndarray[FUSED_TYPES_t, ndim=1] apply_kernel_along_array(np.ndarray[FUSED_TYPES_t, ndim=1] power, np.int startindex,np.int endindex,np.ndarray[FUSED_TYPES_t, ndim=1] k,np.int exclude=1,np.float smooth_length=None):
+
+    if smooth_length == 0:
+        # No smoothing requested, just return the input array.
+        return power
+
+    cdef np.ndarray[FUSED_TYPES_t, ndim=1] excluded_power = np.array([])
+    if (exclude > 0):
+        k = k[exclude:]
+        excluded_power = np.copy(power[:exclude])
+        power = power[exclude:]
+
+    if (smooth_length is None) or (smooth_length < 0):
+        smooth_length = k[1]-k[0]
+
+    cdef np.ndarray[FUSED_TYPES_t, ndim=1] p_smooth = np.empty(endindex-startindex)
+    cdef np.ndarray[FUSED_TYPES_t, ndim=1] l,u
+    cdef np.int i
+    for i in xrange(startindex, endindex):
+        l = max(i-int(2*smooth_length)-1,0)
+        u = min(i+int(2*smooth_length)+2,len(p_smooth))
+        p_smooth[i-startindex] = GaussianKernel(power[l:u], k[l:u], k[i],
+                                          smooth_length)
+
+    if (exclude > 0):
+        p_smooth = np.r_[excluded_power,p_smooth]
+    return p_smooth
+
+
+cpdef np.ndarray[FUSED_TYPES_t] apply_along_axis(tuple axis,np.ndarray[FUSED_TYPES_t] arr, np.int startindex,np.int endindex,np.ndarray[FUSED_TYPES_t, ndim=1] distances,np.int exclude=1,np.float smooth_length=None):
+    cdef tuple nd = arr.ndim
+    if axis < 0:
+        axis += nd
+    if (axis >= nd):
+        raise ValueError("axis must be less than arr.ndim; axis=%d, rank=%d."
+            % (axis, nd))
+    cdef np.ndarray ind = np.zeros(nd-1)
+    cdef np.ndarray i = np.zeros(nd, 'O')
+    indlist = list(range(nd))
+    indlist.remove(axis)
+    i[axis] = slice(None, None)
+    cdef np.ndarray outshape = np.asarray(arr.shape).take(indlist)
+    i.put(indlist, ind)
+    cdef np.ndarray res = apply_kernel_along_array(arr[tuple(i.tolist())], startindex=startindex,endindex=endindex,k=distances,exclude=exclude,smooth_length=smooth_length)
+
+    cdef np.int Ntot = np.product(outshape)
+    holdshape = outshape
+    outshape = list(arr.shape)
+    outshape[axis] = len(res)
+    outarr = np.zeros(outshape, np.asarray(res).dtype)
+    outarr[tuple(i.tolist())] = res
+    cdef np.int k = 1
+    cdef np.int n
+    while k < Ntot:
+        # increment the index
+        ind[-1] += 1
+        n = -1
+        while (ind[n] >= holdshape[n]) and (n > (1-nd)):
+            ind[n-1] += 1
+            ind[n] = 0
+            n -= 1
+        i.put(indlist, ind)
+        res = apply_kernel_along_array(arr[tuple(i.tolist())], startindex=startindex,endindex=endindex,k=distances, exclude=exclude,smooth_length=smooth_length)
+        outarr[tuple(i.tolist())] = res
+        k += 1
+    return outarr
\ No newline at end of file

nifty/operators/smoothing_operator/Cython/util.py

-import numpy as np
-
-def apply_along_axis(func1d, axis, arr, *args, **kwargs):
-    """
-    Apply a function to 1-D slices along the given axis.
-
-    Execute `func1d(a, *args)` where `func1d` operates on 1-D arrays and `a`
-    is a 1-D slice of `arr` along `axis`.
-
-    Parameters
-    ----------
-    func1d : function
-        This function should accept 1-D arrays. It is applied to 1-D
-        slices of `arr` along the specified axis.
-    axis : integer
-        Axis along which `arr` is sliced.
-    arr : ndarray
-        Input array.
-    args : any
-        Additional arguments to `func1d`.
-    kwargs: any
-        Additional named arguments to `func1d`.
-
-        .. versionadded:: 1.9.0
-
-
-    Returns
-    -------
-    apply_along_axis : ndarray
-        The output array. The shape of `outarr` is identical to the shape of
-        `arr`, except along the `axis` dimension, where the length of `outarr`
-        is equal to the size of the return value of `func1d`.  If `func1d`
-        returns a scalar `outarr` will have one fewer dimensions than `arr`.
-    """
-    arr = np.asarray(arr)
-    nd = arr.ndim
-    if axis < 0:
-        axis += nd
-    if (axis >= nd):
-        raise ValueError("axis must be less than arr.ndim; axis=%d, rank=%d."
-            % (axis, nd))
-    ind = [0]*(nd-1)
-    i = np.zeros(nd, 'O')
-    indlist = list(range(nd))
-    indlist.remove(axis)
-    i[axis] = slice(None, None)
-    outshape = np.asarray(arr.shape).take(indlist)
-    i.put(indlist, ind)
-    res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
-
-    Ntot = np.product(outshape)
-    holdshape = outshape
-    outshape = list(arr.shape)
-    outshape[axis] = len(res)
-    outarr = np.zeros(outshape, np.asarray(res).dtype)
-    outarr[tuple(i.tolist())] = res
-    k = 1
-    while k < Ntot:
-        # increment the index
-        ind[-1] += 1
-        n = -1
-        while (ind[n] >= holdshape[n]) and (n > (1-nd)):
-            ind[n-1] += 1
-            ind[n] = 0
-            n -= 1
-        i.put(indlist, ind)
-        res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
-        outarr[tuple(i.tolist())] = res
-        k += 1
-    return outarr
\ No newline at end of file