nonlinear_critical_filter.py 4.25 KB
Newer Older
Philipp Arras's avatar
Philipp Arras committed
1
2
import nifty5 as ift
from nifty5.library.nonlinearities import Exponential
Martin Reinecke's avatar
Martin Reinecke committed
3
4
5
6
7
import numpy as np
np.random.seed(42)


def adjust_zero_mode(m0, t0):
Martin Reinecke's avatar
Martin Reinecke committed
8
    mtmp = m0.to_global_data().copy()
Martin Reinecke's avatar
Martin Reinecke committed
9
    zero_position = len(m0.shape)*(0,)
10
11
    zero_mode = mtmp[zero_position]
    mtmp[zero_position] = zero_mode / abs(zero_mode)
Martin Reinecke's avatar
Martin Reinecke committed
12
    ttmp = t0.to_global_data().copy()
13
    ttmp[0] += 2 * np.log(abs(zero_mode))
14
15
    return (ift.Field.from_global_data(m0.domain, mtmp),
            ift.Field.from_global_data(t0.domain, ttmp))
Martin Reinecke's avatar
Martin Reinecke committed
16
17


Martin Reinecke's avatar
Martin Reinecke committed
18
if __name__ == "__main__":
Martin Reinecke's avatar
Martin Reinecke committed
19
20
21
22
23
24
25
26
    noise_level = 1.
    p_spec = (lambda k: (1. / (k + 1) ** 2))

    # nonlinearity = Linear()
    nonlinearity = Exponential()
    # Set up position space
    # s_space = ift.RGSpace([1024])
    s_space = ift.HPSpace(32)
Martin Reinecke's avatar
Martin Reinecke committed
27
    h_space = s_space.get_default_codomain()
Martin Reinecke's avatar
Martin Reinecke committed
28
29

    # Define harmonic transformation and associated harmonic space
Martin Reinecke's avatar
Martin Reinecke committed
30
    HT = ift.HarmonicTransformOperator(h_space, target=s_space)
Martin Reinecke's avatar
Martin Reinecke committed
31
32
33
34
35

    # Setting up power space
    p_space = ift.PowerSpace(h_space,
                             binbounds=ift.PowerSpace.useful_binbounds(
                                 h_space, logarithmic=True))
36
37
    # Choosing the prior correlation structure and defining
    # correlation operator
Martin Reinecke's avatar
Martin Reinecke committed
38
    p = ift.PS_field(p_space, p_spec)
Martin Reinecke's avatar
Martin Reinecke committed
39
    log_p = ift.log(p)
40
    S = ift.create_power_operator(h_space, lambda k: 1.)
Martin Reinecke's avatar
Martin Reinecke committed
41
42

    # Drawing a sample sh from the prior distribution in harmonic space
43
    sh = S.draw_sample()
Martin Reinecke's avatar
Martin Reinecke committed
44
45
46
47
48

    # Choosing the measurement instrument
    # Instrument = SmoothingOperator(s_space, sigma=0.01)
    mask = np.ones(s_space.shape)
    mask[6000:8000] = 0.
49
    mask = ift.Field.from_global_data(s_space, mask)
Martin Reinecke's avatar
Martin Reinecke committed
50
51

    MaskOperator = ift.DiagonalOperator(mask)
Martin Reinecke's avatar
Martin Reinecke committed
52
53
    R = ift.GeometryRemover(s_space)
    R = R*MaskOperator
Martin Reinecke's avatar
Martin Reinecke committed
54
55
56
    # R = R*HT
    # R = R * ift.create_harmonic_smoothing_operator((harmonic_space,), 0,
    #                                                response_sigma)
Martin Reinecke's avatar
Martin Reinecke committed
57
    MeasurementOperator = R
Martin Reinecke's avatar
Martin Reinecke committed
58

Martin Reinecke's avatar
Martin Reinecke committed
59
60
    d_space = MeasurementOperator.target

Martin Reinecke's avatar
Martin Reinecke committed
61
62
    Distributor = ift.PowerDistributor(target=h_space, power_space=p_space)
    power = Distributor(ift.exp(0.5*log_p))
