getting_started_1.py 5.3 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.

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)
22
###############################################################################
23

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

26
import numpy as np
27

Philipp Arras's avatar
Philipp Arras committed
28 29
import nifty5 as ift

30

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):
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

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

Philipp Arras's avatar
Philipp Arras committed
47

48
if __name__ == '__main__':
Philipp Arras's avatar
Philipp Arras committed
49
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
50

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

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

70
    # Specify harmonic space corresponding to signal
71
    harmonic_space = position_space.get_default_codomain()
72

73 74
    # Harmonic transform from harmonic space to position space
    HT = ift.HarmonicTransformOperator(harmonic_space, target=position_space)
75

76 77
    # Set prior correlation covariance with a power spectrum leading to
    # homogeneous and isotropic statistics
78
    def power_spectrum(k):
79 80 81 82 83 84 85 86 87
        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
88
    prior_correlation_structure = PD(ift.PS_field(power_space, power_spectrum))
89 90

    # Insert the result into the diagonal of an harmonic space operator
91
    S = ift.DiagonalOperator(prior_correlation_structure)
92
    # S is the prior field covariance
93

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

97
    # Masking operator to model that parts of the field have not been observed
98
    mask = ift.Field.from_global_data(position_space, mask)
99 100
    Mask = ift.MaskOperator(mask)
    
Martin Reinecke's avatar
Martin Reinecke committed
101
    # The response operator consists of
102
    # - an harmonic transform (to get to image space)
103 104
    # - the application of the mask
    # - the removal of geometric information
105 106
    # The removal of geometric information is included in the MaskOperator
    # it can also be implemented with a GeometryRemover
Martin Reinecke's avatar
Martin Reinecke committed
107
    # Operators can be composed either with parenthesis
108
    R = Mask(HT)
109
    # or with @
110
    R = Mask @ HT
111

112
    data_space = R.target
113

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

Philipp Arras's avatar
Philipp Arras committed
118
    # Create mock data
119 120 121
    MOCK_SIGNAL = S.draw_sample()
    MOCK_NOISE = N.draw_sample()
    data = R(MOCK_SIGNAL) + MOCK_NOISE
122

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

130
    # Calculate WIENER FILTER solution
131 132
    m = D(j)

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