wiener_filter_via_hamiltonian.py 5.12 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 = 'equal'
43

Jakob Knollmueller's avatar
Jakob Knollmueller committed
44
    # Set up position space
45
    signal_space = RGSpace([256,256])
Jakob Knollmueller's avatar
Jakob Knollmueller committed
46
    # s_space = HPSpace(32)
47
    harmonic_space = FFTOperator.get_default_codomain(signal_space)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
48

49
50
    fft = FFTOperator(harmonic_space, target=signal_space,
                      domain_dtype=np.complex, target_dtype=np.float)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
51
52
53
    # Define harmonic transformation and associated harmonic space

    # Setting up power space
54
55
    power_space = PowerSpace(harmonic_space,
                             distribution_strategy=distribution_strategy)
56

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

60
    S = create_power_operator(harmonic_space, power_spectrum=power_spectrum,
61
62
                              distribution_strategy=distribution_strategy)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
63
    # Drawing a sample sh from the prior distribution in harmonic space
64
    sp = Field(power_space, val=power_spectrum,
65
66
               distribution_strategy=distribution_strategy)
    sh = sp.power_synthesize(real_signal=True)
67
    ss = fft(sh)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
68
    # Choosing the measurement instrument
Jakob Knollmueller's avatar
Jakob Knollmueller committed
69
    # Instrument = SmoothingOperator(s_space, sigma=0.05)
70
71
72
73
    mask = Field(signal_space, val=1,
                 distribution_strategy=distribution_strategy)
    mask.val[63:127, 63:127] = 0.
    Instrument = DiagonalOperator(signal_space, diagonal=mask)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
74
75

    #Adding a harmonic transformation to the instrument
76
    R = ComposedOperator([fft, Instrument], default_spaces=[0, 0])
Jakob Knollmueller's avatar
Jakob Knollmueller committed
77
    signal_to_noise = 1.
78
79
80
81
    N = DiagonalOperator(signal_space, diagonal=ss.var()/signal_to_noise,
                         bare=True,
                         distribution_strategy=distribution_strategy)
    n = Field.from_random(domain=signal_space,
82
83
                          random_type='normal',
                          std=ss.std()/np.sqrt(signal_to_noise),
84
                          mean=0, distribution_strategy=distribution_strategy)
85

Jakob Knollmueller's avatar
Jakob Knollmueller committed
86
    # Creating the mock data
87
    d = R(sh) + n
88
    j = R.adjoint_times(N.inverse_times(d))
89

Jakob Knollmueller's avatar
Jakob Knollmueller committed
90
91
92
    # Choosing the minimization strategy

    def convergence_measure(energy, iteration): # returns current energy
93
94
        x = energy.value
        print (x, iteration)
95

96
97
#    minimizer = SteepestDescent(convergence_tolerance=0,
#                                iteration_limit=50,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
98
#                                callback=convergence_measure)
99
100

    minimizer = RelaxedNewton(convergence_tolerance=0,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
101
                              iteration_limit=1,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
102
                              callback=convergence_measure)
103
104
105
106
107
108

    minimizer = VL_BFGS(convergence_tolerance=0,
                       iteration_limit=500,
                       callback=convergence_measure,
                       max_history_length=3)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
109

110
    inverter = ConjugateGradient(convergence_level=3,
Martin Reinecke's avatar
Martin Reinecke committed
111
                                 convergence_tolerance=1e-5,
112
                                 preconditioner=None)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
113
    # Setting starting position
114
115
    m0 = Field(harmonic_space, val=.0,
               distribution_strategy=distribution_strategy)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
116
117

    # Initializing the Wiener Filter energy
118
119
    energy = WienerFilterEnergy(position=m0, d=d, R=R, N=N, S=S,
                                inverter=inverter)
120
    D0 = energy.curvature
121

Jakob Knollmueller's avatar
Jakob Knollmueller committed
122
    # Solving the problem analytically
123
    m0 = D0.inverse_times(j)
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
    # solution, convergence = minimizer(energy)
    # m0 = solution.position
    m0_s = Field(signal_space,val=fft(m0).val.get_full_data().real)

    plotter = plotting.RG2DPlotter()
    plotter.title = 'mock_signal.html';
    plotter(ss)
    plotter.title = 'data.html'
    plotter(Field(signal_space,
                  val=Instrument.adjoint_times(d).val.get_full_data()\
                  .reshape(signal_space.shape)))
    plotter.title = 'map.html'; plotter(m0_s)
    #
    # sample_variance = Field(sh.domain,val=0. + 0j,
    #                         distribution_strategy=distribution_strategy)
    # sample_mean = Field(sh.domain,val=0. + 0j,
    #                         distribution_strategy=distribution_strategy)
    #
    # # 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)**2
Jakob Knollmueller's avatar
Jakob Knollmueller committed
149