getting_started_1.py 4.09 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
# 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
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.

19 20
import nifty5 as ift
import numpy as np
21 22


Philipp Arras's avatar
Philipp Arras committed
23
def make_chess_mask(position_space):
24 25 26
    mask = np.ones(position_space.shape)
    for i in range(4):
        for j in range(4):
27
            if (i+j) % 2 == 0:
Philipp Arras's avatar
Philipp Arras committed
28
                mask[i*128//4:(i+1)*128//4, j*128//4:(j+1)*128//4] = 0
29 30
    return mask

Philipp Arras's avatar
Philipp Arras committed
31

32
def make_random_mask():
33
    mask = ift.from_random('pm1', position_space)
34
    mask = (mask+1)/2
Martin Reinecke's avatar
Martin Reinecke committed
35
    return mask.to_global_data()
36

Philipp Arras's avatar
Philipp Arras committed
37

Philipp Arras's avatar
Philipp Arras committed
38 39 40 41 42 43
def mask_to_nan(mask, field):
    masked_data = field.local_data.copy()
    masked_data[mask.local_data == 0] = np.nan
    return ift.from_local_data(field.domain, masked_data)


44
if __name__ == '__main__':
Philipp Arras's avatar
Philipp Arras committed
45
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
46 47
    # FIXME description of the tutorial

48
    # Choose problem geometry and masking
Martin Reinecke's avatar
Martin Reinecke committed
49
    mode = 0
Philipp Arras's avatar
Philipp Arras committed
50 51 52 53 54 55 56 57 58 59 60 61
    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()
62

63 64
    harmonic_space = position_space.get_default_codomain()
    HT = ift.HarmonicTransformOperator(harmonic_space, target=position_space)
65

Philipp Arras's avatar
Philipp Arras committed
66
    # Set correlation structure with a power spectrum and build
67
    # prior correlation covariance
68 69 70 71 72
    def power_spectrum(k):
        return 100. / (20.+k**3)
    power_space = ift.PowerSpace(harmonic_space)
    PD = ift.PowerDistributor(harmonic_space, power_space)
    prior_correlation_structure = PD(ift.PS_field(power_space, power_spectrum))
73

74
    S = ift.DiagonalOperator(prior_correlation_structure)
75

Philipp Arras's avatar
Philipp Arras committed
76
    # Build instrument response consisting of a discretization, mask
77
    # and harmonic transformaion
78
    GR = ift.GeometryRemover(position_space)
79
    mask = ift.Field.from_global_data(position_space, mask)
80 81 82 83 84
    Mask = ift.DiagonalOperator(mask)
    R = GR * Mask * HT

    data_space = GR.target

Philipp Arras's avatar
Philipp Arras committed
85
    # Set the noise covariance
86 87
    noise = 5.
    N = ift.ScalingOperator(noise, data_space)
88

Philipp Arras's avatar
Philipp Arras committed
89
    # Create mock data
90 91 92
    MOCK_SIGNAL = S.draw_sample()
    MOCK_NOISE = N.draw_sample()
    data = R(MOCK_SIGNAL) + MOCK_NOISE
93

Philipp Arras's avatar
Philipp Arras committed
94
    # Build propagator D and information source j
95 96
    j = R.adjoint_times(N.inverse_times(data))
    D_inv = R.adjoint * N.inverse * R + S.inverse
Philipp Arras's avatar
Philipp Arras committed
97
    # Make it invertible
98
    IC = ift.GradientNormController(iteration_limit=500, tol_abs_gradnorm=1e-3)
99
    D = ift.InversionEnabler(D_inv, IC, approximation=S.inverse).inverse
100 101 102 103

    # WIENER FILTER
    m = D(j)

104
    # PLOTTING
Philipp Arras's avatar
Philipp Arras committed
105 106
    rg = isinstance(position_space, ift.RGSpace)
    if rg and len(position_space.shape) == 1:
Martin Reinecke's avatar
Martin Reinecke committed
107
        ift.plot([HT(MOCK_SIGNAL), GR.adjoint(data), HT(m)],
Philipp Arras's avatar
Philipp Arras committed
108
                 label=['Mock signal', 'Data', 'Reconstruction'],
Martin Reinecke's avatar
Martin Reinecke committed
109 110 111
                 alpha=[1, .3, 1])
        ift.plot(mask_to_nan(mask, HT(m-MOCK_SIGNAL)))
        ift.plot_finish(1, 2, xsize=10, ysize=4, title="getting_started_1")
Philipp Arras's avatar
Philipp Arras committed
112
    else:
Martin Reinecke's avatar
Martin Reinecke committed
113 114 115 116 117
        ift.plot(HT(MOCK_SIGNAL), title='Mock Signal')
        ift.plot(mask_to_nan(mask, (GR*Mask).adjoint(data)), title='Data')
        ift.plot(HT(m), title='Reconstruction')
        ift.plot(mask_to_nan(mask, HT(m-MOCK_SIGNAL)))
        ift.plot_finish(2, 2, xsize=10, ysize=8, title="getting_started_1")