smoothing_operator.py 14.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
Martin Reinecke's avatar
Martin Reinecke 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

Martin Reinecke's avatar
Martin Reinecke committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def smoothingHelper (x,sigma,dxmax=None):
    """ 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.
Martin Reinecke's avatar
Martin Reinecke committed
48
49
50
    wgt: list with the same number of entries as x
        wgt[i] is an array with nval[i] entries containing the
        normalized smoothing weights.
Martin Reinecke's avatar
Martin Reinecke committed
51
52
53
54
55
56
57
58
59
60
    """

    if dxmax==None:
        dxmax=3.01*sigma

    x=np.asarray(x)

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

Martin Reinecke's avatar
Martin Reinecke committed
61
    wgt=[]
Martin Reinecke's avatar
Martin Reinecke committed
62
63
64
65
66
    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)
Martin Reinecke's avatar
Martin Reinecke committed
67
        wgt.append(t)
Martin Reinecke's avatar
Martin Reinecke committed
68

Martin Reinecke's avatar
Martin Reinecke committed
69
    return ibegin,nval,wgt
Martin Reinecke's avatar
Martin Reinecke committed
70
71

def apply_kernel_along_array(power, startindex, endindex, distances,
Martin Reinecke's avatar
Martin Reinecke committed
72
    smooth_length, smoothing_width,ibegin,nval,wgt):
Martin Reinecke's avatar
Martin Reinecke committed
73
74
75
76
77
78
79

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

    p_smooth = np.empty(endindex-startindex, dtype=np.float64)

    for i in xrange(startindex, endindex):
Martin Reinecke's avatar
Martin Reinecke committed
80
        p_smooth[i-startindex]=np.sum(power[ibegin[i]:ibegin[i]+nval[i]]*wgt[i])
Martin Reinecke's avatar
Martin Reinecke committed
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

    return p_smooth

def getShape(a):
    return tuple(a.shape)

def apply_along_axis(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))
Martin Reinecke's avatar
Martin Reinecke committed
96
    ibegin,nval,wgt=smoothingHelper(distances,smooth_length,smooth_length*smoothing_width)
Martin Reinecke's avatar
Martin Reinecke committed
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

    ind = np.zeros(nd-1, dtype=np.int)
    i = np.zeros(nd, dtype=object)
    shape = getShape(arr)
    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 = apply_kernel_along_array(slicedArr,
                                   startindex,
                                   endindex,
                                   distances,
                                   smooth_length,
Martin Reinecke's avatar
Martin Reinecke committed
117
                                   smoothing_width, ibegin,nval,wgt)
Martin Reinecke's avatar
Martin Reinecke committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

    outshape = np.asarray(getShape(arr))
    outshape[axis] = endindex - startindex
    outarr = np.zeros(outshape, dtype=np.float64)
    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 = apply_kernel_along_array(slicedArr,
                                       startindex,
                                       endindex,
                                       distances,
                                       smooth_length,
Martin Reinecke's avatar
Martin Reinecke committed
139
                                       smoothing_width,ibegin,nval,wgt)
Martin Reinecke's avatar
Martin Reinecke committed
140
141
142
143
        outarr[tuple(i.tolist())] = res
        k += 1

    return outarr
144

145
class SmoothingOperator(EndomorphicOperator):
Jait Dixit's avatar
Jait Dixit committed
146
    # ---Overwritten properties and methods---
Martin Reinecke's avatar
Martin Reinecke committed
147
    def __init__(self, domain, sigma, log_distances=False,
148
149
                 default_spaces=None):
        super(SmoothingOperator, self).__init__(default_spaces)
150
151

        self._domain = self._parse_domain(domain)
Jait Dixit's avatar
Jait Dixit committed
152
153
154

        if len(self.domain) != 1:
            raise ValueError(
155
156
                'ERROR: SmoothOperator accepts only exactly one '
                'space as input domain.'
Jait Dixit's avatar
Jait Dixit committed
157
            )
Jait Dixit's avatar
Jait Dixit committed
158

159
160
161
        self.sigma = sigma
        self.log_distances = log_distances

162
        self._direct_smoothing_width = 3.
Jait Dixit's avatar
Jait Dixit committed
163

164
    def _inverse_times(self, x, spaces):
