getting_started_mf.py 6.92 KB
Newer Older
Philipp Frank's avatar
Philipp Frank 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
31
# 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/>.
#
# Copyright(C) 2013-2019 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.

############################################################
# Non-linear tomography
#
# The signal is a sigmoid-normal distributed field.
# The data is the field integrated along lines of sight that are
# randomly (set mode=0) or radially (mode=1) distributed
#
# Demo takes a while to compute
#############################################################

import sys

import numpy as np

Martin Reinecke's avatar
5->6    
Martin Reinecke committed
32
import nifty6 as ift
Philipp Frank's avatar
Philipp Frank committed
33

Philipp Arras's avatar
Philipp Arras committed
34

Philipp Frank's avatar
Philipp Frank committed
35
class SingleDomain(ift.LinearOperator):
Philipp Arras's avatar
Philipp Arras committed
36
    def __init__(self, domain, target):
Philipp Frank's avatar
Philipp Frank committed
37
38
39
        self._domain = ift.makeDomain(domain)
        self._target = ift.makeDomain(target)
        self._capability = self.TIMES | self.ADJOINT_TIMES
Philipp Arras's avatar
Philipp Arras committed
40
41
42

    def apply(self, x, mode):
        self._check_input(x, mode)
Martin Reinecke's avatar
Martin Reinecke committed
43
        return ift.makeField(self._tgt(mode), x.val)
Philipp Arras's avatar
Philipp Arras committed
44

Philipp Frank's avatar
Philipp Frank committed
45
46

def random_los(n_los):
Martin Reinecke's avatar
Martin Reinecke committed
47
48
    starts = list(ift.random.current_rng().random((n_los, 2)).T)
    ends = list(ift.random.current_rng().random((n_los, 2)).T)
Philipp Frank's avatar
Philipp Frank committed
49
50
51
52
    return starts, ends


def radial_los(n_los):
Martin Reinecke's avatar
Martin Reinecke committed
53
54
    starts = list(ift.random.current_rng().random((n_los, 2)).T)
    ends = list(0.5 + 0*ift.random.current_rng().random((n_los, 2)).T)
Philipp Frank's avatar
Philipp Frank committed
55
56
57
58
59
60
61
62
63
64
    return starts, ends


if __name__ == '__main__':
    # Choose between random line-of-sight response (mode=0) and radial lines
    # of sight (mode=1)
    if len(sys.argv) == 2:
        mode = int(sys.argv[1])
    else:
        mode = 0
Lukas Platz's avatar
Lukas Platz committed
65
66

    # Preparing the filename string for store results
Philipp Frank's avatar
Philipp Frank committed
67
68
    filename = "getting_started_mf_mode_{}_".format(mode) + "{}.png"

Lukas Platz's avatar
Lukas Platz committed
69
    # Set up signal domain
Philipp Frank's avatar
Philipp Frank committed
70
71
72
73
74
    npix1, npix2 = 128, 128
    position_space = ift.RGSpace([npix1, npix2])
    sp1 = ift.RGSpace(npix1)
    sp2 = ift.RGSpace(npix2)

Lukas Platz's avatar
Lukas Platz committed
75
    # Set up signal model
Lukas Platz's avatar
Lukas Platz committed
76
    cfmaker = ift.CorrelatedFieldMaker.make(0., 1e-2, 1e-6, '')
77
78
    cfmaker.add_fluctuations(sp1, 0.1, 1e-2, 1, .1, .01, .5, -2, 1., 'amp1')
    cfmaker.add_fluctuations(sp2, 0.1, 1e-2, 1, .1, .01, .5,
Philipp Arras's avatar
Philipp Arras committed
79
                             -1.5, .5, 'amp2')
80
    correlated_field = cfmaker.finalize()
Philipp Arras's avatar
Philipp Arras committed
81

82
83
    A1 = cfmaker.normalized_amplitudes[0]
    A2 = cfmaker.normalized_amplitudes[1]
Philipp Arras's avatar
Philipp Arras committed
84
    DC = SingleDomain(correlated_field.target, position_space)
Philipp Frank's avatar
Philipp Frank committed
85

Lukas Platz's avatar
Lukas Platz committed
86
    ## Apply a nonlinearity
Philipp Frank's avatar
Philipp Frank committed
87
88
89
90
91
92
93
94
95
96
    signal = DC @ ift.sigmoid(correlated_field)

    # Build the line-of-sight response and define signal response
    LOS_starts, LOS_ends = random_los(100) if mode == 0 else radial_los(100)
    R = ift.LOSResponse(position_space, starts=LOS_starts, ends=LOS_ends)
    signal_response = R(signal)

    # Specify noise
    data_space = R.target
    noise = .001
97
    N = ift.ScalingOperator(data_space, noise)
Philipp Frank's avatar
Philipp Frank committed
98
99

    # Generate mock signal and data
100
    mock_position = ift.from_random(signal_response.domain, 'normal')
Philipp Arras's avatar
Philipp Arras committed
101
    data = signal_response(mock_position) + N.draw_sample_with_dtype(dtype=np.float64)
Philipp Frank's avatar
Philipp Frank committed
102

