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

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


55
if __name__ == '__main__':
Philipp Arras's avatar
Philipp Arras committed
56
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
57

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

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

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

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

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

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

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

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

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

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

    data_space = GR.target

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

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

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

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

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