wiener_filter_via_hamiltonian.py 3.3 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
import nifty2go as ift
Martin Reinecke's avatar
Martin Reinecke committed
2
import numpy as np
3

4
np.random.seed(42)
5

Martin Reinecke's avatar
Martin Reinecke committed
6

Martin Reinecke's avatar
Martin Reinecke committed
7
class AdjointFFTResponse(ift.LinearOperator):
8
9
    def __init__(self, FFT, R):
        super(AdjointFFTResponse, self).__init__()
10
11
12
13
14
        self._domain = FFT.target
        self._target = R.target
        self.R = R
        self.FFT = FFT

15
    def _times(self, x):
16
17
        return self.R(self.FFT.adjoint_times(x))

18
    def _adjoint_times(self, x):
19
        return self.FFT(self.R.adjoint_times(x))
Martin Reinecke's avatar
Martin Reinecke committed
20

21
22
23
24
25
26
27
28
29
30
31
32
    @property
    def domain(self):
        return self._domain

    @property
    def target(self):
        return self._target

    @property
    def unitary(self):
        return False

33

34
if __name__ == "__main__":
Jakob Knollmueller's avatar
Jakob Knollmueller committed
35
    # Set up position space
Martin Reinecke's avatar
Martin Reinecke committed
36
    s_space = ift.RGSpace([128, 128])
Jakob Knollmueller's avatar
Jakob Knollmueller committed
37
38
39
    # s_space = HPSpace(32)

    # Define harmonic transformation and associated harmonic space
Martin Reinecke's avatar
Martin Reinecke committed
40
    fft = ift.FFTOperator(s_space)
41
    h_space = fft.target[0]
Jakob Knollmueller's avatar
Jakob Knollmueller committed
42
43

    # Setting up power space
Martin Reinecke's avatar
Martin Reinecke committed
44
    p_space = ift.PowerSpace(h_space)
45

Martin Reinecke's avatar
Martin Reinecke committed
46
47
    # Choosing the prior correlation structure and defining
    # correlation operator
48
    p_spec = (lambda k: (42 / (k + 1) ** 3))
Jakob Knollmueller's avatar
Jakob Knollmueller committed
49

Martin Reinecke's avatar
Martin Reinecke committed
50
    S = ift.create_power_operator(h_space, power_spectrum=p_spec)
51

Jakob Knollmueller's avatar
Jakob Knollmueller committed
52
    # Drawing a sample sh from the prior distribution in harmonic space
Martin Reinecke's avatar
Martin Reinecke committed
53
    sp = ift.Field(p_space, val=p_spec(p_space.k_lengths))
54
    sh = sp.power_synthesize(real_signal=True)
55
56
    ss = fft.adjoint_times(sh)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
57
    # Choosing the measurement instrument
58
59
    #Instrument = ift.FFTSmoothingOperator(s_space, sigma=0.05)
    diag = ift.Field.ones(s_space)
Martin Reinecke's avatar
Martin Reinecke committed
60
    diag.val[20:80, 20:80] = 0
61
    Instrument = ift.DiagonalOperator(diag.weight(-1))
Jakob Knollmueller's avatar
Jakob Knollmueller committed
62

Martin Reinecke's avatar
Martin Reinecke committed
63
    # Adding a harmonic transformation to the instrument
64
    R = AdjointFFTResponse(fft, Instrument)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
65
    signal_to_noise = 1.
66
    ndiag = ift.Field.full(s_space, ss.var()/signal_to_noise)
Martin Reinecke's avatar
Martin Reinecke committed
67
    N = ift.DiagonalOperator(ndiag)
Martin Reinecke's avatar
Martin Reinecke committed
68
    n = ift.Field.from_random(domain=s_space,
69
70
                          random_type='normal',
                          std=ss.std()/np.sqrt(signal_to_noise),
71
                          mean=0)
72

Jakob Knollmueller's avatar
Jakob Knollmueller committed
73
    # Creating the mock data
74
    d = R(sh) + n
75
    j = R.adjoint_times(N.inverse_times(d))
76

Jakob Knollmueller's avatar
Jakob Knollmueller committed
77
78
    # Choosing the minimization strategy

79
    ctrl = ift.GradientNormController(verbose=True,tol_abs_gradnorm=0.1)
80
    inverter = ift.ConjugateGradient(controller=ctrl)
81
82
    controller = ift.GradientNormController(verbose=True,tol_abs_gradnorm=0.1)
    minimizer = ift.RelaxedNewton(controller=controller)
83
    m0 = ift.Field.zeros(h_space)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
84
    # Initializing the Wiener Filter energy
Martin Reinecke's avatar
Martin Reinecke committed
85
86
    energy = ift.library.WienerFilterEnergy(position=m0, d=d, R=R, N=N, S=S,
                                inverter=inverter)
87

Martin Reinecke's avatar
Martin Reinecke committed
88
89
90
    energy, convergence = minimizer(energy)
    m = energy.position
    D = energy.curvature
91
92
    ift.plotting.plot(ss, name="signal.pdf", colormap="Planck-like")
    ift.plotting.plot(fft.inverse_times(m), name="m.pdf", colormap="Planck-like")
93
94

    # sampling the uncertainty map
Martin Reinecke's avatar
Martin Reinecke committed
95
96
97
98
99
100
101
102
103
104
105
106
    sample_variance = ift.Field.zeros(s_space)
    sample_mean = ift.Field.zeros(s_space)

    n_samples = 50
    for i in range(n_samples):
        sample = fft.adjoint_times(ift.sugar.generate_posterior_sample(m, D))
        sample_variance += sample**2
        sample_mean += sample
    sample_mean /= n_samples
    sample_variance /= n_samples
    variance = sample_variance - sample_mean**2
    ift.plotting.plot(variance, name="variance.pdf", colormap="Planck-like")