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.

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

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

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.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

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
98
    mask = ift.Field.from_global_data(position_space, mask)
99 100
    Mask = ift.MaskOperator(mask)
    
Martin Reinecke's avatar
fix  
Martin Reinecke committed
101
    # The response operator consists of
Philipp Arras's avatar
Philipp Arras committed
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
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 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

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