cartesian_wiener_filter.py 5.14 KB
Newer Older
Theo Steininger's avatar
Theo Steininger committed
1
import numpy as np
Martin Reinecke's avatar
Martin Reinecke committed
2
import nifty2go as ift
Theo Steininger's avatar
Theo Steininger committed
3
4

if __name__ == "__main__":
Martin Reinecke's avatar
Martin Reinecke committed
5
    signal_to_noise = 1.5  # The signal to noise ratio
Theo Steininger's avatar
Theo Steininger committed
6

Martin Reinecke's avatar
Martin Reinecke committed
7
8
9
    # Setting up parameters
    correlation_length_1 = 1.  # Typical distance over which the field is correlated
    field_variance_1 = 2.  # Variance of field in position space
Theo Steininger's avatar
Theo Steininger committed
10

Martin Reinecke's avatar
Martin Reinecke committed
11
    response_sigma_1 = 0.05  # Smoothing length of response (in same unit as L)
Theo Steininger's avatar
Theo Steininger committed
12

Martin Reinecke's avatar
Martin Reinecke committed
13
    def power_spectrum_1(k):  # note: field_variance**2 = a*k_0/4.
Theo Steininger's avatar
Theo Steininger committed
14
15
16
        a = 4 * correlation_length_1 * field_variance_1**2
        return a / (1 + k * correlation_length_1) ** 4.

Martin Reinecke's avatar
Martin Reinecke committed
17
18
19
    # Setting up the geometry
    L_1 = 2.  # Total side-length of the domain
    N_pixels_1 = 512  # Grid resolution (pixels per axis)
Theo Steininger's avatar
Theo Steininger committed
20
21

    signal_space_1 = ift.RGSpace([N_pixels_1], distances=L_1/N_pixels_1)
Martin Reinecke's avatar
Martin Reinecke committed
22
    harmonic_space_1 = signal_space_1.get_default_codomain()
Martin Reinecke's avatar
Martin Reinecke committed
23
24
25
    # Setting up the geometry
    L_2 = 2.  # Total side-length of the domain
    N_pixels_2 = 512  # Grid resolution (pixels per axis)
26
27
28
29
30
    signal_space_2 = ift.RGSpace([N_pixels_2], distances=L_2/N_pixels_2)
    harmonic_space_2 = signal_space_2.get_default_codomain()

    signal_domain = ift.DomainTuple.make((signal_space_1, signal_space_2))
    mid_domain = ift.DomainTuple.make((signal_space_1, harmonic_space_2))
Martin Reinecke's avatar
Martin Reinecke committed
31
32
    harmonic_domain = ift.DomainTuple.make((harmonic_space_1,
                                            harmonic_space_2))
33
34

    fft_1 = ift.FFTOperator(harmonic_domain, space=0)
Martin Reinecke's avatar
Martin Reinecke committed
35
    power_space_1 = ift.PowerSpace(harmonic_space_1)
Theo Steininger's avatar
Theo Steininger committed
36

Martin Reinecke's avatar
Martin Reinecke committed
37
    mock_power_1 = ift.PS_field(power_space_1, power_spectrum_1)
Theo Steininger's avatar
Theo Steininger committed
38

Martin Reinecke's avatar
Martin Reinecke committed
39
40
41
    # Setting up parameters
    correlation_length_2 = 1.  # Typical distance over which the field is correlated
    field_variance_2 = 2.  # Variance of field in position space
Theo Steininger's avatar
Theo Steininger committed
42

Martin Reinecke's avatar
Martin Reinecke committed
43
    response_sigma_2 = 0.01  # Smoothing length of response (in same unit as L)
Theo Steininger's avatar
Theo Steininger committed
44

Martin Reinecke's avatar
Martin Reinecke committed
45
    def power_spectrum_2(k):  # note: field_variance**2 = a*k_0/4.
Theo Steininger's avatar
Theo Steininger committed
46
47
48
        a = 4 * correlation_length_2 * field_variance_2**2
        return a / (1 + k * correlation_length_2) ** 2.5

49
    fft_2 = ift.FFTOperator(mid_domain, space=1)
Martin Reinecke's avatar
Martin Reinecke committed
50
    power_space_2 = ift.PowerSpace(harmonic_space_2)
Theo Steininger's avatar
Theo Steininger committed
51

Martin Reinecke's avatar
Martin Reinecke committed
52
    mock_power_2 = ift.PS_field(power_space_2, power_spectrum_2)
Theo Steininger's avatar
Theo Steininger committed
53
54
55
56

    fft = ift.ComposedOperator((fft_1, fft_2))

    mock_power = ift.Field(domain=(power_space_1, power_space_2),
Martin Reinecke's avatar
Martin Reinecke committed
57
                           val=ift.dobj.from_global_data(np.outer(ift.dobj.to_global_data(mock_power_1.val), ift.dobj.to_global_data(mock_power_2.val))))
Theo Steininger's avatar
Theo Steininger committed
58

Martin Reinecke's avatar
adjust    
Martin Reinecke committed
59
    diagonal = ift.power_synthesize_special(mock_power, spaces=(0, 1))**2
