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

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

26
import numpy as np
27

Martin Reinecke's avatar
5->6  
Martin Reinecke committed
28
import nifty6 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

41
def make_random_mask():
Philipp Arras's avatar
Philipp Arras committed
42
    # Random mask for spherical mode
43
    mask = ift.from_random('pm1', position_space)
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

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

Philipp Arras's avatar
Philipp Arras committed
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:
Philipp Arras's avatar
Philipp Arras committed
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:
Philipp Arras's avatar
Philipp Arras committed
62
        # Two-dimensional regular grid with checkerboard mask
Philipp Arras's avatar
Philipp Arras committed
63
        position_space = ift.RGSpace([128, 128])
Philipp Arras's avatar
Philipp Arras committed
64
        mask = make_checkerboard_mask(position_space)
Philipp Arras's avatar
Philipp Arras committed
65
    else:
Philipp Arras's avatar
Philipp Arras committed
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

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

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

Philipp Arras's avatar
Philipp Arras committed
76 77
    # Set prior correlation covariance with a power spectrum leading to
    # homogeneous and isotropic statistics
78
    def power_spectrum(k):
Philipp Arras's avatar
Philipp Arras committed
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))
Philipp Arras's avatar
Philipp Arras committed
89 90

    # Insert the result into the diagonal of an harmonic space operator
91
    S = ift.DiagonalOperator(prior_correlation_structure)
Philipp Arras's avatar
Philipp Arras committed
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

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

Martin Reinecke's avatar
fix  
Martin Reinecke committed
101
    # The response operator consists of
Martin Reinecke's avatar
Martin Reinecke committed
102
    # - a 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
fix  
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
    noise = 5.
116
    N = ift.ScalingOperator(data_space, noise)
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

Philipp Arras's avatar
Philipp Arras committed
123
    # Build inverse propagator D and information source j
Martin Reinecke's avatar
fix  
Martin Reinecke committed
124
    D_inv = R.adjoint @ N.inverse @ R + S.inverse
125
    j = R.adjoint_times(N.inverse_times(data))
Philipp Arras's avatar
Philipp Arras committed
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

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

Philipp Arras's avatar
Philipp Arras committed
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:
Philipp Arras's avatar
Philipp Arras committed
138
        plot.add(
139
            [HT(MOCK_SIGNAL), Mask.adjoint(data),
Philipp Arras's avatar
Philipp Arras committed
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')
Martin Reinecke's avatar
Martin Reinecke committed
149 150
        plot.add(Mask.adjoint(Mask(HT(m) - HT(MOCK_SIGNAL))),
                 title='Residuals')
Lukas Platz's avatar
Lukas Platz committed
151
        plot.output(nx=2, ny=2, xsize=10, ysize=10, name=filename)
Philipp Arras's avatar
Philipp Arras committed
152
    print("Saved results as '{}'.".format(filename))