smoothing_operator.py 13.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# NIFTy
# Copyright (C) 2017  Theo Steininger
#
# Author: Theo Steininger
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

Jait Dixit's avatar
Jait Dixit committed
19
20
21
22
23
import numpy as np

import nifty.nifty_utilities as utilities
from nifty.operators.endomorphic_operator import EndomorphicOperator
from nifty.operators.fft_operator import FFTOperator
24
from d2o import STRATEGIES
Jait Dixit's avatar
Jait Dixit committed
25

26
class SmoothingOperator(EndomorphicOperator):
Jait Dixit's avatar
Jait Dixit committed
27
    # ---Overwritten properties and methods---
Martin Reinecke's avatar
Martin Reinecke committed
28
    def __init__(self, domain, sigma, log_distances=False,
29
30
                 default_spaces=None):
        super(SmoothingOperator, self).__init__(default_spaces)
31
32

        self._domain = self._parse_domain(domain)
Jait Dixit's avatar
Jait Dixit committed
33
34
35

        if len(self.domain) != 1:
            raise ValueError(
36
37
                'ERROR: SmoothOperator accepts only exactly one '
                'space as input domain.'
Jait Dixit's avatar
Jait Dixit committed
38
            )
Jait Dixit's avatar
Jait Dixit committed
39

40
41
42
        self.sigma = sigma
        self.log_distances = log_distances

43
        self._direct_smoothing_width = 3.
Jait Dixit's avatar
Jait Dixit committed
44

45
    def _inverse_times(self, x, spaces):
46
        return self._smoothing_helper(x, spaces, inverse=True)
Jait Dixit's avatar
Jait Dixit committed
47

48
    def _times(self, x, spaces):
49
        return self._smoothing_helper(x, spaces, inverse=False)
Jait Dixit's avatar
Jait Dixit committed
50

Jait Dixit's avatar
Jait Dixit committed
51
    # ---Mandatory properties and methods---
52
53
54
55
    @property
    def domain(self):
        return self._domain

Jait Dixit's avatar
Jait Dixit committed
56
    @property
Martin Reinecke's avatar
Martin Reinecke committed
57
    def self_adjoint(self):
theos's avatar
theos committed
58
        return True
Jait Dixit's avatar
Jait Dixit committed
59

Jait Dixit's avatar
Jait Dixit committed
60
61
62
    @property
    def unitary(self):
        return False
Jait Dixit's avatar
Jait Dixit committed
63
64

    # ---Added properties and methods---
65

Jait Dixit's avatar
Jait Dixit committed
66
67
68
69
    @property
    def sigma(self):
        return self._sigma

70
71
72
73
74
75
76
77
78
79
80
81
    @sigma.setter
    def sigma(self, sigma):
        self._sigma = np.float(sigma)

    @property
    def log_distances(self):
        return self._log_distances

    @log_distances.setter
    def log_distances(self, log_distances):
        self._log_distances = bool(log_distances)

82
    @staticmethod
