smooth_util.pyx 8.63 KB
Newer Older
1 2 3 4 5
#cython: nonecheck=False
#cython: boundscheck=True
#cython: wraparound=False
#cython: cdivision=True

csongor's avatar
csongor committed
6 7
import numpy as np
cimport numpy as np
csongor's avatar
csongor committed
8
#defining NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
csongor's avatar
csongor committed
9

10 11 12 13 14 15 16 17
cdef int SIGMA_RANGE = 3

cpdef np.float32_t gaussianKernel_f(np.ndarray[np.float32_t, ndim=1] mpower,
                                    np.ndarray[np.float32_t, ndim=1] mk,
                                    np.float32_t mu,
                                    np.float32_t smooth_length):
    cdef np.ndarray[np.float32_t, ndim=1] C = \
        np.exp(-(mk - mu) ** 2 / (2. * smooth_length ** 2))
csongor's avatar
csongor committed
18
    return np.sum(C * mpower) / np.sum(C)
csongor's avatar
csongor committed
19

20 21 22 23 24 25
cpdef np.float64_t gaussianKernel(np.ndarray[np.float64_t, ndim=1] mpower,
                                  np.ndarray[np.float64_t, ndim=1] mk,
                                  np.float64_t mu,
                                  np.float64_t smooth_length):
    cdef np.ndarray[np.float64_t, ndim=1] C = \
        np.exp(-(mk - mu) ** 2 / (2. * smooth_length ** 2))
csongor's avatar
csongor committed
26 27 28
    return np.sum(C * mpower) / np.sum(C)


29 30 31 32 33 34 35
cpdef np.ndarray[np.float64_t, ndim=1] apply_kernel_along_array(
                                    np.ndarray[np.float64_t, ndim=1] power,
                                    np.int_t startindex,
                                    np.int_t endindex,
                                    np.ndarray[np.float64_t, ndim=1] distances,
                                    np.float64_t smooth_length,
                                    np.float64_t smoothing_width):
csongor's avatar
csongor committed
36

csongor's avatar
csongor committed
37 38
    if smooth_length == 0.0:
        return power[startindex:endindex]
csongor's avatar
csongor committed
39 40

    if (smooth_length is None) or (smooth_length < 0):
csongor's avatar
csongor committed
41
        smooth_length = distances[1]-distances[0]
csongor's avatar
csongor committed
42

43 44 45
    cdef np.ndarray[np.float64_t, ndim=1] p_smooth = \
        np.empty(endindex-startindex, dtype=np.float64)
    cdef np.int64_t i, l, u
csongor's avatar
csongor committed
46
    for i in xrange(startindex, endindex):
47 48 49 50 51 52 53 54 55
        current_location = distances[i]
        l = np.searchsorted(distances,
                            current_location - smoothing_width*smooth_length)
        u = np.searchsorted(distances,
                            current_location + smoothing_width*smooth_length)
        p_smooth[i-startindex] = gaussianKernel(power[l:u],
                                                distances[l:u],
                                                distances[i],
                                                smooth_length)
csongor's avatar
csongor committed
56 57 58

    return p_smooth

59 60 61 62 63 64 65
cpdef np.ndarray[np.float32_t, ndim=1] apply_kernel_along_array_f(
                                    np.ndarray[np.float32_t, ndim=1] power,
                                    np.int_t startindex,
                                    np.int_t endindex,
                                    np.ndarray[np.float32_t, ndim=1] distances,
                                    np.float64_t smooth_length,
                                    np.float64_t smoothing_width):
csongor's avatar
csongor committed
66 67 68 69 70 71 72

    if smooth_length == 0.0:
        return power[startindex:endindex]

    if (smooth_length is None) or (smooth_length < 0):
        smooth_length = distances[1]-distances[0]

73 74 75
    cdef np.ndarray[np.float32_t, ndim=1] p_smooth = \
        np.empty(endindex-startindex, dtype=np.float64)
    cdef np.int64_t i, l, u
csongor's avatar
csongor committed
76
    for i in xrange(startindex, endindex):
77 78 79 80 81 82 83 84 85
        current_location = distances[i]
        l = np.searchsorted(distances,
                            current_location - smoothing_width*smooth_length)
        u = np.searchsorted(distances,
                            current_location + smoothing_width*smooth_length)
        p_smooth[i-startindex] = gaussianKernel(power[l:u],
                                                distances[l:u],
                                                distances[i],
                                                smooth_length)
csongor's avatar
csongor committed
86

csongor's avatar
csongor committed
87 88
    return p_smooth

csongor's avatar
csongor committed
89 90
def getShape(a):
    return tuple(a.shape)
csongor's avatar
csongor committed
91

92 93 94 95 96 97 98 99
cpdef np.ndarray[np.float64_t, ndim=1] apply_along_axis(
                                    np.int_t axis,
                                    np.ndarray arr,
                                    np.int_t startindex,
                                    np.int_t endindex,
                                    np.ndarray[np.float64_t, ndim=1] distances,
                                    np.float64_t smooth_length,
                                    np.float64_t smoothing_width):
csongor's avatar
csongor committed
100
    cdef np.int_t nd = arr.ndim
csongor's avatar
csongor committed
101 102 103 104 105
    if axis < 0:
        axis += nd
    if (axis >= nd):
        raise ValueError("axis must be less than arr.ndim; axis=%d, rank=%d."
            % (axis, nd))
