wiener_filter_via_hamiltonian.py 2.55 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
import nifty4 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

7
if __name__ == "__main__":
Jakob Knollmueller's avatar
Jakob Knollmueller committed
8
    # Set up position space
Martin Reinecke's avatar
Martin Reinecke committed
9
    s_space = ift.RGSpace([128, 128])
Martin Reinecke's avatar
Martin Reinecke committed
10
    # s_space = ift.HPSpace(32)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
11

Martin Reinecke's avatar
Martin Reinecke committed
12
13
    # Define associated harmonic space and harmonic transformation
    h_space = s_space.get_default_codomain()
Martin Reinecke's avatar
Martin Reinecke committed
14
    ht = ift.HarmonicTransformOperator(h_space, s_space)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
15
16

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

Martin Reinecke's avatar
Martin Reinecke committed
19
20
    # Choosing the prior correlation structure and defining
    # correlation operator
Martin Reinecke's avatar
Martin Reinecke committed
21
    p_spec = (lambda k: (42/(k+1)**3))
Jakob Knollmueller's avatar
Jakob Knollmueller committed
22

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

Jakob Knollmueller's avatar
Jakob Knollmueller committed
25
    # Drawing a sample sh from the prior distribution in harmonic space
Martin Reinecke's avatar
Martin Reinecke committed
26
    sp = ift.PS_field(p_space, p_spec)
Martin Reinecke's avatar
adjust    
Martin Reinecke committed
27
    sh = ift.power_synthesize(sp, real_signal=True)
28

Jakob Knollmueller's avatar
Jakob Knollmueller committed
29
    # Choosing the measurement instrument
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
30
31
    diag = np.ones(s_space.shape)
    diag[20:80, 20:80] = 0
Martin Reinecke's avatar
Martin Reinecke committed
32
    diag = ift.Field(s_space, ift.dobj.from_global_data(diag))
33
    Instrument = ift.DiagonalOperator(diag)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
34

Martin Reinecke's avatar
Martin Reinecke committed
35
    # Adding a harmonic transformation to the instrument
Martin Reinecke's avatar
Martin Reinecke committed
36
37
    R = Instrument*ht
    noiseless_data = R(sh)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
38
    signal_to_noise = 1.
Martin Reinecke's avatar
Martin Reinecke committed
39
    noise_amplitude = noiseless_data.val.std()/signal_to_noise
Martin Reinecke's avatar
Martin Reinecke committed
40
    N = ift.DiagonalOperator(ift.Field.full(s_space, noise_amplitude**2))
Martin Reinecke's avatar
Martin Reinecke committed
41
    n = ift.Field.from_random(domain=s_space,
Martin Reinecke's avatar
Martin Reinecke committed
42
                              random_type='normal',
Martin Reinecke's avatar
Martin Reinecke committed
43
                              std=noise_amplitude,
Martin Reinecke's avatar
Martin Reinecke committed
44
                              mean=0)
45

Jakob Knollmueller's avatar
Jakob Knollmueller committed
46
    # Creating the mock data
Martin Reinecke's avatar
Martin Reinecke committed
47
    d = noiseless_data + n
48
    j = R.adjoint_times(N.inverse_times(d))
49

Jakob Knollmueller's avatar
Jakob Knollmueller committed
50
51
    # Choosing the minimization strategy

52
    ctrl = ift.GradientNormController(name="inverter", tol_abs_gradnorm=0.1)
53
    inverter = ift.ConjugateGradient(controller=ctrl)
54
    controller = ift.GradientNormController(name="min", tol_abs_gradnorm=0.1)
55
    minimizer = ift.RelaxedNewton(controller=controller)
56
    m0 = ift.Field.zeros(h_space)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
57
    # Initializing the Wiener Filter energy
Martin Reinecke's avatar
Martin Reinecke committed
58
    energy = ift.library.WienerFilterEnergy(position=m0, d=d, R=R, N=N, S=S,
Martin Reinecke's avatar
Martin Reinecke committed
59
                                            inverter=inverter)
60

Martin Reinecke's avatar
Martin Reinecke committed
61
62
63
    energy, convergence = minimizer(energy)
    m = energy.position
    D = energy.curvature
Martin Reinecke's avatar
Martin Reinecke committed
64
65
    ift.plot(ht(sh), name="signal.png", colormap="Planck-like")
    ift.plot(ht(m), name="m.png", colormap="Planck-like")
66
67

    # sampling the uncertainty map
Martin Reinecke's avatar
Martin Reinecke committed
68
69
70
    sample_variance = ift.Field.zeros(s_space)
    sample_mean = ift.Field.zeros(s_space)

71
    mean, variance = ift.probe_with_posterior_samples(D, ht, 50)
72
    ift.plot(variance, name="variance.png", colormap="Planck-like")
73
    ift.plot(mean, name="mean.png", colormap="Planck-like")