smoothing_operator.py 3.58 KB
Newer Older
Jait Dixit's avatar
Jait Dixit committed
1
2
3
4
5
6
7
import numpy as np

import nifty.nifty_utilities as utilities
from nifty import RGSpace, LMSpace
from nifty.operators.endomorphic_operator import EndomorphicOperator
from nifty.operators.fft_operator import FFTOperator

8

9
class SmoothingOperator(EndomorphicOperator):
Jait Dixit's avatar
Jait Dixit committed
10
11

    # ---Overwritten properties and methods---
Jait Dixit's avatar
Jait Dixit committed
12
    def __init__(self, domain=(), field_type=(), sigma=None):
13
14
15

        self._domain = self._parse_domain(domain)
        self._field_type = self._parse_field_type(field_type)
Jait Dixit's avatar
Jait Dixit committed
16
17
18

        if len(self.domain) != 1:
            raise ValueError(
19

Jait Dixit's avatar
Jait Dixit committed
20
                    'ERROR: SmoothOperator accepts only exactly one '
21
                    'space as input domain.'
Jait Dixit's avatar
Jait Dixit committed
22
            )
Jait Dixit's avatar
Jait Dixit committed
23
24

        if self.field_type != ():
25
            raise ValueError(
Jait Dixit's avatar
Jait Dixit committed
26
                'ERROR: SmoothOperator field-type must be an '
Jait Dixit's avatar
Jait Dixit committed
27
                'empty tuple.'
28
            )
Jait Dixit's avatar
Jait Dixit committed
29
30

        self._sigma = sigma
Jait Dixit's avatar
Jait Dixit committed
31
32
33

    def _inverse_times(self, x, spaces, types):
        return self._smooth_helper(x, spaces, types, inverse=True)
Jait Dixit's avatar
Jait Dixit committed
34
35

    def _times(self, x, spaces, types):
Jait Dixit's avatar
Jait Dixit committed
36
        return self._smooth_helper(x, spaces, types)
Jait Dixit's avatar
Jait Dixit committed
37

Jait Dixit's avatar
Jait Dixit committed
38
    # ---Mandatory properties and methods---
39
40
41
42
43
44
45
46
    @property
    def domain(self):
        return self._domain

    @property
    def field_type(self):
        return self._field_type

Jait Dixit's avatar
Jait Dixit committed
47
48
49
    @property
    def implemented(self):
        return True
Jait Dixit's avatar
Jait Dixit committed
50

Jait Dixit's avatar
Jait Dixit committed
51
52
53
    @property
    def symmetric(self):
        return False
Jait Dixit's avatar
Jait Dixit committed
54

Jait Dixit's avatar
Jait Dixit committed
55
56
57
    @property
    def unitary(self):
        return False
Jait Dixit's avatar
Jait Dixit committed
58
59
60
61
62
63

    # ---Added properties and methods---
    @property
    def sigma(self):
        return self._sigma

Jait Dixit's avatar
Jait Dixit committed
64
    def _smooth_helper(self, x, spaces, types, inverse=False):
theos's avatar
theos committed
65
66
67
68
69
70
71
72
73
74
        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
75
76
            spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain))

theos's avatar
theos committed
77
        Transformator = FFTOperator(x.domain[spaces[0]])
Jait Dixit's avatar
Jait Dixit committed
78

theos's avatar
theos committed
79
80
81
82
83
        # 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]]
84

theos's avatar
theos committed
85
86
87
        # 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
88

theos's avatar
theos committed
89
90
91
92
93
        kernel = codomain.distance_array(
                        distribution_strategy=axes_local_distribution_strategy)
        kernel.apply_scalar_function(
            codomain.get_smoothing_kernel_function(self.sigma),
            inplace=True)
Jait Dixit's avatar
Jait Dixit committed
94

theos's avatar
theos committed
95
96
97
98
99
        # 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
100

theos's avatar
theos committed
101
102
103
        reshaper = [transformed_x.shape[i] if i in coaxes else 1
                    for i in xrange(len(transformed_x.shape))]
        local_kernel = np.reshape(local_kernel, reshaper)
Jait Dixit's avatar
Jait Dixit committed
104

theos's avatar
theos committed
105
106
107
108
109
        # apply the kernel
        if inverse:
            local_transformed_x /= local_kernel
        else:
            local_transformed_x *= local_kernel
Jait Dixit's avatar
Jait Dixit committed
110

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

theos's avatar
theos committed
113
        result = Transformator.inverse_times(transformed_x, spaces=spaces)
Jait Dixit's avatar
Jait Dixit committed
114

theos's avatar
theos committed
115
        return result