smoothing_operator.py 10.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
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):
Jait Dixit's avatar
Jait Dixit committed
29
    # ---Overwritten properties and methods---
30
    def __init__(self, domain=(), sigma=0, log_distances=False):
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
57
58
    @property
    def implemented(self):
        return True
Jait Dixit's avatar
Jait Dixit committed
59

Jait Dixit's avatar
Jait Dixit committed
60
61
    @property
    def symmetric(self):
theos's avatar
theos committed
62
        return True
Jait Dixit's avatar
Jait Dixit committed
63

Jait Dixit's avatar
Jait Dixit committed
64
65
66
    @property
    def unitary(self):
        return False
Jait Dixit's avatar
Jait Dixit committed
67
68

    # ---Added properties and methods---
69

Jait Dixit's avatar
Jait Dixit committed
70
71
72
73
    @property
    def sigma(self):
        return self._sigma

74
75
76
77
78
79
80
81
82
83
84
85
86
    @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
87
88
89
90
91
92
93
94
95
96
        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
97
98
            spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain))

99
100
101
102
103
104
105
        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
106
        Transformator = FFTOperator(x.domain[spaces[0]])
Jait Dixit's avatar
Jait Dixit committed
107

theos's avatar
theos committed
108
109
110
111
112
        # 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]]
113

theos's avatar
theos committed
114
115
116
        # 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
117

118
        kernel = codomain.get_distance_array(
119
120
121
122
123
            distribution_strategy=axes_local_distribution_strategy)

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

theos's avatar
theos committed
124
        kernel.apply_scalar_function(
125
            codomain.get_fft_smoothing_kernel_function(self.sigma),
theos's avatar
theos committed
126
            inplace=True)
Jait Dixit's avatar
Jait Dixit committed
127

theos's avatar
theos committed
128
129
130
131
132
        # 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
133

134
        reshaper = [transformed_x.shape[i] if i in coaxes else 1
theos's avatar
theos committed
135
136
                    for i in xrange(len(transformed_x.shape))]
        local_kernel = np.reshape(local_kernel, reshaper)
Jait Dixit's avatar
Jait Dixit committed
137

theos's avatar
theos committed
138
139
140
141
142
        # apply the kernel
        if inverse:
            local_transformed_x /= local_kernel
        else:
            local_transformed_x *= local_kernel
Jait Dixit's avatar
Jait Dixit committed
143

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

theos's avatar
theos committed
146
147
148
149
        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
150

theos's avatar
theos committed
151
        return result
152
153
154
155
156
157

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

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

        affected_axis = affected_axes[0]
163
164

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

        if self.log_distances:
169
            np.log(distance_array, out=distance_array)
170
171
172
173
174
175
176
177
178

        # 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.
179
            if 0 != affected_axis:
180
181
                local_data_Q = True
            else:
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
                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)

199
200
201
202
203
                augmented_data = x.val.get_data(augmented_slice,
                                                local_keys=True,
                                                copy=False)
                augmented_data = augmented_data.get_local_data(copy=False)

204
                augmented_distance_array = distance_array[augmented_slice]
205
206

        else:
207
208
            raise ValueError("Direct smoothing not implemented for given"
                             "distribution strategy.")
209
210
211
212
213

        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)
214
215
216
            augmented_distance_array = distance_array
            true_start = 0
            true_end = x.shape[affected_axis]
217
218

        # perform the convolution along the affected axes
219
220
221
222
223
224
225
226
227
        # 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)
228
229
230
231
232
        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,
233
                                      true_start, true_end, inverse):
234
        if inverse:
235
            true_sigma = 1. / self.sigma
236
237
238
        else:
            true_sigma = self.sigma

239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
        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

272
        else:
273
274
            raise TypeError("Dtype %s not supported" % str(data.dtype))

275
        return smoothed_data