Martin Reinecke's avatar
Martin Reinecke committed
83
    def _precompute (x,sigma,dxmax=None):
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
        """ Performs the 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==None:
            dxmax=3.01*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):
            t=x[ibegin[i]:ibegin[i]+nval[i]]-x[i]
            t=np.exp(-t*t*expfac)
            t*=1./np.sum(t)
            wgt.append(t)

        return ibegin,nval,wgt

    @staticmethod
Martin Reinecke's avatar
Martin Reinecke committed
128
    def _apply_kernel_along_array(power, startindex, endindex, distances,
129
130
131
132
133
        smooth_length, smoothing_width,ibegin,nval,wgt):

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

Martin Reinecke's avatar
Martin Reinecke committed
134
        p_smooth = np.zeros(endindex-startindex, dtype=power.dtype)
135
        for i in xrange(startindex, endindex):
Martin Reinecke's avatar
Martin Reinecke committed
136
137
138
            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]]
139
140
141
142

        return p_smooth

    @staticmethod
Martin Reinecke's avatar
Martin Reinecke committed
143
    def _getShape(a):
144
145
146
        return tuple(a.shape)

    @staticmethod
Martin Reinecke's avatar
Martin Reinecke committed
147
    def _apply_along_axis(axis, arr, startindex, endindex, distances,
148
149
150
151
152
153
154
155
        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))
Martin Reinecke's avatar
Martin Reinecke committed
156
        ibegin,nval,wgt=SmoothingOperator._precompute(distances,smooth_length,smooth_length*smoothing_width)
157
158
159

        ind = np.zeros(nd-1, dtype=np.int)
        i = np.zeros(nd, dtype=object)
Martin Reinecke's avatar
Martin Reinecke committed
160
        shape = SmoothingOperator._getShape(arr)
161
162
163
164
165
166
167
168
169
170
        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())]
Martin Reinecke's avatar
Martin Reinecke committed
171
        res = SmoothingOperator._apply_kernel_along_array(slicedArr,
172
173
174
175
176
177
                                       startindex,
                                       endindex,
                                       distances,
                                       smooth_length,
                                       smoothing_width, ibegin,nval,wgt)

Martin Reinecke's avatar
Martin Reinecke committed
178
        outshape = np.asarray(SmoothingOperator._getShape(arr))
179
        outshape[axis] = endindex - startindex
Martin Reinecke's avatar
updates    
Martin Reinecke committed
180
        outarr = np.zeros(outshape, dtype=arr.dtype)
181
182
183
184
185
186
187
188
189
190
191
192
        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())]
Martin Reinecke's avatar
Martin Reinecke committed
193
            res = SmoothingOperator._apply_kernel_along_array(slicedArr,
194
195
196
197
198
199
200
201
202
203
                                           startindex,
                                           endindex,
                                           distances,
                                           smooth_length,
                                           smoothing_width,ibegin,nval,wgt)
            outarr[tuple(i.tolist())] = res
            k += 1

        return outarr

204
    def _smoothing_helper(self, x, spaces, inverse):
theos's avatar
theos committed
205
206
207
208
209
210
211
212
213
214
        if self.sigma == 0:
            return x.copy()

        # the domain of the smoothing operator contains exactly one space.
        # Hence, if spaces is None, but we passed LinearOperator's
        # _check_input_compatibility, we know that x is also solely defined
        # on that space
        if spaces is None:
            spaces = (0,)
        else:
Jait Dixit's avatar
Jait Dixit committed
215
216
            spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain))

217
218
219
220
221
222
223
        try:
            result = self._fft_smoothing(x, spaces, inverse)
        except ValueError:
            result = self._direct_smoothing(x, spaces, inverse)
        return result

    def _fft_smoothing(self, x, spaces, inverse):
theos's avatar
theos committed
224
        Transformator = FFTOperator(x.domain[spaces[0]])
Jait Dixit's avatar
Jait Dixit committed
225

theos's avatar
theos committed
226
227
228
229
230
        # transform to the (global-)default codomain and perform all remaining
        # steps therein
        transformed_x = Transformator(x, spaces=spaces)
        codomain = transformed_x.domain[spaces[0]]
        coaxes = transformed_x.domain_axes[spaces[0]]
231

theos's avatar
theos committed
232
233
234
        # create the kernel using the knowledge of codomain about itself
        axes_local_distribution_strategy = \
            transformed_x.val.get_axes_local_distribution_strategy(axes=coaxes)
Jait Dixit's avatar
Jait Dixit committed
235

236
        kernel = codomain.get_distance_array(
237
238
239
240
241
            distribution_strategy=axes_local_distribution_strategy)

        if self.log_distances:
            kernel.apply_scalar_function(np.log, inplace=True)

theos's avatar
theos committed
242
        kernel.apply_scalar_function(
243
            codomain.get_fft_smoothing_kernel_function(self.sigma),
theos's avatar
theos committed
244
            inplace=True)
Jait Dixit's avatar
Jait Dixit committed
245

theos's avatar
theos committed
246
247
248
249
250
        # now, apply the kernel to transformed_x
        # this is done node-locally utilizing numpys reshaping in order to
        # apply the kernel to the correct axes
        local_transformed_x = transformed_x.val.get_local_data(copy=False)
        local_kernel = kernel.get_local_data(copy=False)
Jait Dixit's avatar
Jait Dixit committed
251

252
        reshaper = [transformed_x.shape[i] if i in coaxes else 1
theos's avatar
theos committed
253
254
                    for i in xrange(len(transformed_x.shape))]
        local_kernel = np.reshape(local_kernel, reshaper)
Jait Dixit's avatar
Jait Dixit committed
255

theos's avatar
theos committed
256
257
258
259
260
        # apply the kernel
        if inverse:
            local_transformed_x /= local_kernel
        else:
            local_transformed_x *= local_kernel
Jait Dixit's avatar
Jait Dixit committed
261

theos's avatar
theos committed
262
        transformed_x.val.set_local_data(local_transformed_x, copy=False)
Jait Dixit's avatar
Jait Dixit committed
263

theos's avatar
theos committed
264
265
266
267
        smoothed_x = Transformator.inverse_times(transformed_x, spaces=spaces)

        result = x.copy_empty()
        result.set_val(smoothed_x, copy=False)
Jait Dixit's avatar
Jait Dixit committed
268

theos's avatar
theos committed
269
        return result
270
271
272
273
274
275

    def _direct_smoothing(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]]

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

        affected_axis = affected_axes[0]
281
282

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

        if self.log_distances:
287
            np.log(distance_array, out=distance_array)
288
289
290
291
292
293
294
295
296

        # 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.
297
            if 0 != affected_axis:
298
299
                local_data_Q = True
            else:
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
                start_index = x.val.distributor.local_start
                start_distance = distance_array[start_index]
                augmented_start_distance = \
                    (start_distance - self._direct_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._direct_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)

317
318
319
320
321
                augmented_data = x.val.get_data(augmented_slice,
                                                local_keys=True,
                                                copy=False)
                augmented_data = augmented_data.get_local_data(copy=False)

322
                augmented_distance_array = distance_array[augmented_slice]
323
324

        else:
325
326
            raise ValueError("Direct smoothing not implemented for given"
                             "distribution strategy.")
327
328
329
330
331

        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)
332
333
334
            augmented_distance_array = distance_array
            true_start = 0
            true_end = x.shape[affected_axis]
335
336

        # perform the convolution along the affected axes
337
338
339
340
341
342
343
344
345
        # currently only one axis is supported
        data_axis = affected_axes[0]
        local_result = self._direct_smoothing_single_axis(
                                                    augmented_data,
                                                    data_axis,
                                                    augmented_distance_array,
                                                    true_start,
                                                    true_end,
                                                    inverse)
346
347
348
349
350
        result = x.copy_empty()
        result.val.set_local_data(local_result, copy=False)
        return result

    def _direct_smoothing_single_axis(self, data, data_axis, distances,
351
                                      true_start, true_end, inverse):
352
        if inverse:
353
            true_sigma = 1. / self.sigma
354
355
356
        else:
            true_sigma = self.sigma

Martin Reinecke's avatar
Martin Reinecke committed
357
        smoothed_data = SmoothingOperator._apply_along_axis(
358
359
360
361
362
363
                              data_axis, data,
                              startindex=true_start,
                              endindex=true_end,
                              distances=distances,
                              smooth_length=true_sigma,
                              smoothing_width=self._direct_smoothing_width)
364

365
        return smoothed_data