getting_started_1.py 5.53 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
# 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

24
import numpy as np
25

Philipp Arras's avatar
Philipp Arras committed
26
27
import nifty5 as ift

28

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

Philipp Arras's avatar
Philipp Arras committed
38

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

Philipp Arras's avatar
Philipp Arras committed
45

Philipp Arras's avatar
Philipp Arras committed
46
def mask_to_nan(mask, field):
Philipp Arras's avatar
Philipp Arras committed
47
    # Set masked pixels to nan for plotting
Philipp Arras's avatar
Philipp Arras committed
48
49
50
51
52
    masked_data = field.local_data.copy()
    masked_data[mask.local_data == 0] = np.nan
    return ift.from_local_data(field.domain, masked_data)


53
if __name__ == '__main__':
Lukas Platz's avatar
Lukas Platz committed
54
    import sys
Philipp Arras's avatar
Philipp Arras committed
55
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
56

Philipp Arras's avatar
Philipp Arras committed
57
    # Choose space on which the signal field is defined
Lukas Platz's avatar
Lukas Platz committed
58
59
60
61
62
    if len(sys.argv) == 2:
        mode = int(sys.argv[1])
    else:
        mode = 1

Philipp Arras's avatar
Philipp Arras committed
63
    if mode == 0:
Philipp Arras's avatar
Philipp Arras committed
64
        # One-dimensional regular grid
Philipp Arras's avatar
Philipp Arras committed
65
66
67
        position_space = ift.RGSpace([1024])
        mask = np.ones(position_space.shape)
    elif mode == 1:
Philipp Arras's avatar
Philipp Arras committed
68
        # Two-dimensional regular grid with checkerboard mask
Philipp Arras's avatar
Philipp Arras committed
69
        position_space = ift.RGSpace([128, 128])
Philipp Arras's avatar
Philipp Arras committed
70
        mask = make_checkerboard_mask(position_space)
Philipp Arras's avatar
Philipp Arras committed
71
    else:
Philipp Arras's avatar
Philipp Arras committed
72
        # Sphere with half of its pixels randomly masked
Philipp Arras's avatar
Philipp Arras committed
73
74
        position_space = ift.HPSpace(128)
        mask = make_random_mask()
75

Philipp Arras's avatar
Philipp Arras committed
76
    # Specify harmonic space corresponding to signal
77
    harmonic_space = position_space.get_default_codomain()
78

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

Philipp Arras's avatar
Philipp Arras committed
82
83
    # Set prior correlation covariance with a power spectrum leading to
    # homogeneous and isotropic statistics
84
    def power_spectrum(k):
Philipp Arras's avatar
Philipp Arras committed
85
86
87
88
89
90
91
92
93
        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
94
    prior_correlation_structure = PD(ift.PS_field(power_space, power_spectrum))
Philipp Arras's avatar
Philipp Arras committed
95
96

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

Philipp Arras's avatar
Philipp Arras committed
100
    # Build instrument response consisting of a discretization, mask
101
    # and harmonic transformaion
102

Philipp Arras's avatar
Philipp Arras committed
103
    # Data is defined on a geometry-free space, thus the geometry is removed
104
    GR = ift.GeometryRemover(position_space)
Philipp Arras's avatar
Philipp Arras committed
105
106

    # Masking operator to model that parts of the field have not been observed
107
    mask = ift.Field.from_global_data(position_space, mask)
108
    Mask = ift.DiagonalOperator(mask)
Philipp Arras's avatar
Philipp Arras committed
109

Martin Reinecke's avatar
fix    
Martin Reinecke committed
110
    # The response operator consists of
Philipp Arras's avatar
Philipp Arras committed
111
    # - an harmonic transform (to get to image space)
112
113
    # - the application of the mask
    # - the removal of geometric information
Martin Reinecke's avatar
fix    
Martin Reinecke committed
114
    # Operators can be composed either with parenthesis
Martin Reinecke's avatar
Martin Reinecke committed
115
    R = GR(Mask(HT))
116
117
    # or with @
    R = GR @ Mask @ HT
118
119
120

    data_space = GR.target

121
    # Set the noise covariance N
122
123
    noise = 5.
    N = ift.ScalingOperator(noise, data_space)
124

Philipp Arras's avatar
Philipp Arras committed
125
    # Create mock data
126
127
128
    MOCK_SIGNAL = S.draw_sample()
    MOCK_NOISE = N.draw_sample()
    data = R(MOCK_SIGNAL) + MOCK_NOISE
129

Philipp Arras's avatar
Philipp Arras committed
130
    # Build inverse propagator D and information source j
Martin Reinecke's avatar
fix    
Martin Reinecke committed
131
    D_inv = R.adjoint @ N.inverse @ R + S.inverse
132
    j = R.adjoint_times(N.inverse_times(data))
Philipp Arras's avatar
Philipp Arras committed
133
    # Make D_inv invertible (via Conjugate Gradient)
134
    IC = ift.GradientNormController(iteration_limit=500, tol_abs_gradnorm=1e-3)
135
    D = ift.InversionEnabler(D_inv, IC, approximation=S.inverse).inverse
136

Philipp Arras's avatar
Philipp Arras committed
137
    # Calculate WIENER FILTER solution
138
139
    m = D(j)

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