smoothing_operator.py 11.9 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
Mihai Baltac's avatar
Mihai Baltac committed
24
from nifty.operators.smoothing_operator import smooth_util as su
25
from d2o import STRATEGIES
Jait Dixit's avatar
Jait Dixit committed
26

27

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

    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
33
34
    via a gaussian convolution. Fields living over the PowerSpace are directly
    smoothed.
35
36
37

    Parameters
    ----------
Theo Steininger's avatar
Theo Steininger committed
38
    domain : tuple of DomainObjects, i.e. Spaces and FieldTypes
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
78
        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
79
    # ---Overwritten properties and methods---
80
81
82
    def __init__(self, domain=(), sigma=0, log_distances=False,
                 default_spaces=None):
        super(SmoothingOperator, self).__init__(default_spaces)
83
84

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

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

92
93
94
        self.sigma = sigma
        self.log_distances = log_distances

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

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

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

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

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

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

    # ---Added properties and methods---
117

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

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

    def _smoothing_helper(self, x, spaces, inverse):
theos's avatar
theos committed
135
136
137
138
139
140
141
142
143
144
        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
145
146
            spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain))

147
148
149
150
151
152
153
        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
154
        Transformator = FFTOperator(x.domain[spaces[0]])
Jait Dixit's avatar
Jait Dixit committed
155

theos's avatar
theos committed
156
157
158
159
160
        # 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]]
161

theos's avatar
theos committed
162
163
164
        # 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
165

166
        kernel = codomain.get_distance_array(
167
168
169
170
171
            distribution_strategy=axes_local_distribution_strategy)

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

theos's avatar
theos committed
172
        kernel.apply_scalar_function(
173
            codomain.get_fft_smoothing_kernel_function(self.sigma),
theos's avatar
theos committed
174
            inplace=True)
Jait Dixit's avatar
Jait Dixit committed
175

theos's avatar
theos committed
176
177
178
179
180
        # 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
181

182
        reshaper = [transformed_x.shape[i] if i in coaxes else 1
theos's avatar
theos committed
183
184
                    for i in xrange(len(transformed_x.shape))]
        local_kernel = np.reshape(local_kernel, reshaper)
Jait Dixit's avatar
Jait Dixit committed
185

theos's avatar
theos committed
186
187
188
189
190
        # apply the kernel
        if inverse:
            local_transformed_x /= local_kernel
        else:
            local_transformed_x *= local_kernel
Jait Dixit's avatar
Jait Dixit committed
191

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

194
        smoothed_x = Transformator.adjoint_times(transformed_x, spaces=spaces)
theos's avatar
theos committed
195
196
197

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

theos's avatar
theos committed
199
        return result
200
201
202
203
204
205

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

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

        affected_axis = affected_axes[0]
211
212

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

        if self.log_distances:
217
            np.log(distance_array, out=distance_array)
218
219
220
221
222
223
224
225
226

        # 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.
227
            if 0 != affected_axis:
228
229
                local_data_Q = True
            else:
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
                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)

247
248
249
250
251
                augmented_data = x.val.get_data(augmented_slice,
                                                local_keys=True,
                                                copy=False)
                augmented_data = augmented_data.get_local_data(copy=False)

252
                augmented_distance_array = distance_array[augmented_slice]
253
254

        else:
255
256
            raise ValueError("Direct smoothing not implemented for given"
                             "distribution strategy.")
257
258
259
260
261

        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)
262
263
264
            augmented_distance_array = distance_array
            true_start = 0
            true_end = x.shape[affected_axis]
265
266

        # perform the convolution along the affected axes
267
268
269
270
271
272
273
274
275
        # 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)
276
277
278
279
280
        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,
281
                                      true_start, true_end, inverse):
282
        if inverse:
283
            true_sigma = 1. / self.sigma
284
285
286
        else:
            true_sigma = self.sigma

287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        if data.dtype is np.dtype('float32'):
            distances = distances.astype(np.float32, copy=False)
            smoothed_data = su.apply_along_axis_f(
                                  data_axis, data,
                                  startindex=true_start,
                                  endindex=true_end,
                                  distances=distances,
                                  smooth_length=true_sigma,
                                  smoothing_width=self._direct_smoothing_width)
        elif data.dtype is np.dtype('float64'):
            distances = distances.astype(np.float64, copy=False)
            smoothed_data = su.apply_along_axis(
                                  data_axis, data,
                                  startindex=true_start,
                                  endindex=true_end,
                                  distances=distances,
                                  smooth_length=true_sigma,
                                  smoothing_width=self._direct_smoothing_width)

        elif np.issubdtype(data.dtype, np.complexfloating):
            real = self._direct_smoothing_single_axis(data.real,
                                                      data_axis,
                                                      distances,
                                                      true_start,
                                                      true_end, inverse)
            imag = self._direct_smoothing_single_axis(data.imag,
                                                      data_axis,
                                                      distances,
                                                      true_start,
                                                      true_end, inverse)

            return real + 1j*imag

320
        else:
321
322
            raise TypeError("Dtype %s not supported" % str(data.dtype))

323
        return smoothed_data