Martin Reinecke's avatar
Martin Reinecke committed
60
    diagonal = diagonal.real
Theo Steininger's avatar
Theo Steininger committed
61

62
    S = ift.DiagonalOperator(diagonal)
Theo Steininger's avatar
Theo Steininger committed
63
64

    np.random.seed(10)
Martin Reinecke's avatar
adjust    
Martin Reinecke committed
65
    mock_signal = fft(ift.power_synthesize(mock_power, real_signal=True))
Theo Steininger's avatar
Theo Steininger committed
66
67
68

    # Setting up a exemplary response
    N1_10 = int(N_pixels_1/10)
Martin Reinecke's avatar
Martin Reinecke committed
69
70
71
    mask_1 = np.ones(signal_space_1.shape)
    mask_1[N1_10*7:N1_10*9] = 0.
    mask_1 = ift.Field(signal_space_1, ift.dobj.from_global_data(mask_1))
Theo Steininger's avatar
Theo Steininger committed
72
73

    N2_10 = int(N_pixels_2/10)
Martin Reinecke's avatar
Martin Reinecke committed
74
75
76
    mask_2 = np.ones(signal_space_2.shape)
    mask_2[N2_10*7:N2_10*9] = 0.
    mask_2 = ift.Field(signal_space_2, ift.dobj.from_global_data(mask_2))
Theo Steininger's avatar
Theo Steininger committed
77

Martin Reinecke's avatar
Martin Reinecke committed
78
    R = ift.ResponseOperator(signal_domain, spaces=(0, 1),
Theo Steininger's avatar
Theo Steininger committed
79
                             sigma=(response_sigma_1, response_sigma_2),
Martin Reinecke's avatar
Martin Reinecke committed
80
                             exposure=(mask_1, mask_2))
Theo Steininger's avatar
Theo Steininger committed
81
    data_domain = R.target
82
    R_harmonic = ift.ComposedOperator([fft, R])
Theo Steininger's avatar
Theo Steininger committed
83
84

    # Setting up the noise covariance and drawing a random noise realization
85
    ndiag = ift.Field.full(data_domain, mock_signal.var()/signal_to_noise)
86
    N = ift.DiagonalOperator(ndiag.weight(1))
Theo Steininger's avatar
Theo Steininger committed
87
88
    noise = ift.Field.from_random(domain=data_domain, random_type='normal',
                                  std=mock_signal.std()/np.sqrt(signal_to_noise),
Martin Reinecke's avatar
Martin Reinecke committed
89
                                  mean=0)
Martin Reinecke's avatar
Martin Reinecke committed
90
    data = R(mock_signal) + noise
Theo Steininger's avatar
Theo Steininger committed
91
92
93

    # Wiener filter
    j = R_harmonic.adjoint_times(N.inverse_times(data))
94
    ctrl = ift.GradientNormController(verbose=True, tol_abs_gradnorm=0.1)
95
    inverter = ift.ConjugateGradient(controller=ctrl)
Martin Reinecke's avatar
Martin Reinecke committed
96
97
    wiener_curvature = ift.library.WienerFilterCurvature(S=S, N=N,
                                                         R=R_harmonic)
Martin Reinecke's avatar
adjust    
Martin Reinecke committed
98
    wiener_curvature = ift.InversionEnabler(wiener_curvature, inverter)
Theo Steininger's avatar
Theo Steininger committed
99

Martin Reinecke's avatar
Martin Reinecke committed
100
    m_k = wiener_curvature.inverse_times(j)
Theo Steininger's avatar
Theo Steininger committed
101
102
103
    m = fft(m_k)

    # Probing the variance
Martin Reinecke's avatar
Martin Reinecke committed
104
105
106
    class Proby(ift.DiagonalProberMixin, ift.Prober):
        pass
    proby = Proby((signal_space_1, signal_space_2), probe_count=1, ncpu=1)
Theo Steininger's avatar
Theo Steininger committed
107
108
109
110
111
112
    proby(lambda z: fft(wiener_curvature.inverse_times(fft.inverse_times(z))))
#    sm = SmoothingOperator(signal_space, sigma=0.02)
#    variance = sm(proby.diagonal.weight(-1))
    variance = proby.diagonal.weight(-1)

    plot_space = ift.RGSpace((N_pixels_1, N_pixels_2))
113
    sm = ift.FFTSmoothingOperator(plot_space, sigma=0.03)
Martin Reinecke's avatar
Martin Reinecke committed
114
115
116
117
    ift.plotting.plot(ift.log(ift.sqrt(sm(ift.Field(plot_space, val=variance.val.real)))), name='uncertainty.png',zmin=0.,zmax=3.,title="Uncertainty map",colormap="Planck-like")
    ift.plotting.plot(ift.Field(plot_space, val=mock_signal.val.real), name='mock_signal.png',colormap="Planck-like")
    ift.plotting.plot(ift.Field(plot_space, val=data.val.real), name='data.png',colormap="Planck-like")
    ift.plotting.plot(ift.Field(plot_space, val=m.val.real), name='map.png',colormap="Planck-like")