There is a maintenance of MPCDF Gitlab on Thursday, April 22st 2020, 9:00 am CEST - Expect some service interruptions during this time

demo_excaliwir.py 9.61 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
    p = problem(x, log = False)
##    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
289 290 291

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