wiener_filter.py 3.27 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
import nifty2go 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
Theo Steininger's avatar
Theo Steininger committed
7
8
9
10
11
    correlation_length_scale = 1.  # Typical distance over which the field is correlated
    fluctuation_scale = 2.         # Variance of field in position space
    response_sigma = 0.05          # Smoothing length of response (in same unit as L)
    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
12

Theo Steininger's avatar
Theo Steininger committed
13
14
15
16
    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

Martin Reinecke's avatar
Martin Reinecke committed
17
    # Setting up the geometry
Theo Steininger's avatar
Theo Steininger committed
18
19
20
    L = 2.  # Total side-length of the domain
    N_pixels = 512  # Grid resolution (pixels per axis)
    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
    fft = ift.FFTOperator(harmonic_space, target=signal_space)
Martin Reinecke's avatar
Martin Reinecke committed
23
    power_space = ift.PowerSpace(harmonic_space, binbounds=ift.PowerSpace.useful_binbounds(harmonic_space,logarithmic=True))
Theo Steininger's avatar
Theo Steininger committed
24

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

    # Setting up an exemplary response
Martin Reinecke's avatar
Martin Reinecke committed
32
    mask = np.ones(signal_space.shape)
Theo Steininger's avatar
Theo Steininger committed
33
    N10 = int(N_pixels/10)
Martin Reinecke's avatar
Martin Reinecke committed
34
35
36
37
    mask[N10*5:N10*9, N10*5:N10*9] = 0.
    mask = ift.Field(signal_space, ift.dobj.from_global_data(mask))
    R = ift.ResponseOperator(signal_space, sigma=(response_sigma,),
                             exposure=(mask,))
Theo Steininger's avatar
Theo Steininger committed
38
    data_domain = R.target[0]
39
    R_harmonic = ift.ComposedOperator([fft, R])
Theo Steininger's avatar
Theo Steininger committed
40
41

    # Setting up the noise covariance and drawing a random noise realization
42
    ndiag = ift.Field.full(data_domain, mock_signal.var()/signal_to_noise)
43
    N = ift.DiagonalOperator(ndiag.weight(1))
Theo Steininger's avatar
Theo Steininger committed
44
45
    noise = ift.Field.from_random(domain=data_domain, random_type='normal',
                                  std=mock_signal.std()/np.sqrt(signal_to_noise), mean=0)
Martin Reinecke's avatar
Martin Reinecke committed
46
    data = R(mock_signal) + noise
Theo Steininger's avatar
Theo Steininger committed
47
48
49

    # Wiener filter
    j = R_harmonic.adjoint_times(N.inverse_times(data))
Martin Reinecke's avatar
Martin Reinecke committed
50
    ctrl = ift.GradientNormController(verbose=True, tol_abs_gradnorm=0.1)
51
    inverter = ift.ConjugateGradient(controller=ctrl)
52
53
    wiener_curvature = ift.library.WienerFilterCurvature(
        S=S, N=N, R=R_harmonic, inverter=inverter)
Martin Reinecke's avatar
Martin Reinecke committed
54
    m_k = wiener_curvature.inverse_times(j)
Theo Steininger's avatar
Theo Steininger committed
55
56
    m = fft(m_k)

Martin Reinecke's avatar
Martin Reinecke committed
57
58
59
60
61
    # Probing the uncertainty
    class Proby(ift.DiagonalProberMixin, ift.Prober):
        pass
    proby = Proby(signal_space, probe_count=1, ncpu=1)
    proby(lambda z: fft(wiener_curvature.inverse_times(fft.inverse_times(z))))
Theo Steininger's avatar
Theo Steininger committed
62

Martin Reinecke's avatar
Martin Reinecke committed
63
    sm = ift.FFTSmoothingOperator(signal_space, sigma=0.03)
Martin Reinecke's avatar
Martin Reinecke committed
64
    variance = ift.sqrt(sm(proby.diagonal.weight(-1)))
65

Martin Reinecke's avatar
Martin Reinecke committed
66
    # Plotting
Martin Reinecke's avatar
Martin Reinecke committed
67
68
69
70
    ift.plotting.plot(variance,name="uncertainty.png",xlabel='Pixel index', ylabel='Pixel index')
    ift.plotting.plot(mock_signal,name="mock_signal.png",xlabel='Pixel index', ylabel='Pixel index')
    ift.plotting.plot(ift.Field(signal_space, val=data.val),name="data.png",xlabel='Pixel index', ylabel='Pixel index')
    ift.plotting.plot(m,name="map.png",xlabel='Pixel index', ylabel='Pixel index')