demo_excaliwir.py 9.63 KB
Newer Older
Marco Selig's avatar
Marco Selig committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
## NIFTY (Numerical Information Field Theory) has been developed at the
## Max-Planck-Institute for Astrophysics.
##
## Copyright (C) 2013 Max-Planck-Society
##
## Author: Marco Selig
## Project homepage: <http://www.mpa-garching.mpg.de/ift/nifty/>
##
## 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/>.

"""
    ..                  __   ____   __
    ..                /__/ /   _/ /  /_
    ..      __ ___    __  /  /_  /   _/  __   __
    ..    /   _   | /  / /   _/ /  /   /  / /  /
    ..   /  / /  / /  / /  /   /  /_  /  /_/  /
    ..  /__/ /__/ /__/ /__/    \___/  \___   /  demo
    ..                               /______/

Marco Selig's avatar
Marco Selig committed
31
    NIFTY demo for (extended) critical Wiener filtering of Gaussian random signals.
Marco Selig's avatar
Marco Selig committed
32
33
34

"""
from __future__ import division
Ultima's avatar
Ultima committed
35
36
37
38
39
40
41
import matplotlib as mpl
mpl.use('Agg')

import imp
#nifty = imp.load_module('nifty', None,
#                        '/home/steininger/Downloads/nifty', ('','',5))

Marco Selig's avatar
Marco Selig committed
42
43
44
from nifty import *


Marco Selig's avatar
Marco Selig committed
45
##=============================================================================
Marco Selig's avatar
Marco Selig committed
46

Marco Selig's avatar
Marco Selig committed
47
class problem(object):
Marco Selig's avatar
Marco Selig committed
48

Ultima's avatar
Ultima committed
49
    def __init__(self, x_space, s2n=6, **kwargs):
Marco Selig's avatar
Marco Selig committed
50
51
52
53
54
55
56
57
58
        """
            Sets up a Wiener filter problem.

            Parameters
            ----------
            x_space : space
                Position space the signal lives in.
            s2n : float, *optional*
                Signal-to-noise ratio (default: 2).
Marco Selig's avatar
Marco Selig committed
59
60

        """
Ultima's avatar
Ultima committed
61
        self.store = []
Marco Selig's avatar
Marco Selig committed
62
63
64
65
        ## set signal space
        self.z = x_space
        ## set conjugate space
        self.k = self.z.get_codomain()
Ultima's avatar
Ultima committed
66
        #self.k.power_indices.set_default()
67
        #self.k.set_power_indices(**kwargs)
Marco Selig's avatar
Marco Selig committed
68
69

        ## set some power spectrum
Ultima's avatar
Ultima committed
70
        self.power = (lambda k: 42 / (k + 1) ** 2)
Marco Selig's avatar
Marco Selig committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84

        ## define signal covariance
        self.S = power_operator(self.k, spec=self.power, bare=True)
        ## define projector to spectral bands
        self.Sk = self.S.get_projection_operator()
        ## generate signal
        self.s = self.S.get_random_field(domain=self.z)

        ## define response
        self.R = response_operator(self.z, sigma=0.0, mask=1.0)
        ## get data space
        d_space = self.R.target

        ## define noise covariance
Ultima's avatar
Ultima committed
85
86
        #self.N = diagonal_operator(d_space, diag=abs(s2n) * self.s.var(), bare=True)
        self.N = diagonal_operator(d_space, diag=abs(s2n), bare=True)
Marco Selig's avatar
Marco Selig committed
87
88
89
90
91
92
93
94
95
        ## define (plain) projector
        self.Nj = projection_operator(d_space)
        ## generate noise
        n = self.N.get_random_field(domain=d_space)

        ## compute data
        self.d = self.R(self.s) + n

        ## define information source
96
97
        #self.j = self.R.adjoint_times(self.N.inverse_times(self.d), target=self.k)
        self.j = self.R.adjoint_times(self.N.inverse_times(self.d))
Marco Selig's avatar
Marco Selig committed
98
        ## define information propagator
Ultima's avatar
Ultima committed
99
100
101
        self.D = propagator_operator(S=self.S,
                                     N=self.N,
                                     R=self.R)
