cartesian_wiener_filter.py 5.44 KB
Newer Older
Theo Steininger's avatar
Theo Steininger committed
1
2
3
# -*- coding: utf-8 -*-

import numpy as np
Martin Reinecke's avatar
Martin Reinecke committed
4
5
import nifty2go as ift
from nifty2go import plotting
Theo Steininger's avatar
Theo Steininger committed
6
7

if __name__ == "__main__":
Martin Reinecke's avatar
Martin Reinecke committed
8
    signal_to_noise = 1.5 # The signal to noise ratio
Theo Steininger's avatar
Theo Steininger committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

    # Setting up parameters    |\label{code:wf_parameters}|
    correlation_length_1 = 1. # Typical distance over which the field is correlated
    field_variance_1 = 2. # Variance of field in position space

    response_sigma_1 = 0.05 # Smoothing length of response (in same unit as L)

    def power_spectrum_1(k): # note: field_variance**2 = a*k_0/4.
        a = 4 * correlation_length_1 * field_variance_1**2
        return a / (1 + k * correlation_length_1) ** 4.

    # Setting up the geometry |\label{code:wf_geometry}|
    L_1 = 2. # Total side-length of the domain
    N_pixels_1 = 512 # Grid resolution (pixels per axis)

    signal_space_1 = ift.RGSpace([N_pixels_1], distances=L_1/N_pixels_1)
Martin Reinecke's avatar
Martin Reinecke committed
25
    harmonic_space_1 = signal_space_1.get_default_codomain()
Martin Reinecke's avatar
Martin Reinecke committed
26
    fft_1 = ift.FFTOperator(harmonic_space_1, target=signal_space_1)
Martin Reinecke's avatar
Martin Reinecke committed
27
    power_space_1 = ift.PowerSpace(harmonic_space_1)
Theo Steininger's avatar
Theo Steininger committed
28

Martin Reinecke's avatar
Martin Reinecke committed
29
    mock_power_1 = ift.Field(power_space_1, val=power_spectrum_1(power_space_1.kindex))
Theo Steininger's avatar
Theo Steininger committed
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47



    # Setting up parameters    |\label{code:wf_parameters}|
    correlation_length_2 = 1. # Typical distance over which the field is correlated
    field_variance_2 = 2. # Variance of field in position space

    response_sigma_2 = 0.01 # Smoothing length of response (in same unit as L)

    def power_spectrum_2(k): # note: field_variance**2 = a*k_0/4.
        a = 4 * correlation_length_2 * field_variance_2**2
        return a / (1 + k * correlation_length_2) ** 2.5

    # Setting up the geometry |\label{code:wf_geometry}|
    L_2 = 2. # Total side-length of the domain
    N_pixels_2 = 512 # Grid resolution (pixels per axis)

    signal_space_2 = ift.RGSpace([N_pixels_2], distances=L_2/N_pixels_2)
Martin Reinecke's avatar
Martin Reinecke committed
48
    harmonic_space_2 = signal_space_2.get_default_codomain()
Martin Reinecke's avatar
Martin Reinecke committed
49
50
    fft_2 = ift.FFTOperator(harmonic_space_2, target=signal_space_2)
    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.Field(power_space_2, val=power_spectrum_2(power_space_2.kindex))
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=np.outer(mock_power_1.val, mock_power_2.val))
Theo Steininger's avatar
Theo Steininger committed
58
59

    diagonal = mock_power.power_synthesize(spaces=(0, 1), mean=1, std=0,
Martin Reinecke's avatar
Martin Reinecke committed
60
                                           real_signal=False)**2
Martin Reinecke's avatar
Martin Reinecke committed
61
    diagonal = diagonal.real
Theo Steininger's avatar
Theo Steininger committed
62
63
64
65
66
67

    S = ift.DiagonalOperator(domain=(harmonic_space_1, harmonic_space_2),
                             diagonal=diagonal)


    np.random.seed(10)
Martin Reinecke's avatar
Martin Reinecke committed
68
    mock_signal = fft(mock_power.power_synthesize(real_signal=True))
Theo Steininger's avatar
Theo Steininger committed
69
70
71

    # Setting up a exemplary response
    N1_10 = int(N_pixels_1/10)
