wiener_filter.py 3.08 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
import nifty4 as ift
Theo Steininger's avatar
Theo Steininger committed
2
3
4
5
import numpy as np


if __name__ == "__main__":
Martin Reinecke's avatar
Martin Reinecke committed
6
    # Setting up parameters
Martin Reinecke's avatar
Martin Reinecke committed
7
8
    L = 2.                         # Total side-length of the domain
    N_pixels = 512                 # Grid resolution (pixels per axis)
Martin Reinecke's avatar
Martin Reinecke committed
9
    correlation_length_scale = .2  # Typical distance over which the field is
Martin Reinecke's avatar
Martin Reinecke committed
10
                                   # correlated
Theo Steininger's avatar
Theo Steininger committed
11
    fluctuation_scale = 2.         # Variance of field in position space
Martin Reinecke's avatar
Martin Reinecke committed
12
    response_sigma = 0.05          # Smoothing length of response
Theo Steininger's avatar
Theo Steininger committed
13
14
    signal_to_noise = 1.5          # The signal to noise ratio
    np.random.seed(43)             # Fixing the random seed
Martin Reinecke's avatar
Martin Reinecke committed
15

Theo Steininger's avatar
Theo Steininger committed
16
17
18
19
20
    def power_spectrum(k):         # Defining the power spectrum
        a = 4 * correlation_length_scale * fluctuation_scale**2
        return a / (1 + (k * correlation_length_scale)**2) ** 2

    signal_space = ift.RGSpace([N_pixels, N_pixels], distances=L/N_pixels)
Martin Reinecke's avatar
Martin Reinecke committed
21
    harmonic_space = signal_space.get_default_codomain()
Martin Reinecke's avatar
Martin Reinecke committed
22
    ht = ift.HarmonicTransformOperator(harmonic_space, target=signal_space)
Martin Reinecke's avatar
Martin Reinecke committed
23
24
25
    power_space = ift.PowerSpace(
        harmonic_space, binbounds=ift.PowerSpace.useful_binbounds(
            harmonic_space, logarithmic=True))
Theo Steininger's avatar
Theo Steininger committed
26

Martin Reinecke's avatar
Martin Reinecke committed
27
28
29
    # Creating the mock signal
    S = ift.create_power_operator(harmonic_space,
                                  power_spectrum=power_spectrum)
Martin Reinecke's avatar
Martin Reinecke committed
30
    mock_power = ift.PS_field(power_space, power_spectrum)
31
    mock_signal = ift.power_synthesize(mock_power, real_signal=True)
Theo Steininger's avatar
Theo Steininger committed
32
33

    # Setting up an exemplary response
Martin Reinecke's avatar
Martin Reinecke committed
34
    mask = np.ones(signal_space.shape)
Theo Steininger's avatar
Theo Steininger committed
35
    N10 = int(N_pixels/10)
Martin Reinecke's avatar
Martin Reinecke committed
36
    mask[N10*5:N10*9, N10*5:N10*9] = 0.
37
    mask = ift.Field.from_global_data(signal_space, mask).lock()
38
39
    R = ift.GeometryRemover(signal_space)
    R = R*ift.DiagonalOperator(mask)
Martin Reinecke's avatar
Martin Reinecke committed
40
    R = R*ht
Martin Reinecke's avatar
Martin Reinecke committed
41
42
    R = R * ift.create_harmonic_smoothing_operator((harmonic_space,), 0,
                                                   response_sigma)
Theo Steininger's avatar
Theo Steininger committed
43
44
    data_domain = R.target[0]

Martin Reinecke's avatar
Martin Reinecke committed
45
    noiseless_data = R(mock_signal)
Martin Reinecke's avatar
Martin Reinecke committed
46
    noise_amplitude = noiseless_data.val.std()/signal_to_noise
Theo Steininger's avatar
Theo Steininger committed
47
    # Setting up the noise covariance and drawing a random noise realization
48
    N = ift.ScalingOperator(noise_amplitude**2, data_domain)
Martin Reinecke's avatar
Martin Reinecke committed
49
50
    noise = ift.Field.from_random(
        domain=data_domain, random_type='normal',
Martin Reinecke's avatar
Martin Reinecke committed
51
52
        std=noise_amplitude, mean=0)
    data = noiseless_data + noise
Theo Steininger's avatar
Theo Steininger committed
53
54

    # Wiener filter
55
    j = R.adjoint_times(N.inverse_times(data))
Martin Reinecke's avatar
Martin Reinecke committed
56
    ctrl = ift.GradientNormController(name="inverter", tol_abs_gradnorm=1e-2)
57
    inverter = ift.ConjugateGradient(controller=ctrl)
58
    wiener_curvature = ift.library.WienerFilterCurvature(
59
        S=S, N=N, R=R, inverter=inverter)
Martin Reinecke's avatar
Martin Reinecke committed
60
    m_k = wiener_curvature.inverse_times(j)
Martin Reinecke's avatar
Martin Reinecke committed
61
    m = ht(m_k)
Theo Steininger's avatar
Theo Steininger committed
62

63
    plotdict = {"colormap": "Planck-like"}
Martin Reinecke's avatar
Martin Reinecke committed
64
    ift.plot(ht(mock_signal), name="mock_signal.png", **plotdict)
Martin Reinecke's avatar
Martin Reinecke committed
65
    ift.plot(data.cast_domain(signal_space), name="data.png", **plotdict)
66
    ift.plot(m, name="map.png", **plotdict)
Theo Steininger's avatar
Theo Steininger committed
67

Martin Reinecke's avatar
Martin Reinecke committed
68
    # sampling the uncertainty map
69
    mean, variance = ift.probe_with_posterior_samples(wiener_curvature, ht, 5)
Martin Reinecke's avatar
Martin Reinecke committed
70
    ift.plot(ift.sqrt(variance), name="uncertainty.png", **plotdict)
71
    ift.plot(mean+m, name="posterior_mean.png", **plotdict)