Marco Selig's avatar
Marco Selig committed
102
103
104
105
106
107
108

        ## reserve map
        self.m = None

    ##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    def solve(self, newspec=None):
Marco Selig's avatar
Marco Selig committed
109
        """
Marco Selig's avatar
Marco Selig committed
110
111
            Solves the Wiener filter problem for a given power spectrum
            reconstructing a signal estimate.
Marco Selig's avatar
Marco Selig committed
112

Marco Selig's avatar
Marco Selig committed
113
114
            Parameters
            ----------
115
            newspace : {scalar, list, array, Field, function}, *optional*
Marco Selig's avatar
Marco Selig committed
116
                Assumed power spectrum (default: k ** -2).
Marco Selig's avatar
Marco Selig committed
117

Marco Selig's avatar
Marco Selig committed
118
119
120
121
122
123
124
        """
        ## set (given) power spectrum
        if(newspec is None):
            newspec = np.r_[1, 1 / self.k.power_indices["kindex"][1:] ** 2] ## Laplacian
        elif(newspec is False):
            newspec = self.power ## assumed to be known
        self.S.set_power(newspec, bare=True)
Marco Selig's avatar
Marco Selig committed
125

Marco Selig's avatar
Marco Selig committed
126
127
        ## reconstruct map
        self.m = self.D(self.j, W=self.S, tol=1E-3, note=False)
Marco Selig's avatar
Marco Selig committed
128

Marco Selig's avatar
Marco Selig committed
129
    ##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Marco Selig's avatar
Marco Selig committed
130

Marco Selig's avatar
Marco Selig committed
131
132
133
134
135
136
137
    def solve_critical(self, newspec=None, q=0, alpha=1, delta=1, epsilon=0):
        """
            Solves the (generalised) Wiener filter problem
            reconstructing a signal estimate and a power spectrum.

            Parameters
            ----------
138
            newspace : {scalar, list, array, Field, function}, *optional*
Marco Selig's avatar
Marco Selig committed
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
                Initial power spectrum (default: k ** -2).
            q : {scalar, list, array}, *optional*
                Spectral scale parameter of the assumed inverse-Gamme prior
                (default: 0).
            alpha : {scalar, list, array}, *optional*
                Spectral shape parameter of the assumed inverse-Gamme prior
                (default: 1).
            delta : float, *optional*
                First filter perception parameter (default: 1).
            epsilon : float, *optional*
                Second filter perception parameter (default: 0).

            See Also
            --------
            infer_power
Marco Selig's avatar
Marco Selig committed
154

Marco Selig's avatar
Marco Selig committed
155
156
157
158
159
160
161
        """
        ## set (initial) power spectrum
        if(newspec is None):
            newspec = np.r_[1, 1 / self.k.power_indices["kindex"][1:] ** 2] ## Laplacian
        elif(newspec is False):
            newspec = self.power ## assumed to be known
        self.S.set_power(newspec, bare=True)
Marco Selig's avatar
Marco Selig committed
162

Marco Selig's avatar
Marco Selig committed
163
164
        ## pre-compute denominator
        denominator = self.k.power_indices["rho"] + 2 * (alpha - 1 + abs(epsilon))
Marco Selig's avatar
Marco Selig committed
165

Ultima's avatar
Ultima committed
166
167
        self.save_signal_and_data()

Marco Selig's avatar
Marco Selig committed
168
        ## iterate
Ultima's avatar
Ultima committed
169
        i = 0
Marco Selig's avatar
Marco Selig committed
170
171
172
173
        iterating = True
        while(iterating):

            ## reconstruct map
174
            self.m = self.D(self.j, W=self.S, tol=1E-3, note=True)
Marco Selig's avatar
Marco Selig committed
175
176
            if(self.m is None):
                break
Ultima's avatar
Ultima committed
177
            #print'Reconstructed m'
Marco Selig's avatar
Marco Selig committed
178
179
            ## reconstruct power spectrum
            tr_B1 = self.Sk.pseudo_tr(self.m) ## == Sk(m).pseudo_dot(m)
180
181
            print 'Calculated trace B1'
            print ('tr_b1', tr_B1)
Marco Selig's avatar
Marco Selig committed
182
            tr_B2 = self.Sk.pseudo_tr(self.D, loop=True)
