smoothing_operator.py 14.5 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

27
class SmoothingOperator(EndomorphicOperator):
Theo Steininger's avatar
Theo Steininger committed
28
    """ NIFTY class for smoothing operators.
29
30
31

    The NIFTy SmoothingOperator smooths Fields, with a given kernel length.
    Fields which are not living over a PowerSpace are smoothed
Theo Steininger's avatar
Theo Steininger committed
32
33
    via a gaussian convolution. Fields living over the PowerSpace are directly
    smoothed.
34
35
36

    Parameters
    ----------
Theo Steininger's avatar
Theo Steininger committed
37
    domain : tuple of DomainObjects, i.e. Spaces and FieldTypes
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
        The Space on which the operator acts
    sigma : float
        Sets the length of the Gaussian convolution kernel
    log_distances : boolean
        States whether the convolution happens on the logarithmic grid or not.

    Attributes
    ----------
    sigma : float
        Sets the length of the Gaussian convolution kernel
    log_distances : boolean
        States whether the convolution happens on the logarithmic grid or not.

    Raises
    ------
    ValueError
        Raised if
            * the given domain inherits more than one space. The
              SmoothingOperator acts only on one Space.

    Notes
    -----

    Examples
    --------
    >>> x = RGSpace(5)
    >>> S = SmoothingOperator(x, sigma=1.)
    >>> f = Field(x, val=[1,2,3,4,5])
    >>> S.times(f).val
    <distributed_data_object>
    array([ 3.,  3.,  3.,  3.,  3.])

    See Also
    --------
    DiagonalOperator, SmoothingOperator,
    PropagatorOperator, ProjectionOperator,
    ComposedOperator

    """

Jait Dixit's avatar
Jait Dixit committed
78
    # ---Overwritten properties and methods---
Martin Reinecke's avatar
Martin Reinecke committed
79
    def __init__(self, domain, sigma, log_distances=False,
80
81
                 default_spaces=None):
        super(SmoothingOperator, self).__init__(default_spaces)
82
83

        self._domain = self._parse_domain(domain)
Jait Dixit's avatar
Jait Dixit committed
84
85
86

        if len(self.domain) != 1:
            raise ValueError(
87
88
                'ERROR: SmoothOperator accepts only exactly one '
                'space as input domain.'
Jait Dixit's avatar
Jait Dixit committed
89
            )
Jait Dixit's avatar
Jait Dixit committed
90

91
92
93
        self.sigma = sigma
        self.log_distances = log_distances

94
        self._direct_smoothing_width = 3.
Jait Dixit's avatar
Jait Dixit committed
95

96
    def _inverse_times(self, x, spaces):
97
        return self._smoothing_helper(x, spaces, inverse=True)
Jait Dixit's avatar
Jait Dixit committed
98

99
    def _times(self, x, spaces):
100
        return self._smoothing_helper(x, spaces, inverse=False)
Jait Dixit's avatar
Jait Dixit committed
101

Jait Dixit's avatar
Jait Dixit committed
102
    # ---Mandatory properties and methods---
103
104
105
106
    @property
    def domain(self):
        return self._domain

Jait Dixit's avatar
Jait Dixit committed
107
    @property
Martin Reinecke's avatar
Martin Reinecke committed
108
    def self_adjoint(self):
theos's avatar
theos committed
109
        return True
Jait Dixit's avatar
Jait Dixit committed
110

Jait Dixit's avatar
Jait Dixit committed
111
112
113
    @property
    def unitary(self):
        return False
Jait Dixit's avatar
Jait Dixit committed
114
115

    # ---Added properties and methods---
116

Jait Dixit's avatar
Jait Dixit committed
117
118
119
120
    @property
    def sigma(self):
        return self._sigma

121
122
123
124
125
126
127
128
129
130
131
132
    @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)

133
    @staticmethod
134
    def _precompute(x, sigma, dxmax=None):
135
136
137
138
        """ Performs the precomputations for Gaussian smoothing on a 1D irregular grid.

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

142
143
144
145
        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.
146
147
148
149
150
151
152
153
154
155
156
157
158
159

        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.
        """

160
161
        if dxmax is None:
            dxmax = 3.01*sigma
162

163
        x = np.asarray(x)
164

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

168
169
        wgt = []
        expfac = 1. / (2. * sigma*sigma)
170
        for i in range(x.size):
171
172
173
            t = x[ibegin[i]:ibegin[i]+nval[i]]-x[i]
            t = np.exp(-t*t*expfac)
            t *= 1./np.sum(t)
174
175
            wgt.append(t)

176
        return ibegin, nval, wgt
177
178

    @staticmethod
179
180
181
    def _apply_kernel_along_array(
            power, startindex, endindex, distances, smooth_length,
            smoothing_width, ibegin, nval, wgt):
182
183
184
185

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

Martin Reinecke's avatar
Martin Reinecke committed
186
        p_smooth = np.zeros(endindex-startindex, dtype=power.dtype)
187
        for i in xrange(startindex, endindex):
188
189
190
191
            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]])
192
193
194
195

        return p_smooth

    @staticmethod
Martin Reinecke's avatar
Martin Reinecke committed
196
    def _getShape(a):
197
198
199
        return tuple(a.shape)

    @staticmethod
