P_Sparsity.py 7.21 KB
Newer Older
1
from numpy import zeros_like, unravel_index, sum, max
robin.requadt's avatar
robin.requadt committed
2
import numpy as np
jansen31's avatar
jansen31 committed
3
from .proxoperators import ProxOperator
4
from proxtoolbox.Utilities.OrbitalTomog import tile_array, bin_array
jansen31's avatar
jansen31 committed
5
6


robin.requadt's avatar
robin.requadt committed
7
class P_Sparsity(ProxOperator):
8
    """
jansen31's avatar
jansen31 committed
9
    Projection subroutine for projecting onto a sparsity constraint
10
    USAGE: p_Sparsity = P_Sparsity(config)
11
           p_Sparsity.work(u)
12
    """
jansen31's avatar
jansen31 committed
13

jansen31's avatar
jansen31 committed
14
15
16
17
    def __init__(self, config):
        """
        Initialization

robin.requadt's avatar
robin.requadt committed
18
19
        Parameters
        ----------
20
        config : dict - Dictionary containing the problem configuration. It must contain the following mappings:
jansen31's avatar
jansen31 committed
21
22
23
        'sparsity_parameter' : int
        """
        self.sparsity_parameter = config['sparsity_parameter']
24

jansen31's avatar
jansen31 committed
25
26
        if 'sparsity_support' in config:
            self.support = config['sparsity_support'].real.astype(np.uint)
27
            assert (sum(self.support) > 0), 'Invalid empty sparsity_support given'
jansen31's avatar
jansen31 committed
28
29
30
        else:
            self.support = 1

Matthijs's avatar
Matthijs committed
31
        if self.sparsity_parameter > 30 or len(config['u0'].shape) != 2:
32
            def value_selection(original, indices, sparsity_parameter):
33
                idx_for_threshold = unravel_index(indices[-sparsity_parameter], original.shape)
34
35
36
37
38
39
                threshold_val = abs(original[idx_for_threshold].get())
                return (abs(original) >= threshold_val) * original
        else:
            def value_selection(original, indices, sparsity_parameter):
                out = zeros_like(original)
                hits = indices[-sparsity_parameter:].get()
40
                hit_idx = [unravel_index(hit, original.shape) for hit in hits]
41
42
43
44
45
                for _idx in hit_idx:
                    out[_idx[0], _idx[1]] = original[_idx[0], _idx[1]]
                return out

        self.value_selection = value_selection
jansen31's avatar
jansen31 committed
46

jansen31's avatar
jansen31 committed
47
48
49
    def work(self, u):
        """
        Parameters
robin.requadt's avatar
robin.requadt committed
50
        ----------
51
        u : array_like - Input, array to be projected
jansen31's avatar
jansen31 committed
52

robin.requadt's avatar
robin.requadt committed
53
54
        Returns
        -------
55
        p_Sparsity : array_like, the projection
jansen31's avatar
jansen31 committed
56
        """
Matthijs's avatar
Matthijs committed
57
58
        u *= self.support  # apply support (simple 1 if no support)
        p_sparse = 0 * u
59
        sorting = np.argsort(abs(u), axis=None)  # gives indices of sorted array in ascending order
Matthijs's avatar
bugfix    
Matthijs committed
60
        indices = np.asarray([unravel_index(sorting[i], u.shape) for i in range(-1 * self.sparsity_parameter, 0)])
Matthijs's avatar
Matthijs committed
61
        p_sparse[indices[:, 0], indices[:, 1]] = u[indices[:, 0], indices[:, 1]]
robin.requadt's avatar
robin.requadt committed
62

Matthijs's avatar
Matthijs committed
63
        return p_sparse
64
65
66


class P_Sparsity_real(P_Sparsity):
jansen31's avatar
jansen31 committed
67
68
69
70
71
    """
    Projection subroutine for projecting onto a combined sparsity and realness constraint

    order: realness, sparsity
    """
jansen31's avatar
jansen31 committed
72

73
74
    def work(self, u):
        return super(P_Sparsity_real, self).work(u.real)
75
76
77


class P_Sparsity_Symmetric(P_Sparsity):
jansen31's avatar
jansen31 committed
78
79
80
81
82
    """
    Projection subroutine for projecting onto a combined sparsity and symmetry constraint

    order: symmetry, sparsity
    """
jansen31's avatar
jansen31 committed
83

84
85
    def __init__(self, config):
        super(P_Sparsity_Symmetric, self).__init__(config=config)
jansen31's avatar
jansen31 committed
86
        self.symmetry = config['symmetry_type']  # antisymmetric = -1, symmetric = 1
jansen31's avatar
jansen31 committed
87
        self.symmetry_axis = config['symmetry_axis']  # -1 for last, 0 for first, 'both', 'all' or None for full flip