Lukas Platz's avatar
Lukas Platz committed
103
104
105
106
107
108
109
    plot = ift.Plot()
    plot.add(signal(mock_position), title='Ground Truth')
    plot.add(R.adjoint_times(data), title='Data')
    plot.add([A1.force(mock_position)], title='Power Spectrum 1')
    plot.add([A2.force(mock_position)], title='Power Spectrum 2')
    plot.output(ny=2, nx=2, xsize=10, ysize=10, name=filename.format("setup"))

Philipp Frank's avatar
Philipp Frank committed
110
    # Minimization parameters
Philipp Arras's avatar
Philipp Arras committed
111
112
113
114
115
116
    ic_sampling = ift.AbsDeltaEnergyController(name='Sampling',
                                               deltaE=0.01,
                                               iteration_limit=100)
    ic_newton = ift.AbsDeltaEnergyController(name='Newton',
                                             deltaE=0.01,
                                             iteration_limit=35)
Philipp Arras's avatar
Philipp Arras committed
117
118
119
    ic_sampling.enable_logging()
    ic_newton.enable_logging()
    minimizer = ift.NewtonCG(ic_newton, activate_logging=True)
Philipp Frank's avatar
Philipp Frank committed
120

Lukas Platz's avatar
Lukas Platz committed
121
122
123
    ## number of samples used to estimate the KL
    N_samples = 20

Philipp Frank's avatar
Philipp Frank committed
124
    # Set up likelihood and information Hamiltonian
Philipp Arras's avatar
Fixes    
Philipp Arras committed
125
    likelihood = ift.GaussianEnergy(mean=data, inverse_covariance=N.inverse) @ signal_response
Philipp Frank's avatar
Philipp Frank committed
126
127
    H = ift.StandardHamiltonian(likelihood, ic_sampling)

Lukas Platz's avatar
Lukas Platz committed
128
    # Begin minimization
Philipp Frank's avatar
Philipp Frank committed
129
130
131
132
133
134
135
136
137
138
139
140
141
    initial_mean = ift.MultiField.full(H.domain, 0.)
    mean = initial_mean

    for i in range(10):
        # Draw new samples and minimize KL
        KL = ift.MetricGaussianKL(mean, H, N_samples)
        KL, convergence = minimizer(KL)
        mean = KL.position

        # Plot current reconstruction
        plot = ift.Plot()
        plot.add(signal(mock_position), title="ground truth")
        plot.add(signal(KL.position), title="reconstruction")
Philipp Arras's avatar
Philipp Arras committed
142
143
144
145
146
147
        plot.add([A1.force(KL.position),
                  A1.force(mock_position)],
                 title="power1")
        plot.add([A2.force(KL.position),
                  A2.force(mock_position)],
                 title="power2")
Philipp Arras's avatar
Philipp Arras committed
148
149
150
151
152
153
        plot.add((ic_newton.history, ic_sampling.history,
                  minimizer.inversion_history),
                 label=['KL', 'Sampling', 'Newton inversion'],
                 title='Cumulative energies', s=[None, None, 1],
                 alpha=[None, 0.2, None])
        plot.output(nx=3,
Philipp Arras's avatar
Philipp Arras committed
154
155
                    ny=2,
                    ysize=10,
Philipp Arras's avatar
Philipp Arras committed
156
                    xsize=15,
Philipp Frank's avatar
Philipp Frank committed
157
158
                    name=filename.format("loop_{:02d}".format(i)))

Lukas Platz's avatar
Lukas Platz committed
159
    # Done, draw posterior samples
Philipp Frank's avatar
Philipp Frank committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
    Nsamples = 20
    KL = ift.MetricGaussianKL(mean, H, N_samples)
    sc = ift.StatCalculator()
    scA1 = ift.StatCalculator()
    scA2 = ift.StatCalculator()
    powers1 = []
    powers2 = []
    for sample in KL.samples:
        sc.add(signal(sample + KL.position))
        p1 = A1.force(sample + KL.position)
        p2 = A2.force(sample + KL.position)
        scA1.add(p1)
        powers1.append(p1)
        scA2.add(p2)
        powers2.append(p2)

    # Plotting
    filename_res = filename.format("results")
    plot = ift.Plot()
    plot.add(sc.mean, title="Posterior Mean")
    plot.add(ift.sqrt(sc.var), title="Posterior Standard Deviation")

    powers1 = [A1.force(s + KL.position) for s in KL.samples]
    powers2 = [A2.force(s + KL.position) for s in KL.samples]
Philipp Arras's avatar
Philipp Arras committed
184
185
186
187
188
189
    plot.add(powers1 + [scA1.mean, A1.force(mock_position)],
             title="Sampled Posterior Power Spectrum 1",
             linewidth=[1.]*len(powers1) + [3., 3.])
    plot.add(powers2 + [scA2.mean, A2.force(mock_position)],
             title="Sampled Posterior Power Spectrum 2",
             linewidth=[1.]*len(powers2) + [3., 3.])
Philipp Frank's avatar
Philipp Frank committed
190
191
    plot.output(ny=2, nx=2, xsize=15, ysize=15, name=filename_res)
    print("Saved results as '{}'.".format(filename_res))