csongor's avatar
csongor committed
106 107
    cdef np.ndarray[np.int_t, ndim=1] ind = np.zeros(nd-1, dtype=np.int)
    cdef np.ndarray i = np.zeros(nd, dtype=object)
csongor's avatar
csongor committed
108
    cdef tuple shape = getShape(arr)
csongor's avatar
csongor committed
109 110
    cdef np.ndarray[np.int_t, ndim=1] indlist = np.asarray(range(nd))
    indlist = np.delete(indlist,axis)
csongor's avatar
csongor committed
111
    i[axis] = slice(None, None)
112 113
    cdef np.ndarray[np.int_t, ndim=1] outshape = \
        np.asarray(shape).take(indlist)
csongor's avatar
csongor committed
114

csongor's avatar
csongor committed
115 116
    i.put(indlist, ind)

csongor's avatar
csongor committed
117 118 119 120 121 122
    cdef np.ndarray[np.float64_t, ndim=1] slicedArr
    cdef np.ndarray[np.float64_t, ndim=1] res

    cdef np.int_t Ntot = np.product(outshape)
    cdef np.ndarray[np.int_t, ndim=1] holdshape = outshape
    slicedArr = arr[tuple(i.tolist())]
csongor's avatar
csongor committed
123

124 125 126 127 128 129
    res = apply_kernel_along_array(slicedArr,
                                   startindex,
                                   endindex,
                                   distances,
                                   smooth_length,
                                   smoothing_width)
csongor's avatar
csongor committed
130

csongor's avatar
csongor committed
131 132 133
    outshape = np.asarray(getShape(arr))
    outshape[axis] = endindex - startindex
    outarr = np.zeros(outshape, dtype=np.float64)
csongor's avatar
csongor committed
134
    outarr[tuple(i.tolist())] = res
csongor's avatar
csongor committed
135
    cdef np.int_t n, k = 1
csongor's avatar
csongor committed
136 137
    while k < Ntot:
        # increment the index
138
        ind[nd-1] += 1
csongor's avatar
csongor committed
139 140 141 142 143 144
        n = -1
        while (ind[n] >= holdshape[n]) and (n > (1-nd)):
            ind[n-1] += 1
            ind[n] = 0
            n -= 1
        i.put(indlist, ind)
csongor's avatar
csongor committed
145
        slicedArr = arr[tuple(i.tolist())]
146 147 148 149 150 151
        res = apply_kernel_along_array(slicedArr,
                                       startindex,
                                       endindex,
                                       distances,
                                       smooth_length,
                                       smoothing_width)
csongor's avatar
csongor committed
152 153 154 155 156 157
        outarr[tuple(i.tolist())] = res
        k += 1

    return outarr


158 159 160 161 162 163 164 165
cpdef np.ndarray[np.float32_t, ndim=1] apply_along_axis_f(
                                    np.int_t axis,
                                    np.ndarray arr,
                                    np.int_t startindex,
                                    np.int_t endindex,
                                    np.ndarray[np.float32_t, ndim=1] distances,
                                    np.float64_t smooth_length,
                                    np.float64_t smoothing_width):
csongor's avatar
csongor committed
166 167 168 169 170 171 172 173 174 175 176 177
    cdef np.int_t 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[np.int_t, ndim=1] ind = np.zeros(nd-1, dtype=np.int)
    cdef np.ndarray i = np.zeros(nd, dtype=object)
    cdef tuple shape = getShape(arr)
    cdef np.ndarray[np.int_t, ndim=1] indlist = np.asarray(range(nd))
    indlist = np.delete(indlist,axis)
    i[axis] = slice(None, None)
178 179
    cdef np.ndarray[np.int_t, ndim=1] outshape = \
        np.asarray(shape).take(indlist)
csongor's avatar
csongor committed
180 181 182 183 184 185 186 187 188 189

    i.put(indlist, ind)

    cdef np.ndarray[np.float32_t, ndim=1] slicedArr
    cdef np.ndarray[np.float32_t, ndim=1] res

    cdef np.int_t Ntot = np.product(outshape)
    cdef np.ndarray[np.int_t, ndim=1] holdshape = outshape
    slicedArr = arr[tuple(i.tolist())]

190 191 192 193 194 195
    res = apply_kernel_along_array_f(slicedArr,
                                     startindex,
                                     endindex,
                                     distances,
                                     smooth_length,
                                     smoothing_width)
csongor's avatar
csongor committed
196 197 198

    outshape = np.asarray(getShape(arr))
    outshape[axis] = endindex - startindex
199
    outarr = np.zeros(outshape, dtype=np.float32)
csongor's avatar
csongor committed
200 201 202 203
    outarr[tuple(i.tolist())] = res
    cdef np.int_t n, k = 1
    while k < Ntot:
        # increment the index
204
        ind[nd-1] += 1
csongor's avatar
csongor committed
205 206 207 208 209 210 211
        n = -1
        while (ind[n] >= holdshape[n]) and (n > (1-nd)):
            ind[n-1] += 1
            ind[n] = 0
            n -= 1
        i.put(indlist, ind)
        slicedArr = arr[tuple(i.tolist())]
212 213 214 215 216 217
        res = apply_kernel_along_array(slicedArr,
                                       startindex,
                                       endindex,
                                       distances,
                                       smooth_length,
                                       smoothing_width)
csongor's avatar
csongor committed
218 219
        outarr[tuple(i.tolist())] = res
        k += 1
csongor's avatar
csongor committed
220 221

    return outarr
csongor's avatar
csongor committed
222