165
        return self._smoothing_helper(x, spaces, inverse=True)
Jait Dixit's avatar
Jait Dixit committed
166

167
    def _times(self, x, spaces):
168
        return self._smoothing_helper(x, spaces, inverse=False)
Jait Dixit's avatar
Jait Dixit committed
169

Jait Dixit's avatar
Jait Dixit committed
170
    # ---Mandatory properties and methods---
171
172
173
174
    @property
    def domain(self):
        return self._domain

Jait Dixit's avatar
Jait Dixit committed
175
    @property
Martin Reinecke's avatar
Martin Reinecke committed
176
    def self_adjoint(self):
theos's avatar
theos committed
177
        return True
Jait Dixit's avatar
Jait Dixit committed
178

Jait Dixit's avatar
Jait Dixit committed
179
180
181
    @property
    def unitary(self):
        return False
Jait Dixit's avatar
Jait Dixit committed
182
183

    # ---Added properties and methods---
184

Jait Dixit's avatar
Jait Dixit committed
185
186
187
188
    @property
    def sigma(self):
        return self._sigma

189
190
191
192
193
194
195
196
197
198
199
200
201
    @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
202
203
204
205
206
207
208
209
210
211
        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
212
213
            spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain))

214
215
216
217
218
219
220
        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
221
        Transformator = FFTOperator(x.domain[spaces[0]])
Jait Dixit's avatar
Jait Dixit committed
222

theos's avatar
theos committed
223
224
225
226
227
        # 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]]
228

theos's avatar
theos committed
229
230
231
        # 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
232

233
        kernel = codomain.get_distance_array(
234
235
236
237
238
            distribution_strategy=axes_local_distribution_strategy)

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

theos's avatar
theos committed
239
        kernel.apply_scalar_function(
240
            codomain.get_fft_smoothing_kernel_function(self.sigma),
theos's avatar
theos committed
241
            inplace=True)
Jait Dixit's avatar
Jait Dixit committed
242

theos's avatar
theos committed
243
244
245
246
247
        # 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
248

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

theos's avatar
theos committed
253
254
255
256
257
        # apply the kernel
        if inverse:
            local_transformed_x /= local_kernel
        else:
            local_transformed_x *= local_kernel
Jait Dixit's avatar
Jait Dixit committed
258

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

theos's avatar
theos committed
261
262
263
264
        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
265

theos's avatar
theos committed
266
        return result
267
268
269
270
271
272

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

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

        affected_axis = affected_axes[0]
278
279

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

        if self.log_distances:
284
            np.log(distance_array, out=distance_array)
285
286
287
288
289
290
291
292
293

        # 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.
294
            if 0 != affected_axis:
295
296
                local_data_Q = True
            else:
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
                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)

314
315
316
317
318
                augmented_data = x.val.get_data(augmented_slice,
                                                local_keys=True,
                                                copy=False)
                augmented_data = augmented_data.get_local_data(copy=False)

319
                augmented_distance_array = distance_array[augmented_slice]
320
321

        else:
322
323
            raise ValueError("Direct smoothing not implemented for given"
                             "distribution strategy.")
324
325
326
327
328

        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)
329
330
331
            augmented_distance_array = distance_array
            true_start = 0
            true_end = x.shape[affected_axis]
332
333

        # perform the convolution along the affected axes
334
335
336
337
338
339
340
341
342
        # 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)
343
344
345
346
347
        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,
348
                                      true_start, true_end, inverse):
349
        if inverse:
350
            true_sigma = 1. / self.sigma
351
352
353
        else:
            true_sigma = self.sigma

Martin Reinecke's avatar
Martin Reinecke committed
354
        #MR: do not split for complex arrays!
355
356
        if data.dtype is np.dtype('float32'):
            distances = distances.astype(np.float32, copy=False)
Martin Reinecke's avatar
Martin Reinecke committed
357
            smoothed_data = apply_along_axis(
358
359
360
361
362
363
364
365
                                  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)
Martin Reinecke's avatar
Martin Reinecke committed
366
            smoothed_data = apply_along_axis(
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
                                  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

388
        else:
389
390
            raise TypeError("Dtype %s not supported" % str(data.dtype))

391
        return smoothed_data