Martin Reinecke's avatar
Martin Reinecke committed
72
    mask_1 = ift.Field(signal_space_1, val=1.)
Theo Steininger's avatar
Theo Steininger committed
73
74
75
    mask_1.val[N1_10*7:N1_10*9] = 0.

    N2_10 = int(N_pixels_2/10)
Martin Reinecke's avatar
Martin Reinecke committed
76
    mask_2 = ift.Field(signal_space_2, val=1.)
Theo Steininger's avatar
Theo Steininger committed
77
78
79
80
81
82
    mask_2.val[N2_10*7:N2_10*9] = 0.

    R = ift.ResponseOperator((signal_space_1, signal_space_2),
                             sigma=(response_sigma_1, response_sigma_2),
                             exposure=(mask_1, mask_2)) #|\label{code:wf_response}|
    data_domain = R.target
Theo Steininger's avatar
Theo Steininger committed
83
    R_harmonic = ift.ComposedOperator([fft, R], default_spaces=(0, 1, 0, 1))
Theo Steininger's avatar
Theo Steininger committed
84
85
86

    # Setting up the noise covariance and drawing a random noise realization
    N = ift.DiagonalOperator(data_domain, diagonal=mock_signal.var()/signal_to_noise,
Martin Reinecke's avatar
Martin Reinecke committed
87
                             bare=True)
Theo Steininger's avatar
Theo Steininger committed
88
89
    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
90
                                  mean=0)
Theo Steininger's avatar
Theo Steininger committed
91
92
93
94
    data = R(mock_signal) + noise #|\label{code:wf_mock_data}|

    # Wiener filter
    j = R_harmonic.adjoint_times(N.inverse_times(data))
95
    ctrl = ift.DefaultIterationController(verbose=True, tol_custom=1e-3, convergence_level=3)
96
    inverter = ift.ConjugateGradient(controller=ctrl)
Martin Reinecke's avatar
Martin Reinecke committed
97
    wiener_curvature = ift.library.WienerFilterCurvature(S=S, N=N, R=R_harmonic, inverter=inverter)
Theo Steininger's avatar
Theo Steininger committed
98
99
100
101
102
103

    m_k = wiener_curvature.inverse_times(j) #|\label{code:wf_wiener_filter}|
    m = fft(m_k)

    # Probing the variance
    class Proby(ift.DiagonalProberMixin, ift.Prober): pass
Martin Reinecke's avatar
Martin Reinecke committed
104
    proby = Proby((signal_space_1, signal_space_2), probe_count=2,ncpu=2)
Theo Steininger's avatar
Theo Steininger committed
105
106
107
108
109
110
    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))
111
    sm = ift.FFTSmoothingOperator(plot_space, sigma=0.03)
Martin Reinecke's avatar
Martin Reinecke committed
112
113
114
115
116
    ift.plotting.plot(ift.log(ift.sqrt(sm(ift.Field(plot_space, val=variance.val.real)))), name='uncertainty.pdf',zmin=0.,zmax=3.,title="Uncertainty map",xlabel="x")
    ift.plotting.plot(ift.Field(plot_space, val=mock_signal.val.real), name='mock_signal.pdf')
    ift.plotting.plot(ift.Field(plot_space, val=data.val.real), name='data.pdf')
    ift.plotting.plot(ift.Field(plot_space, val=m.val.real), name='map.pdf')
    exit()
Theo Steininger's avatar
Theo Steininger committed
117
    plotter(ift.log(ift.sqrt(sm(ift.Field(plot_space, val=variance.val.real)))), path='uncertainty.html')
Theo Steininger's avatar
Theo Steininger committed
118
119
120
121

    plotter.plot.zmin = np.real(mock_signal.min());
    plotter.plot.zmax = np.real(mock_signal.max());
    plotter(ift.Field(plot_space, val=mock_signal.val.real), path='mock_signal.html')
Martin Reinecke's avatar
Martin Reinecke committed
122
    plotter(ift.Field(plot_space, val=data.val.real), path = 'data.html')
Theo Steininger's avatar
Theo Steininger committed
123
124
    plotter(ift.Field(plot_space, val=m.val.real), path = 'map.html')