getting_started_3.py 5.67 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
# 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/>.
#
14
# Copyright(C) 2013-2019 Max-Planck-Society
15
#
16 17 18
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.

############################################################
Philipp Arras's avatar
Philipp Arras committed
19
# Non-linear tomography
Torsten Ensslin's avatar
Torsten Ensslin committed
20 21 22 23 24 25
#
# 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
26
#############################################################
27

Jakob Knollmueller's avatar
Jakob Knollmueller committed
28 29
import numpy as np

Philipp Arras's avatar
Philipp Arras committed
30 31
import nifty5 as ift

Jakob Knollmueller's avatar
Jakob Knollmueller committed
32

Philipp Arras's avatar
Philipp Arras committed
33
def random_los(n_los):
34
    starts = list(np.random.uniform(0, 1, (n_los, 2)).T)
Torsten Ensslin's avatar
Torsten Ensslin committed
35
    ends = list(np.random.uniform(0, 1, (n_los, 2)).T)
Philipp Arras's avatar
Philipp Arras committed
36 37 38 39 40
    return starts, ends


def radial_los(n_los):
    starts = list(np.random.uniform(0, 1, (n_los, 2)).T)
Torsten Ensslin's avatar
Torsten Ensslin committed
41
    ends = list(0.5 + 0*np.random.uniform(0, 1, (n_los, 2)).T)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
42 43
    return starts, ends

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
44

Jakob Knollmueller's avatar
Jakob Knollmueller committed
45
if __name__ == '__main__':
Lukas Platz's avatar
Lukas Platz committed
46
    import sys
Torsten Ensslin's avatar
Torsten Ensslin committed
47
    np.random.seed(420)  # picked for a nice field realization
Philipp Arras's avatar
Philipp Arras committed
48

Torsten Ensslin's avatar
Torsten Ensslin committed
49 50
    # Choose between random line-of-sight response (mode=0) and radial lines
    # of sight (mode=1)
Lukas Platz's avatar
Lukas Platz committed
51 52 53 54 55
    if len(sys.argv) == 2:
        mode = int(sys.argv[1])
    else:
        mode = 0
    filename = f"getting_started_3_mode_{mode}_" + "{}.png"
Philipp Arras's avatar
Philipp Arras committed
56

57
    position_space = ift.RGSpace([128, 128])
58 59 60
    harmonic_space = position_space.get_default_codomain()
    ht = ift.HarmonicTransformOperator(harmonic_space, position_space)
    power_space = ift.PowerSpace(harmonic_space)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
61

Philipp Arras's avatar
Philipp Arras committed
62
    # Set up an amplitude operator for the field
63 64 65 66 67 68 69 70 71 72 73
    dct = {
        'target': power_space,
        'n_pix': 64,  # 64 spectral bins

        # Spectral smoothness (affects Gaussian process part)
        'a': 3,  # relatively high variance of spectral curbvature
        'k0': .4,  # quefrency mode below which cepstrum flattens

        # Power-law part of spectrum:
        'sm': -5,  # preferred power-law slope
        'sv': .5,  # low variance of power-law slope
Torsten Ensslin's avatar
Torsten Ensslin committed
74 75
        'im':  0,  # y-intercept mean, in-/decrease for more/less contrast
        'iv': .3   # y-intercept variance
76
    }
77
    A = ift.SLAmplitude(**dct)
Philipp Arras's avatar
Philipp Arras committed
78

Philipp Arras's avatar
Philipp Arras committed
79
    # Build the operator for a correlated signal
Jakob Knollmueller's avatar
Jakob Knollmueller committed
80
    power_distributor = ift.PowerDistributor(harmonic_space, power_space)
81 82 83
    vol = harmonic_space.scalar_dvol**-0.5
    xi = ift.ducktape(harmonic_space, None, 'xi')
    correlated_field = ht(vol*power_distributor(A)*xi)
Philipp Arras's avatar
Philipp Arras committed
84 85
    # Alternatively, one can use:
    # correlated_field = ift.CorrelatedField(position_space, A)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
86

Philipp Arras's avatar
Philipp Arras committed
87
    # Apply a nonlinearity
Jakob Knollmueller's avatar
Jakob Knollmueller committed
88
    signal = ift.sigmoid(correlated_field)
Martin Reinecke's avatar
Martin Reinecke committed
89

Philipp Arras's avatar
Philipp Arras committed
90
    # Build the line-of-sight response and define signal response
Torsten Ensslin's avatar
Torsten Ensslin committed
91
    LOS_starts, LOS_ends = random_los(100) if mode == 0 else radial_los(100)
Philipp Arras's avatar
Philipp Arras committed
92
    R = ift.LOSResponse(position_space, starts=LOS_starts, ends=LOS_ends)
Martin Reinecke's avatar
Martin Reinecke committed
93
    signal_response = R(signal)
Philipp Arras's avatar
Philipp Arras committed
94 95

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