Martin Reinecke's avatar
Martin Reinecke committed
63
    # Creating the mock data
Martin Reinecke's avatar
Martin Reinecke committed
64
65
66
    true_sky = nonlinearity(HT(power*sh))
    noiseless_data = MeasurementOperator(true_sky)
    noise_amplitude = noiseless_data.val.std()*noise_level
67
    N = ift.ScalingOperator(noise_amplitude**2, d_space)
Martin Reinecke's avatar
Martin Reinecke committed
68
    n = N.draw_sample()
Martin Reinecke's avatar
Martin Reinecke committed
69
70
    # Creating the mock data
    d = noiseless_data + n
Martin Reinecke's avatar
Martin Reinecke committed
71

Martin Reinecke's avatar
step 1  
Martin Reinecke committed
72
73
    m0 = ift.full(h_space, 1e-7)
    t0 = ift.full(p_space, -4.)
Martin Reinecke's avatar
Martin Reinecke committed
74
    power0 = Distributor.times(ift.exp(0.5 * t0))
Martin Reinecke's avatar
Martin Reinecke committed
75

76
    IC1 = ift.GradientNormController(name="IC1", iteration_limit=100,
Martin Reinecke's avatar
Martin Reinecke committed
77
                                     tol_abs_gradnorm=1e-3)
78
79
    LS = ift.LineSearchStrongWolfe(c2=0.02)
    minimizer = ift.RelaxedNewton(IC1, line_searcher=LS)
Martin Reinecke's avatar
Martin Reinecke committed
80

81
    ICI = ift.GradientNormController(iteration_limit=500,
Martin Reinecke's avatar
Martin Reinecke committed
82
83
84
85
                                     tol_abs_gradnorm=1e-3)
    inverter = ift.ConjugateGradient(controller=ICI)

    for i in range(20):
Martin Reinecke's avatar
Martin Reinecke committed
86
        power0 = Distributor(ift.exp(0.5*t0))
87
        map0_energy = ift.library.NonlinearWienerFilterEnergy(
Martin Reinecke's avatar
Martin Reinecke committed
88
            m0, d, MeasurementOperator, nonlinearity, HT, power0, N, S,
89
            inverter=inverter)
Martin Reinecke's avatar
Martin Reinecke committed
90
91
92
93
94
95
96
97
98

        # Minimization with chosen minimizer
        map0_energy, convergence = minimizer(map0_energy)
        m0 = map0_energy.position

        # Updating parameters for correlation structure reconstruction
        D0 = map0_energy.curvature

        # Initializing the power energy with updated parameters
99
        power0_energy = ift.library.NonlinearPowerEnergy(
Martin Reinecke's avatar
Martin Reinecke committed
100
            position=t0, d=d, N=N, xi=m0, D=D0, ht=HT,
101
            Instrument=MeasurementOperator, nonlinearity=nonlinearity,
Martin Reinecke's avatar
Martin Reinecke committed
102
            Distributor=Distributor, sigma=1., samples=2, inverter=inverter)
Martin Reinecke's avatar
Martin Reinecke committed
103

Martin Reinecke's avatar
Martin Reinecke committed
104
        power0_energy = minimizer(power0_energy)[0]
Martin Reinecke's avatar
Martin Reinecke committed
105
106
107
108

        # Setting new power spectrum
        t0 = power0_energy.position

109
110
111
        # break degeneracy between amplitude and excitation by setting
        # excitation monopole to 1
        m0, t0 = adjust_zero_mode(m0, t0)
Martin Reinecke's avatar
Martin Reinecke committed
112

Martin Reinecke's avatar
Martin Reinecke committed
113
114
    plotdict = {"colormap": "Planck-like"}
    ift.plot(true_sky, name="true_sky.png", **plotdict)
Martin Reinecke's avatar
Martin Reinecke committed
115
    ift.plot(nonlinearity(HT(power0*m0)),
Martin Reinecke's avatar
Martin Reinecke committed
116
117
             name="reconstructed_sky.png", **plotdict)
    ift.plot(MeasurementOperator.adjoint_times(d), name="data.png", **plotdict)
118
    ift.plot([ift.exp(t0), p], name="ps.png")