getting_started_1.py 5.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 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/>.
#
# Copyright(C) 2013-2018 Max-Planck-Society
#
16
17
18
19
20
21
22
23
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.


############################################################
# Wiener filter demo code 
# showing how measurement gaps are filled in 
# 1D (set mode=0), 2D (mode=1), or on the sphere (mode=2)
#############################################################
24

25
26
import nifty5 as ift
import numpy as np
27
28


Philipp Arras's avatar
Philipp Arras committed
29
def make_chess_mask(position_space):
30
    # set up a checkerboard mask for 2D mode
31
32
33
    mask = np.ones(position_space.shape)
    for i in range(4):
        for j in range(4):
34
            if (i+j) % 2 == 0:
Philipp Arras's avatar
Philipp Arras committed
35
                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():
40
    # set up random mask for spherical mode
41
    mask = ift.from_random('pm1', position_space)
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):
47
    # mask masked pixels for plotting data
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__':
Philipp Arras's avatar
Philipp Arras committed
54
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
55

56
    # Choose position space of signal and masking in measurement
Martin Reinecke's avatar
Martin Reinecke committed
57
    mode = 1
Philipp Arras's avatar
Philipp Arras committed
58
59
60
61
62
63
64
65
66
67
68
69
    if mode == 0:
        # One dimensional regular grid
        position_space = ift.RGSpace([1024])
        mask = np.ones(position_space.shape)
    elif mode == 1:
        # Two dimensional regular grid with chess mask
        position_space = ift.RGSpace([128, 128])
        mask = make_chess_mask(position_space)
    else:
        # Sphere with half of its locations randomly masked
        position_space = ift.HPSpace(128)
        mask = make_random_mask()
70

71
72
    # specify harmonic space corresponding to signal
    # the signal degree of freedom will live here
73
    harmonic_space = position_space.get_default_codomain()
74
75
 
    # harmonic transform from harmonic space to position space
76
    HT = ift.HarmonicTransformOperator(harmonic_space, target=position_space)
77

Philipp Arras's avatar
Philipp Arras committed
78
    # Set correlation structure with a power spectrum and build
79
    # prior correlation covariance
80
81

    # define 1D isotropic power spectrum
82
83
    def power_spectrum(k):
        return 100. / (20.+k**3)
84
85
86
87
88
    # 1D spectral space in which power spectrum lives
    power_space = ift.PowerSpace(harmonic_space) 
    # its mapping to (higher dimentional) harmonic space
    PD = ift.PowerDistributor(harmonic_space, power_space) 
    # applying the mapping to the 1D spectrum
89
    prior_correlation_structure = PD(ift.PS_field(power_space, power_spectrum))
90
    # inserting the result into the diagonal of an harmonic space operator
91
    S = ift.DiagonalOperator(prior_correlation_structure)
92
    # S is now the 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

    # data lives in geometry free space, thus we need to remove geometry
98
    GR = ift.GeometryRemover(position_space)
99
    # masking operator, to model that parts of the field where not observed
100
    mask = ift.Field.from_global_data(position_space, mask)
101
    Mask = ift.DiagonalOperator(mask)
102
103
104
105
    # compile response operator out of 
    # - an harmic transform (to get to image space)
    # - the application of the mask
    # - the removal of geometric information
Martin Reinecke's avatar
Martin Reinecke committed
106
    R = GR(Mask(HT))
107

108
    # find out where the data then lives
109
110
    data_space = GR.target

111
    # Set the noise covariance N
112
113
    noise = 5.
    N = ift.ScalingOperator(noise, data_space)
114

Philipp Arras's avatar
Philipp Arras committed
115
    # Create mock data
116
117
118
    MOCK_SIGNAL = S.draw_sample()
    MOCK_NOISE = N.draw_sample()
    data = R(MOCK_SIGNAL) + MOCK_NOISE
119

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

    # WIENER FILTER
128
    # m = D j
129
130
    m = D(j)

131
    # PLOTTING
Philipp Arras's avatar
Philipp Arras committed
132
    rg = isinstance(position_space, ift.RGSpace)
133
    plot = ift.Plot()
Philipp Arras's avatar
Philipp Arras committed
134
    if rg and len(position_space.shape) == 1:
135
        plot.add([HT(MOCK_SIGNAL), GR.adjoint(data), HT(m)],
Philipp Arras's avatar
Philipp Arras committed
136
                 label=['Mock signal', 'Data', 'Reconstruction'],
Martin Reinecke's avatar
Martin Reinecke committed
137
                 alpha=[1, .3, 1])
138
139
        plot.add(mask_to_nan(mask, HT(m-MOCK_SIGNAL)), title='Residuals')
        plot.output(nx=2, ny=1, xsize=10, ysize=4, title="getting_started_1")
Philipp Arras's avatar
Philipp Arras committed
140
    else:
141
142
        plot.add(HT(MOCK_SIGNAL), title='Mock Signal')
        plot.add(mask_to_nan(mask, (GR(Mask)).adjoint(data)),
143
                 title='Data')
144
145
146
        plot.add(HT(m), title='Reconstruction')
        plot.add(mask_to_nan(mask, HT(m-MOCK_SIGNAL)), title='Residuals')
        plot.output(nx=2, ny=2, xsize=10, ysize=10, title="getting_started_1")