direct_smoothing_operator.py 8.48 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
# -*- coding: utf8 -*-

import numpy as np

from d2o import STRATEGIES

from .smoothing_operator import SmoothingOperator


class DirectSmoothingOperator(SmoothingOperator):
    def __init__(self, domain, sigma, log_distances=False,
                 default_spaces=None):
        super(DirectSmoothingOperator, self).__init__(domain,
                                                      sigma,
                                                      log_distances,
                                                      default_spaces)
        self.effective_smoothing_width = 3.01

    def _precompute(self, x, sigma, dxmax=None):
        """ Does precomputations for Gaussian smoothing on a 1D irregular grid.

        Parameters
        ----------
        x: 1D floating point array or list containing the individual grid
            positions. Points must be given in ascending order.

        sigma: The sigma of the Gaussian with which the function living on x
            should be smoothed, in the same units as x.
        dxmax: (optional) The maximum distance up to which smoothing is
            performed, in the same units as x. Default is 3.01*sigma.

        Returns
        -------
        ibegin: integer array of the same size as x
            ibegin[i] is the minimum grid index to consider when computing the
            smoothed value at grid index i
        nval: integer array of the same size as x
            nval[i] is the number of indices to consider when computing the
            smoothed value at grid index i.
        wgt: list with the same number of entries as x
            wgt[i] is an array with nval[i] entries containing the
            normalized smoothing weights.
        """

        if dxmax is None:
            dxmax = self.effective_smoothing_width*sigma

        x = np.asarray(x)

        ibegin = np.searchsorted(x, x-dxmax)
        nval = np.searchsorted(x, x+dxmax) - ibegin

        wgt = []
        expfac = 1. / (2. * sigma*sigma)
        for i in range(x.size):
56
57
58
59
60
61
62
            if nval[i]>0:
                t = x[ibegin[i]:ibegin[i]+nval[i]]-x[i]
                t = np.exp(-t*t*expfac)
                t *= 1./np.sum(t)
                wgt.append(t)
            else:
                wgt.append(np.array([]))
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

        return ibegin, nval, wgt

    def _apply_kernel_along_array(self, power, startindex, endindex,
                                  distances, smooth_length, smoothing_width,
                                  ibegin, nval, wgt):

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

        p_smooth = np.zeros(endindex-startindex, dtype=power.dtype)
        for i in xrange(startindex, endindex):
            imin = max(startindex, ibegin[i])
            imax = min(endindex, ibegin[i]+nval[i])
            p_smooth[imin:imax] += (power[i] *
                                    wgt[i][imin-ibegin[i]:imax-imin+ibegin[i]])

        return p_smooth

    def _apply_along_axis(self, axis, arr, startindex, endindex, distances,
                          smooth_length, smoothing_width):

        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))
        ibegin, nval, wgt = self._precompute(
                distances, smooth_length, smooth_length*smoothing_width)

        ind = np.zeros(nd-1, dtype=np.int)
        i = np.zeros(nd, dtype=object)
        shape = arr.shape
        indlist = np.asarray(range(nd))
        indlist = np.delete(indlist, axis)
        i[axis] = slice(None, None)
        outshape = np.asarray(shape).take(indlist)

        i.put(indlist, ind)

        Ntot = np.product(outshape)
        holdshape = outshape
        slicedArr = arr[tuple(i.tolist())]
        res = self._apply_kernel_along_array(
                    slicedArr, startindex, endindex, distances,
                    smooth_length, smoothing_width, ibegin, nval, wgt)

        outshape = np.asarray(arr.shape)
        outshape[axis] = endindex - startindex
        outarr = np.zeros(outshape, dtype=arr.dtype)
        outarr[tuple(i.tolist())] = res
        k = 1
        while k < Ntot:
            # increment the index
            ind[nd-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)
            slicedArr = arr[tuple(i.tolist())]
            res = self._apply_kernel_along_array(
                    slicedArr, startindex, endindex, distances,
                    smooth_length, smoothing_width, ibegin, nval, wgt)
            outarr[tuple(i.tolist())] = res
            k += 1

        return outarr

    def _smooth(self, x, spaces, inverse):
        # infer affected axes
        # we rely on the knowledge, that `spaces` is a tuple with length 1.
        affected_axes = x.domain_axes[spaces[0]]

        if len(affected_axes) > 1:
            raise ValueError("By this implementation only one-dimensional "
                             "spaces can be smoothed directly.")

        affected_axis = affected_axes[0]

        distance_array = x.domain[spaces[0]].get_distance_array(
            distribution_strategy='not')
        distance_array = distance_array.get_local_data(copy=False)

Martin Reinecke's avatar
Martin Reinecke committed
150
        #MR FIXME: this causes calls of log(0.) which should probably be avoided
151
        if self.log_distances:
152
            np.log(np.maximum(distance_array,1e-15), out=distance_array)
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222

        # collect the local data + ghost cells
        local_data_Q = False

        if x.distribution_strategy == 'not':
            local_data_Q = True
        elif x.distribution_strategy in STRATEGIES['slicing']:
            # infer the local start/end based on the slicing information of
            # x's d2o. Only gets non-trivial for axis==0.
            if 0 != affected_axis:
                local_data_Q = True
            else:
                start_index = x.val.distributor.local_start
                start_distance = distance_array[start_index]
                augmented_start_distance = \
                    (start_distance -
                     self.effective_smoothing_width*self.sigma)
                augmented_start_index = \
                    np.searchsorted(distance_array, augmented_start_distance)
                true_start = start_index - augmented_start_index
                end_index = x.val.distributor.local_end
                end_distance = distance_array[end_index-1]
                augmented_end_distance = \
                    (end_distance + self.effective_smoothing_width*self.sigma)
                augmented_end_index = \
                    np.searchsorted(distance_array, augmented_end_distance)
                true_end = true_start + x.val.distributor.local_length
                augmented_slice = slice(augmented_start_index,
                                        augmented_end_index)

                augmented_data = x.val.get_data(augmented_slice,
                                                local_keys=True,
                                                copy=False)
                augmented_data = augmented_data.get_local_data(copy=False)

                augmented_distance_array = distance_array[augmented_slice]

        else:
            raise ValueError("Direct smoothing not implemented for given"
                             "distribution strategy.")

        if local_data_Q:
            # if the needed data resides on the nodes already, the necessary
            # are the same; no matter what the distribution strategy was.
            augmented_data = x.val.get_local_data(copy=False)
            augmented_distance_array = distance_array
            true_start = 0
            true_end = x.shape[affected_axis]

        # perform the convolution along the affected axes
        # currently only one axis is supported
        data_axis = affected_axes[0]

        if inverse:
            true_sigma = 1. / self.sigma
        else:
            true_sigma = self.sigma

        local_result = self._apply_along_axis(
                              data_axis,
                              augmented_data,
                              startindex=true_start,
                              endindex=true_end,
                              distances=augmented_distance_array,
                              smooth_length=true_sigma,
                              smoothing_width=self.effective_smoothing_width)

        result = x.copy_empty()
        result.val.set_local_data(local_result, copy=False)
        return result