getting_started_1.py 5.36 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/>.
#
Philipp Arras's avatar
Philipp Arras committed
14
# Copyright(C) 2013-2020 Max-Planck-Society
15
#
16
17
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.

Philipp Arras's avatar
Philipp Arras committed
18
19
20
###############################################################################
# Compute a Wiener filter solution with NIFTy
# Shows how measurement gaps are filled in
21
# 1D (set mode=0), 2D (mode=1), or on the sphere (mode=2)
Philipp Arras's avatar
Philipp Arras committed
22
###############################################################################
23

Philipp Arras's avatar
Philipp Arras committed
24
25
import sys

26
import numpy as np
27

Martin Reinecke's avatar
Martin Reinecke committed
28
import nifty7 as ift
Philipp Arras's avatar
Philipp Arras committed
29

30

Philipp Arras's avatar
Philipp Arras committed
31
32
def make_checkerboard_mask(position_space):
    # Checkerboard mask for 2D mode
33
34
35
    mask = np.ones(position_space.shape)
    for i in range(4):
        for j in range(4):
Philipp Arras's avatar
Philipp Arras committed
36
37
            if (i + j) % 2 == 0:
                mask[i*128//4:(i + 1)*128//4, j*128//4:(j + 1)*128//4] = 0
38
39
    return mask

Philipp Arras's avatar
Philipp Arras committed
40

Philipp Arras's avatar
Philipp Arras committed
41
def make_random_mask(domain):
Philipp Arras's avatar
Philipp Arras committed
42
    # Random mask for spherical mode
Philipp Arras's avatar
Philipp Arras committed
43
    mask = ift.from_random(domain, 'pm1')
Philipp Arras's avatar
Philipp Arras committed
44
    mask = (mask + 1)/2
Martin Reinecke's avatar
Martin Reinecke committed
45
    return mask.val
46

Philipp Arras's avatar
Philipp Arras committed
47

Philipp Arras's avatar
Philipp Arras committed
48
def main():
Philipp Arras's avatar
Philipp Arras committed
49
    # Choose space on which the signal field is defined
Lukas Platz's avatar
Lukas Platz committed
50
51
52
53
54
    if len(sys.argv) == 2:
        mode = int(sys.argv[1])
    else:
        mode = 1

Philipp Arras's avatar
Philipp Arras committed
55
    if mode == 0:
Philipp Arras's avatar
Philipp Arras committed
56
        # One-dimensional regular grid
Philipp Arras's avatar
Philipp Arras committed
57
        position_space = ift.RGSpace([1024])
58
        mask = np.zeros(position_space.shape)
Philipp Arras's avatar
Philipp Arras committed
59
    elif mode == 1:
Philipp Arras's avatar
Philipp Arras committed
60
        # Two-dimensional regular grid with checkerboard mask
Philipp Arras's avatar
Philipp Arras committed
61
        position_space = ift.RGSpace([128, 128])
Philipp Arras's avatar
Philipp Arras committed
62
        mask = make_checkerboard_mask(position_space)
Philipp Arras's avatar
Philipp Arras committed
63
    else:
Philipp Arras's avatar
Philipp Arras committed
64
        # Sphere with half of its pixels randomly masked
Philipp Arras's avatar
Philipp Arras committed
65
        position_space = ift.HPSpace(128)
Philipp Arras's avatar
Philipp Arras committed
66
        mask = make_random_mask(position_space)
67

Philipp Arras's avatar
Philipp Arras committed
68
    # Specify harmonic space corresponding to signal
69
    harmonic_space = position_space.get_default_codomain()
70

Philipp Arras's avatar
Philipp Arras committed
71
72
    # Harmonic transform from harmonic space to position space
    HT = ift.HarmonicTransformOperator(harmonic_space, target=position_space)
73

Philipp Arras's avatar
Philipp Arras committed
74
75
    # Set prior correlation covariance with a power spectrum leading to
    # homogeneous and isotropic statistics
76
    def power_spectrum(k):
Philipp Arras's avatar
Philipp Arras committed
77
78
79
80
81
82
83
84
85
        return 100./(20. + k**3)

    # 1D spectral space on which the power spectrum is defined
    power_space = ift.PowerSpace(harmonic_space)

    # Mapping to (higher dimensional) harmonic space
    PD = ift.PowerDistributor(harmonic_space, power_space)

    # Apply the mapping
86
    prior_correlation_structure = PD(ift.PS_field(power_space, power_spectrum))
Philipp Arras's avatar
Philipp Arras committed
87
88

    # Insert the result into the diagonal of an harmonic space operator
89
    S = ift.DiagonalOperator(prior_correlation_structure)
Philipp Arras's avatar
Philipp Arras committed
90
    # S is the prior field covariance
91

Philipp Arras's avatar
Philipp Arras committed
92
    # Build instrument response consisting of a discretization, mask
93
    # and harmonic transformaion
94

Philipp Arras's avatar
Philipp Arras committed
95
    # Masking operator to model that parts of the field have not been observed
Martin Reinecke's avatar
Martin Reinecke committed
96
    mask = ift.Field.from_raw(position_space, mask)
97
    Mask = ift.MaskOperator(mask)
Martin Reinecke's avatar
Martin Reinecke committed
98

Martin Reinecke's avatar
fix    
Martin Reinecke committed
99
    # The response operator consists of
Martin Reinecke's avatar
Martin Reinecke committed
100
    # - a harmonic transform (to get to image space)
101
102
    # - the application of the mask
    # - the removal of geometric information
103
104
    # The removal of geometric information is included in the MaskOperator
    # it can also be implemented with a GeometryRemover
Martin Reinecke's avatar
fix    
Martin Reinecke committed
105
    # Operators can be composed either with parenthesis
106
    R = Mask(HT)
107
    # or with @
108
    R = Mask @ HT
109

110
    data_space = R.target
111

112
    # Set the noise covariance N
113
    noise = 5.
114
    N = ift.ScalingOperator(data_space, noise)
115

Philipp Arras's avatar
Philipp Arras committed
116
    # Create mock data
Philipp Arras's avatar
Philipp Arras committed
117
118
    MOCK_SIGNAL = S.draw_sample_with_dtype(dtype=np.float64)
    MOCK_NOISE = N.draw_sample_with_dtype(dtype=np.float64)
119
    data = R(MOCK_SIGNAL) + MOCK_NOISE
120

Philipp Arras's avatar
Philipp Arras committed
121
    # Build inverse propagator D and information source j
Martin Reinecke's avatar
fix    
Martin Reinecke committed
122
    D_inv = R.adjoint @ N.inverse @ R + S.inverse
123
    j = R.adjoint_times(N.inverse_times(data))
Philipp Arras's avatar
Philipp Arras committed
124
    # Make D_inv invertible (via Conjugate Gradient)
125
    IC = ift.GradientNormController(iteration_limit=500, tol_abs_gradnorm=1e-3)
126
    D = ift.InversionEnabler(D_inv, IC, approximation=S.inverse).inverse
127

Philipp Arras's avatar
Philipp Arras committed
128
    # Calculate WIENER FILTER solution
129
130
    m = D(j)

Philipp Arras's avatar
Philipp Arras committed
131
    # Plotting
Philipp Arras's avatar
Philipp Arras committed
132
    rg = isinstance(position_space, ift.RGSpace)
133
    plot = ift.Plot()
Philipp Arras's avatar
Philipp Arras committed
134
    filename = "getting_started_1_mode_{}.png".format(mode)
Philipp Arras's avatar
Philipp Arras committed
135
    if rg and len(position_space.shape) == 1:
Philipp Arras's avatar
Philipp Arras committed
136
        plot.add(
137
            [HT(MOCK_SIGNAL), Mask.adjoint(data),
Philipp Arras's avatar
Philipp Arras committed
138
139
140
             HT(m)],
            label=['Mock signal', 'Data', 'Reconstruction'],
            alpha=[1, .3, 1])
141
        plot.add(Mask.adjoint(Mask(HT(m - MOCK_SIGNAL))), title='Residuals')
Lukas Platz's avatar
Lukas Platz committed
142
        plot.output(nx=2, ny=1, xsize=10, ysize=4, name=filename)
Philipp Arras's avatar
Philipp Arras committed
143
    else:
144
        plot.add(HT(MOCK_SIGNAL), title='Mock Signal')
145
        plot.add(Mask.adjoint(data), title='Data')
146
        plot.add(HT(m), title='Reconstruction')
Martin Reinecke's avatar
Martin Reinecke committed
147
148
        plot.add(Mask.adjoint(Mask(HT(m) - HT(MOCK_SIGNAL))),
                 title='Residuals')
Lukas Platz's avatar
Lukas Platz committed
149
        plot.output(nx=2, ny=2, xsize=10, ysize=10, name=filename)
Philipp Arras's avatar
Philipp Arras committed
150
    print("Saved results as '{}'.".format(filename))
Philipp Arras's avatar
Philipp Arras committed
151
152
153
154


if __name__ == '__main__':
    main()