Martin Reinecke's avatar
Martin Reinecke committed
200
    def _apply_along_axis(axis, arr, startindex, endindex, distances,
201
                          smooth_length, smoothing_width):
202
203
204
205
206

        nd = arr.ndim
        if axis < 0:
            axis += nd
        if (axis >= nd):
207
208
            raise ValueError(
                "axis must be less than arr.ndim; axis=%d, rank=%d."
209
                % (axis, nd))
210
211
        ibegin, nval, wgt = SmoothingOperator._precompute(
                distances, smooth_length, smooth_length*smoothing_width)
212
213
214

        ind = np.zeros(nd-1, dtype=np.int)
        i = np.zeros(nd, dtype=object)
Martin Reinecke's avatar
Martin Reinecke committed
215
        shape = SmoothingOperator._getShape(arr)
216
        indlist = np.asarray(range(nd))
217
        indlist = np.delete(indlist, axis)
218
219
220
221
222
        i[axis] = slice(None, None)
        outshape = np.asarray(shape).take(indlist)

        i.put(indlist, ind)

223
        Ntot = np.product(outshape)
224
225
        holdshape = outshape
        slicedArr = arr[tuple(i.tolist())]
226
227
228
        res = SmoothingOperator._apply_kernel_along_array(
                    slicedArr, startindex, endindex, distances,
                    smooth_length, smoothing_width, ibegin, nval, wgt)
229

Martin Reinecke's avatar
Martin Reinecke committed
230
        outshape = np.asarray(SmoothingOperator._getShape(arr))
231
        outshape[axis] = endindex - startindex
Martin Reinecke's avatar
updates    
Martin Reinecke committed
232
        outarr = np.zeros(outshape, dtype=arr.dtype)
233
234
235
236
237
238
239
240
241
242
243
244
        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())]
245
246
247
            res = SmoothingOperator._apply_kernel_along_array(
                    slicedArr, startindex, endindex, distances,
                    smooth_length, smoothing_width, ibegin, nval, wgt)
248
249
250
251
252
            outarr[tuple(i.tolist())] = res
            k += 1

        return outarr

253
    def _smoothing_helper(self, x, spaces, inverse):
theos's avatar
theos committed
254
255
256
257
258
259
260
261
262
263
        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
264
265
            spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain))

266
267
268
269
270
271
272
        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
273
        Transformator = FFTOperator(x.domain[spaces[0]])
Jait Dixit's avatar
Jait Dixit committed
274

theos's avatar
theos committed
275
276
277
278
279
        # 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]]
280

theos's avatar
theos committed
281
282
283
        # 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
284

285
        kernel = codomain.get_distance_array(
286
287
288
289
290
            distribution_strategy=axes_local_distribution_strategy)

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

theos's avatar
theos committed
291
        kernel.apply_scalar_function(
292
            codomain.get_fft_smoothing_kernel_function(self.sigma),
theos's avatar
theos committed
293
            inplace=True)
Jait Dixit's avatar
Jait Dixit committed
294

theos's avatar
theos committed
295
296
297
298
299
        # 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
300

301
        reshaper = [transformed_x.shape[i] if i in coaxes else 1
theos's avatar
theos committed
302
303
                    for i in xrange(len(transformed_x.shape))]
        local_kernel = np.reshape(local_kernel, reshaper)
Jait Dixit's avatar
Jait Dixit committed
304

theos's avatar
theos committed
305
306
307
308
309
        # apply the kernel
        if inverse:
            local_transformed_x /= local_kernel
        else:
            local_transformed_x *= local_kernel
Jait Dixit's avatar
Jait Dixit committed
310

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

313
        smoothed_x = Transformator.adjoint_times(transformed_x, spaces=spaces)
theos's avatar
theos committed
314
315
316

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

theos's avatar
theos committed
318
        return result
319
320
321
322
323
324

    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]]

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

        affected_axis = affected_axes[0]
330
331

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

        if self.log_distances:
336
            np.log(distance_array, out=distance_array)
337
338
339
340
341
342
343
344
345

        # 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.
346
            if 0 != affected_axis:
347
348
                local_data_Q = True
            else:
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
                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)

366
367
368
369
370
                augmented_data = x.val.get_data(augmented_slice,
                                                local_keys=True,
                                                copy=False)
                augmented_data = augmented_data.get_local_data(copy=False)

371
                augmented_distance_array = distance_array[augmented_slice]
372
373

        else:
374
375
            raise ValueError("Direct smoothing not implemented for given"
                             "distribution strategy.")
376
377
378
379
380

        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)
381
382
383
            augmented_distance_array = distance_array
            true_start = 0
            true_end = x.shape[affected_axis]
384
385

        # perform the convolution along the affected axes
386
387
388
389
390
391
392
393
394
        # 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)
395
396
397
398
399
        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,
400
                                      true_start, true_end, inverse):
401
        if inverse:
402
            true_sigma = 1. / self.sigma
403
404
405
        else:
            true_sigma = self.sigma

Martin Reinecke's avatar
Martin Reinecke committed
406
        smoothed_data = SmoothingOperator._apply_along_axis(
407
408
409
410
411
412
                              data_axis, data,
                              startindex=true_start,
                              endindex=true_end,
                              distances=distances,
                              smooth_length=true_sigma,
                              smoothing_width=self._direct_smoothing_width)
413

414
        return smoothed_data