log_normal_wiener_filter.py 3.96 KB
Newer Older
1 2
# -*- coding: utf-8 -*-

Martin Reinecke's avatar
updates  
Martin Reinecke committed
3 4
import nifty2go as ift
import numpy as np
Theo Steininger's avatar
Theo Steininger committed
5

6
if __name__ == "__main__":
Martin Reinecke's avatar
updates  
Martin Reinecke committed
7
    np.random.seed(42)
8 9 10
    # 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
11 12
    response_sigma = 0.02       # Smoothing length of response (in same unit as L)
    signal_to_noise = 100       # The signal to noise ratio
13 14 15 16 17 18 19
    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
20
    N_pixels = 128 # Grid resolution (pixels per axis)
Martin Reinecke's avatar
Martin Reinecke committed
21
    #signal_space = ift.RGSpace([N_pixels, N_pixels], distances=L/N_pixels)
Martin Reinecke's avatar
updates  
Martin Reinecke committed
22 23 24 25
    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)
26 27

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

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

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

    # Wiener filter
Martin Reinecke's avatar
updates  
Martin Reinecke committed
48
    m0 = ift.Field(harmonic_space, val=0.)
49 50
    ctrl = ift.GradientNormController(verbose=False,tol_abs_gradnorm=1)
    ctrl2 = ift.GradientNormController(verbose=True,tol_abs_gradnorm=0.1, name="outer")
51
    inverter = ift.ConjugateGradient(controller=ctrl)
Martin Reinecke's avatar
updates  
Martin Reinecke committed
52 53 54 55 56
    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
57
    #me1 = minimizer1(energy)
Martin Reinecke's avatar
updates  
Martin Reinecke committed
58
    me2 = minimizer2(energy)
Martin Reinecke's avatar
Martin Reinecke committed
59
    #me3 = minimizer3(energy)
Theo Steininger's avatar
Theo Steininger committed
60

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


    #Plotting #|\label{code:wf_plotting}|
Martin Reinecke's avatar
Martin Reinecke committed
67 68 69 70 71
    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
72

Martin Reinecke's avatar
Martin Reinecke committed
73 74 75 76
    # 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))))
77

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