Philipp Arras's avatar
Philipp Arras committed
100 101 102
    # Generate mock signal and data
    mock_position = ift.from_random('normal', signal_response.domain)
    data = signal_response(mock_position) + N.draw_sample()
Jakob Knollmueller's avatar
Jakob Knollmueller committed
103

Philipp Arras's avatar
Philipp Arras committed
104
    # Minimization parameters
Jakob Knollmueller's avatar
Jakob Knollmueller committed
105
    ic_sampling = ift.GradientNormController(iteration_limit=100)
Martin Reinecke's avatar
Martin Reinecke committed
106
    ic_newton = ift.GradInfNormController(
107
        name='Newton', tol=1e-7, iteration_limit=35)
108
    minimizer = ift.NewtonCG(ic_newton)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
109

Philipp Arras's avatar
Philipp Arras committed
110
    # Set up likelihood and information Hamiltonian
Philipp Arras's avatar
Philipp Arras committed
111
    likelihood = ift.GaussianEnergy(mean=data, covariance=N)(signal_response)
112
    H = ift.StandardHamiltonian(likelihood, ic_sampling)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
113

114 115
    initial_mean = ift.MultiField.full(H.domain, 0.)
    mean = initial_mean
Jakob Knollmueller's avatar
Jakob Knollmueller committed
116

117
    plot = ift.Plot()
Philipp Arras's avatar
Philipp Arras committed
118
    plot.add(signal(mock_position), title='Ground Truth')
Martin Reinecke's avatar
merge  
Martin Reinecke committed
119
    plot.add(R.adjoint_times(data), title='Data')
Philipp Arras's avatar
Philipp Arras committed
120
    plot.add([A.force(mock_position)], title='Power Spectrum')
Lukas Platz's avatar
Lukas Platz committed
121
    plot.output(ny=1, nx=3, xsize=24, ysize=6, name=filename.format("setup"))
Jakob Knollmueller's avatar
Jakob Knollmueller committed
122

Jakob Knollmueller's avatar
Jakob Knollmueller committed
123
    # number of samples used to estimate the KL
124
    N_samples = 20
Philipp Arras's avatar
Philipp Arras committed
125 126

    # Draw new samples to approximate the KL five times
127
    for i in range(5):
Philipp Arras's avatar
Philipp Arras committed
128
        # Draw new samples and minimize KL
129
        KL = ift.MetricGaussianKL(mean, H, N_samples)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
130
        KL, convergence = minimizer(KL)
131
        mean = KL.position
Philipp Arras's avatar
Philipp Arras committed
132 133

        # Plot current reconstruction
134
        plot = ift.Plot()
Martin Reinecke's avatar
merge  
Martin Reinecke committed
135
        plot.add(signal(KL.position), title="reconstruction")
Philipp Arras's avatar
Philipp Arras committed
136
        plot.add([A.force(KL.position), A.force(mock_position)], title="power")
Lukas Platz's avatar
Lukas Platz committed
137 138
        plot.output(ny=1, ysize=6, xsize=16,
                    name=filename.format(f"loop_{i:02}"))
Jakob Knollmueller's avatar
Jakob Knollmueller committed
139

Philipp Arras's avatar
Philipp Arras committed
140
    # Draw posterior samples
141
    KL = ift.MetricGaussianKL(mean, H, N_samples)
Martin Reinecke's avatar
Martin Reinecke committed
142
    sc = ift.StatCalculator()
143
    for sample in KL.samples:
Philipp Arras's avatar
Philipp Arras committed
144
        sc.add(signal(sample + KL.position))
Philipp Arras's avatar
Philipp Arras committed
145 146

    # Plotting
Lukas Platz's avatar
Lukas Platz committed
147
    filename_res = filename.format("results")
Philipp Arras's avatar
Philipp Arras committed
148
    plot = ift.Plot()
Martin Reinecke's avatar
merge  
Martin Reinecke committed
149 150
    plot.add(sc.mean, title="Posterior Mean")
    plot.add(ift.sqrt(sc.var), title="Posterior Standard Deviation")
Martin Reinecke's avatar
Martin Reinecke committed
151

Philipp Arras's avatar
Philipp Arras committed
152
    powers = [A.force(s + KL.position) for s in KL.samples]
Martin Reinecke's avatar
merge  
Martin Reinecke committed
153
    plot.add(
Philipp Arras's avatar
Philipp Arras committed
154 155
        powers + [A.force(KL.position),
                  A.force(mock_position)],
Lukas Platz's avatar
Lukas Platz committed
156 157
        title="Sampled Posterior Power Spectrum",
        linewidth=[1.]*len(powers) + [3., 3.])
Lukas Platz's avatar
Lukas Platz committed
158 159
    plot.output(ny=1, nx=3, xsize=24, ysize=6, name=filename_res)
    print(f"Saved results as '{filename_res}'.")