88
89
90
91
92
93
94
95
96
97
98
99

    def work(self, u):
        if self.symmetry_axis in ['both', 'all']:
            mirror = np.flip(u, axis=None)
        else:
            mirror = np.flip(u, axis=self.symmetry_axis)
        inp = (u + self.symmetry * mirror) / 2
        out = super(P_Sparsity_Symmetric, self).work(inp)
        return out


class P_Sparsity_Symmetric_real(P_Sparsity_Symmetric):
jansen31's avatar
jansen31 committed
100
101
102
103
104
    """
    Projection subroutine for projecting onto a combined symmetry, sparsity and realness constraint

    order: realness, symmetry, sparsity
    """
jansen31's avatar
jansen31 committed
105

106
107
    def work(self, u):
        out = super(P_Sparsity_Symmetric_real, self).work(u.real)
108
        return out
jansen31's avatar
jansen31 committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140


class P_Sparsity_Line(P_Sparsity):
    def __init__(self, config):
        """
        Sparsity in n-1 dimensions, where n is the number of dimensions in the data u
        example: data I(x,y,z) projected to Ip(x,y) -> I(x,y,z) = 0 if Ip(x,y,z) is lower than the threshold set by sparsity
        default projection = linear sum

        Args:
            config: identical to P_Sparse, with extra keys:
                'unconstrained dimension': in which dimension the data is not constrained, defaults to -1
                'dimension reduction': method for dimension reducing projection e.g.:
                    'linear' for sum(u) (default),
                    'abs' for sum(abs(u)),
                    'rms' for sqrt(sum(abs(u)**2))

        """
        super(P_Sparsity_Line, self).__init__(config=config)
        if 'unconstrained dimension' not in config:
            self.unconstrained_axis = -1
        else:
            self.unconstrained_axis = config['unconstrained dimension']

        if 'dimension reduction' in config and config['3d to plane projection'] == 'abs':
            def dim_proj(arr):
                return np.sum(abs(arr), axis=config['unconstrained dimension'])

        elif 'dimension reduction' in config and config['3d to plane projection'] == 'rms':
            def dim_proj(arr):
                return np.sqrt(np.sum(abs(arr) ** 2, axis=config['unconstrained dimension']))

141
        else:  # default to linear
jansen31's avatar
jansen31 committed
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
            def dim_proj(arr):
                return np.sum(arr, axis=config['unconstrained dimension'])

        self.dim_proj = dim_proj

    def work(self, u):
        # projection
        sparsity_input = self.dim_proj(u)
        # calculate sparse map
        out = super(P_Sparsity_Line, self).work(sparsity_input)
        # broadcast to original dimensions, keeping original values
        rolled_u = np.moveaxis(u, self.unconstrained_axis, 0)
        new = np.moveaxis(np.where(out != 0, rolled_u, np.zeros_like(rolled_u)),
                          0, self.unconstrained_axis)
        return new.astype(u.dtype)


class P_Sparsity_Line_Real(P_Sparsity_Line):
    """
    Combining P_Sparsity_Line with a realness constraint
    """
163
164
165
    def work(self, u):
        out = super(P_Sparsity_Line_Real, self).work(u.real)
        return out
jansen31's avatar
jansen31 committed
166

167
168
169
170
171

class P_Sparsity_Superpixel(P_Sparsity):
    """
    Apply sparsity on superpixels, i.e. on the binned array
    """
jansen31's avatar
jansen31 committed
172
    def __init__(self, config):
173
174
175
176
177
178
179
180
181
182
183
        super(P_Sparsity_Superpixel, self).__init__(config=config)
        # Set the superpixel size:
        self.superpixel_size = config['superpixel size']
        if self.support is not 1:
            self.support = bin_array(self.support, self.superpixel_size)
            self.support /= max(self.support)
        # Assert that the binning+upsampling conserves the array size
        test_shape = tile_array(bin_array(config['u_0'], self.superpixel_size, pad_zeros=False),
                                self.superpixel_size).shape
        assert test_shape == config['u_0'].shape, 'Given array size does not allow for binning'
        # TODO: allow for padding, then cut of the remainder after tile_array
jansen31's avatar
jansen31 committed
184
185

    def work(self, u):
186
187
188
189
190
191
192
193
194
195
196
197
        binned = bin_array(u, self.superpixel_size)
        constrained = super(P_Sparsity_Superpixel, self).work(binned)
        return tile_array(constrained, self.superpixel_size)


class P_Sparsity_Superpixel_real(P_Sparsity_Superpixel):
    """
    Apply real-valued sparsity on superpixels, i.e. on the binned array
    """
    def work(self, u):
        super(P_Sparsity_Superpixel, self).work(u.real)