183
184
185
            print 'Calculated trace B2'
            print ('tr_B2', tr_B2)
            numerator = 2 * q + tr_B1 +  tr_B2 * abs(delta)  ## non-bare(!)
Marco Selig's avatar
Marco Selig committed
186
            power = numerator / denominator
187
188
189
190
            print ('numerator', numerator)
            print ('denominator', denominator)
            print ('power', power)
            print 'Calculated power'
Ultima's avatar
Ultima committed
191

Ultima's avatar
Ultima committed
192
            #power = np.clip(power, 0.00000001, np.max(power))
Ultima's avatar
Ultima committed
193
194
195
196
            self.store += [{'tr_B1': tr_B1,
                            'tr_B2': tr_B2,
                            'num': numerator,
                            'denom': denominator}]
Marco Selig's avatar
Marco Selig committed
197
            ## check convergence
Marco Selig's avatar
Marco Selig committed
198
            dtau = log(power / self.S.get_power(), base=self.S.get_power())
Ultima's avatar
Ultima committed
199
            print ('dtau', np.max(np.abs(dtau)))
Marco Selig's avatar
Marco Selig committed
200
            iterating = (np.max(np.abs(dtau)) > 2E-2)
Ultima's avatar
Ultima committed
201
202
203
            #printmax(np.abs(dtau))
            self.save_map(i)
            i += 1
Marco Selig's avatar
Marco Selig committed
204
205
206
207
208
            ## update signal covariance
            self.S.set_power(power, bare=False) ## auto-updates D

    ##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

Ultima's avatar
Ultima committed
209
210
211
212
    def save_signal_and_data(self):
        self.s.plot(title="signal", save="img/signal.png")

        try:
213
            d_ = Field(self.z, val=self.d.val, target=self.k)
Ultima's avatar
Ultima committed
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
            d_.plot(title="data", vmin=self.s.min(), vmax=self.s.max(),
                    save="img/data.png")
        except:
            pass


    def save_map(self, index=0):

        # save map
        if(self.m is None):
            pass
        else:
            self.m.plot(title="reconstructed map",
                        vmin=self.s.min(), vmax=self.s.max(),
                        save="img/map_"+str(index)+".png")
            self.m.plot(power=True, mono=False,
                        other=(self.power, self.S.get_power()),
                        nbin=None, binbounds=None, log=False,
                        save='img/map_power_'+str(index)+".png")

Marco Selig's avatar
Marco Selig committed
234
235
236
237
238
239
240
241
242
    def plot(self):
        """
            Produces plots.

        """
        ## plot signal
        self.s.plot(title="signal")
        ## plot data
        try:
243
            d_ = Field(self.z, val=self.d.val, target=self.k)
Marco Selig's avatar
Marco Selig committed
244
245
246
247
248
249
250
251
252
            d_.plot(title="data", vmin=self.s.min(), vmax=self.s.max())
        except:
            pass
        ## plot map
        if(self.m is None):
            self.s.plot(power=True, mono=False, other=self.power)
        else:
            self.m.plot(title="reconstructed map", vmin=self.s.min(), vmax=self.s.max())
            self.m.plot(power=True, mono=False, other=(self.power, self.S.get_power()))
Marco Selig's avatar
Marco Selig committed
253
254
255
256

##=============================================================================

##-----------------------------------------------------------------------------
257
#
Marco Selig's avatar
Marco Selig committed
258
if(__name__=="__main__"):
259
    x = RGSpace((1280), zerocenter=True)
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
    p = problem(x, log = False)
    about.warnings.off()
##    pl.close("all")
#
#    ## define signal space
#    x_space = rg_space(128)
#
#    ## setup problem
#    p = problem(x_space, log=True)
#    ## solve problem given some power spectrum
#    p.solve()
#    ## solve problem
#    p.solve_critical()
#
#    p.plot()
#
#    ## retrieve objects
#    k_space = p.k
#    power = p.power
#    S = p.S
#    Sk = p.Sk
#    s = p.s
#    R = p.R
#    d_space = p.R.target
#    N = p.N
#    Nj = p.Nj
#    d = p.d
#    j = p.j
#    D = p.D
#    m = p.m
Marco Selig's avatar
Marco Selig committed
290
291
292

##-----------------------------------------------------------------------------