wiener_filter_advanced.py 3.9 KB
Newer Older
1
2

from nifty import *
3

4
5
import plotly.offline as pl
import plotly.graph_objs as go
6
7
8
9
10

from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.rank

11
np.random.seed(42)
12

13
14
class AdjointFFTResponse(LinearOperator):
    def __init__(self, FFT, R, default_spaces=None):
Jakob Knollmueller's avatar
test    
Jakob Knollmueller committed
15
        super(AdjointFFTResponse, self).__init__(default_spaces)
16
        self._domain = FFT.target
Jakob Knollmueller's avatar
test    
Jakob Knollmueller committed
17
        self._target = R.target
18
19
20
        self.R = R
        self.FFT = FFT

Jakob Knollmueller's avatar
test    
Jakob Knollmueller committed
21
    def _times(self, x, spaces=None):
22
23
        return self.R(self.FFT.adjoint_times(x))

Jakob Knollmueller's avatar
test    
Jakob Knollmueller committed
24
    def _adjoint_times(self, x, spaces=None):
25
        return self.FFT(self.R.adjoint_times(x))
Jakob Knollmueller's avatar
test    
Jakob Knollmueller committed
26
27
28
29
30
31
32
33
34
35
36
    @property
    def domain(self):
        return self._domain

    @property
    def target(self):
        return self._target

    @property
    def unitary(self):
        return False
37

38
39


40
41
if __name__ == "__main__":

Martin Reinecke's avatar
Martin Reinecke committed
42
    distribution_strategy = 'not'
43

Jakob Knollmueller's avatar
Jakob Knollmueller committed
44
    # Set up position space
45
    s_space = RGSpace([128,128])
Jakob Knollmueller's avatar
Jakob Knollmueller committed
46
47
48
    # s_space = HPSpace(32)

    # Define harmonic transformation and associated harmonic space
49
50
    fft = FFTOperator(s_space)
    h_space = fft.target[0]
Jakob Knollmueller's avatar
Jakob Knollmueller committed
51
52

    # Setting up power space
53
54
    p_space = PowerSpace(h_space, distribution_strategy=distribution_strategy)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
55
    # Choosing the prior correlation structure and defining correlation operator
Jakob Knollmueller's avatar
Jakob Knollmueller committed
56
57
58
    p_spec = (lambda k: (42 / (k + 1) ** 3))

    S = create_power_operator(h_space, power_spectrum=p_spec,
59
60
                              distribution_strategy=distribution_strategy)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
61
    # Drawing a sample sh from the prior distribution in harmonic space
Jakob Knollmueller's avatar
Jakob Knollmueller committed
62
    sp = Field(p_space, val=p_spec,
63
64
               distribution_strategy=distribution_strategy)
    sh = sp.power_synthesize(real_signal=True)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
65
    ss = fft.adjoint_times(sh)
66

Jakob Knollmueller's avatar
Jakob Knollmueller committed
67
    # Choosing the measurement instrument
Jakob Knollmueller's avatar
Jakob Knollmueller committed
68
69
    # Instrument = SmoothingOperator(s_space, sigma=0.05)
    Instrument = DiagonalOperator(s_space, diagonal=1.)
70
#    Instrument._diagonal.val[200:400, 200:400] = 0
Jakob Knollmueller's avatar
Jakob Knollmueller committed
71
72

    #Adding a harmonic transformation to the instrument
73
    R = AdjointFFTResponse(fft, Instrument)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
74
    signal_to_noise = 1.
75
76
77
78
79
80
    N = DiagonalOperator(s_space, diagonal=ss.var()/signal_to_noise, bare=True)
    n = Field.from_random(domain=s_space,
                          random_type='normal',
                          std=ss.std()/np.sqrt(signal_to_noise),
                          mean=0)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
81
    # Creating the mock data
82
    d = R(sh) + n
83
    j = R.adjoint_times(N.inverse_times(d))
84

Jakob Knollmueller's avatar
Jakob Knollmueller committed
85
86
87
    # Choosing the minimization strategy

    def convergence_measure(energy, iteration): # returns current energy
88
89
        x = energy.value
        print (x, iteration)
90

91
92
#    minimizer = SteepestDescent(convergence_tolerance=0,
#                                iteration_limit=50,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
93
#                                callback=convergence_measure)
94
95

    minimizer = RelaxedNewton(convergence_tolerance=0,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
96
                              iteration_limit=1,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
97
98
99
                              callback=convergence_measure)
    #
    # minimizer = VL_BFGS(convergence_tolerance=0,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
100
    #                    iteration_limit=50,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
101
102
103
    #                    callback=convergence_measure,
    #                    max_history_length=3)
    #
Jakob Knollmueller's avatar
Jakob Knollmueller committed
104

105
106
107
    inverter = ConjugateGradient(convergence_level=3,
                                 convergence_tolerance=10e-5,
                                 preconditioner=None)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
108
    # Setting starting position
Jakob Knollmueller's avatar
Jakob Knollmueller committed
109
    m0 = Field(h_space, val=.0)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
110
111

    # Initializing the Wiener Filter energy
112
113
    energy = WienerFilterEnergy(position=m0, d=d, R=R, N=N, S=S, inverter=inverter)
    D0 = energy.curvature
114

Jakob Knollmueller's avatar
Jakob Knollmueller committed
115
    # Solving the problem analytically
116
117
    m0 = D0.inverse_times(j)

118
119
    sample_variance = Field(sh.power_analyze(logarithmic=False).domain,val=0. + 0j)
    sample_mean = Field(sh.domain,val=0. + 0j)
120
121
122

    # sampling the uncertainty map
    n_samples = 1
123
    for i in range(n_samples):
124
125
        sample = sugar.generate_posterior_sample(m0,D0)
        sample_variance += sample**2
126
        sample_mean += sample
127
    variance = sample_variance/n_samples - (sample_mean/n_samples)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
128