wiener_filter_via_hamiltonian.py 3.9 KB
Newer Older
1
from nifty import *
2

3
4
import plotly.offline as pl
import plotly.graph_objs as go
Martin Reinecke's avatar
Martin Reinecke committed
5
from nifty.library.wiener_filter import *
6
7
8
9
10

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

11
np.random.seed(42)
12

13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class AdjointFFTResponse(LinearOperator):
    def __init__(self, FFT, R, default_spaces=None):
        super(AdjointFFTResponse, self).__init__(default_spaces)
        self._domain = FFT.target
        self._target = R.target
        self.R = R
        self.FFT = FFT

    def _times(self, x, spaces=None):
        return self.R(self.FFT.adjoint_times(x))

    def _adjoint_times(self, x, spaces=None):
        return self.FFT(self.R.adjoint_times(x))
    @property
    def domain(self):
        return self._domain

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

    @property
    def unitary(self):
        return False

38
39


40
41
if __name__ == "__main__":

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
    p_space = PowerSpace(h_space, distribution_strategy=distribution_strategy)
54

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

58
    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
62
    sp = Field(p_space, val=p_spec,
63
64
               distribution_strategy=distribution_strategy)
    sh = sp.power_synthesize(real_signal=True)
65
66
    ss = fft.adjoint_times(sh)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
67
    # Choosing the measurement instrument
Jakob Knollmueller's avatar
Jakob Knollmueller committed
68
    # Instrument = SmoothingOperator(s_space, sigma=0.05)
69
70
    Instrument = DiagonalOperator(s_space, diagonal=1.)
#    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
    N = DiagonalOperator(s_space, diagonal=ss.var()/signal_to_noise, bare=True)
    n = Field.from_random(domain=s_space,
77
78
                          random_type='normal',
                          std=ss.std()/np.sqrt(signal_to_noise),
79
                          mean=0)
80

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
                              callback=convergence_measure)
98
99
100
101
102
103
    #
    # minimizer = VL_BFGS(convergence_tolerance=0,
    #                    iteration_limit=50,
    #                    callback=convergence_measure,
    #                    max_history_length=3)
    #
Jakob Knollmueller's avatar
Jakob Knollmueller committed
104

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

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

Jakob Knollmueller's avatar
Jakob Knollmueller committed
115
    # Solving the problem analytically
116
    m0 = D0.inverse_times(j)
117
118
119
120
121
122
123
124
125
126
127

    sample_variance = Field(sh.domain,val=0. + 0j)
    sample_mean = Field(sh.domain,val=0. + 0j)

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