log_normal_wiener_filter.py 3.89 KB
Newer Older
Martin Reinecke's avatar
updates    
Martin Reinecke committed
1
2
import nifty2go as ift
import numpy as np
Theo Steininger's avatar
Theo Steininger committed
3

4
if __name__ == "__main__":
Martin Reinecke's avatar
updates    
Martin Reinecke committed
5
    np.random.seed(42)
6
7
8
    # Setting up parameters    |\label{code:wf_parameters}|
    correlation_length = 1.     # Typical distance over which the field is correlated
    field_variance = 2.         # Variance of field in position space
Theo Steininger's avatar
Theo Steininger committed
9
10
    response_sigma = 0.02       # Smoothing length of response (in same unit as L)
    signal_to_noise = 100       # The signal to noise ratio
11
12
13
14
15
16
17
    np.random.seed(43)          # Fixing the random seed
    def power_spectrum(k):      # Defining the power spectrum
        a = 4 * correlation_length * field_variance**2
        return a / (1 + k * correlation_length) ** 4

    # Setting up the geometry |\label{code:wf_geometry}|
    L = 2. # Total side-length of the domain
Theo Steininger's avatar
Theo Steininger committed
18
    N_pixels = 128 # Grid resolution (pixels per axis)
Martin Reinecke's avatar
Martin Reinecke committed
19
    #signal_space = ift.RGSpace([N_pixels, N_pixels], distances=L/N_pixels)
Martin Reinecke's avatar
updates    
Martin Reinecke committed
20
21
22
23
    signal_space = ift.HPSpace(16)
    harmonic_space = signal_space.get_default_codomain()
    fft = ift.FFTOperator(harmonic_space, target=signal_space)
    power_space = ift.PowerSpace(harmonic_space)
24
25

    # Creating the mock signal |\label{code:wf_mock_signal}|
Martin Reinecke's avatar
updates    
Martin Reinecke committed
26
    S = ift.create_power_operator(harmonic_space, power_spectrum=power_spectrum)
Martin Reinecke's avatar
Martin Reinecke committed
27
    mock_power = ift.PS_field(power_space, power_spectrum)
Martin Reinecke's avatar
adjust    
Martin Reinecke committed
28
    mock_signal = fft(ift.power_synthesize(mock_power, real_signal=True))
29
30

    # Setting up an exemplary response
31
    mask = ift.Field.ones(signal_space)
32
    N10 = int(N_pixels/10)
Theo Steininger's avatar
Theo Steininger committed
33
    #mask.val[N10*5:N10*9, N10*5:N10*9] = 0.
Martin Reinecke's avatar
updates    
Martin Reinecke committed
34
    R = ift.ResponseOperator(signal_space, sigma=(response_sigma,), exposure=(mask,)) #|\label{code:wf_response}|
35
    data_domain = R.target[0]
36
    R_harmonic = ift.ComposedOperator([fft, R])
37
38

    # Setting up the noise covariance and drawing a random noise realization
39
    ndiag = ift.Field.full(data_domain, mock_signal.var()/signal_to_noise)
40
    N = ift.DiagonalOperator(ndiag.weight(1))
Martin Reinecke's avatar
updates    
Martin Reinecke committed
41
    noise = ift.Field.from_random(domain=data_domain, random_type='normal',
42
                              std=mock_signal.std()/np.sqrt(signal_to_noise), mean=0)
Martin Reinecke's avatar
updates    
Martin Reinecke committed
43
    data = R(ift.exp(mock_signal)) + noise #|\label{code:wf_mock_data}|
44
45

    # Wiener filter
46
    m0 = ift.Field.zeros(harmonic_space)
47
48
    ctrl = ift.GradientNormController(verbose=False,tol_abs_gradnorm=1)
    ctrl2 = ift.GradientNormController(verbose=True,tol_abs_gradnorm=0.1, name="outer")
49
    inverter = ift.ConjugateGradient(controller=ctrl)
Martin Reinecke's avatar
updates    
Martin Reinecke committed
50
51
52
53
54
    energy = ift.library.LogNormalWienerFilterEnergy(m0, data, R_harmonic, N, S, inverter=inverter)
    minimizer1 = ift.VL_BFGS(controller=ctrl2,max_history_length=20)
    minimizer2 = ift.RelaxedNewton(controller=ctrl2)
    minimizer3 = ift.SteepestDescent(controller=ctrl2)

Martin Reinecke's avatar
Martin Reinecke committed
55
    #me1 = minimizer1(energy)
Martin Reinecke's avatar
updates    
Martin Reinecke committed
56
    me2 = minimizer2(energy)
Martin Reinecke's avatar
Martin Reinecke committed
57
    #me3 = minimizer3(energy)
Theo Steininger's avatar
Theo Steininger committed
58

Martin Reinecke's avatar
Martin Reinecke committed
59
    #m1 = fft(me1[0].position)
Martin Reinecke's avatar
updates    
Martin Reinecke committed
60
    m2 = fft(me2[0].position)
Martin Reinecke's avatar
Martin Reinecke committed
61
    #m3 = fft(me3[0].position)
62
63
64


    #Plotting #|\label{code:wf_plotting}|
Martin Reinecke's avatar
Martin Reinecke committed
65
66
67
68
69
    ift.plotting.plot(mock_signal.real,name='mock_signal.pdf', colormap="plasma",xlabel="Pixel Index",ylabel="Pixel Index")
    ift.plotting.plot(ift.Field(signal_space, val=np.log(data.val.real).reshape(signal_space.shape)),name="log_of_data.pdf", colormap="plasma",xlabel="Pixel Index",ylabel="Pixel Index")
    #ift.plotting.plot(m1.real,name='m_LBFGS.pdf', colormap="plasma",xlabel="Pixel Index",ylabel="Pixel Index")
    ift.plotting.plot(m2.real,name='m_Newton.pdf', colormap="plasma",xlabel="Pixel Index",ylabel="Pixel Index")
    #ift.plotting.plot(m3.real,name='m_SteepestDescent.pdf', colormap="plasma",xlabel="Pixel Index",ylabel="Pixel Index")
Theo Steininger's avatar
Theo Steininger committed
70

Martin Reinecke's avatar
Martin Reinecke committed
71
72
73
74
    # Probing the variance
    class Proby(ift.DiagonalProberMixin, ift.Prober): pass
    proby = Proby(signal_space, probe_count=1)
    proby(lambda z: fft(me2[0].curvature.inverse_times(fft.adjoint_times(z))))
75

Martin Reinecke's avatar
Martin Reinecke committed
76
77
78
    sm = ift.FFTSmoothingOperator(signal_space, sigma=0.02)
    variance = sm(proby.diagonal.weight(-1))
    ift.plotting.plot(variance, name = 'variance.pdf')