critical_filtering.py 6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10

from nifty import *

import plotly.offline as pl
import plotly.graph_objs as go

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

Jakob Knollmueller's avatar
Jakob Knollmueller committed
11
np.random.seed(232)
12

Jakob Knollmueller's avatar
Jakob Knollmueller committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

def plot_parameters(m,t,t_true, t_real, t_d):
    m = fft.adjoint_times(m)
    m_data = m.val.get_full_data().real
    t_data = t.val.get_full_data().real
    t_d_data = t_d.val.get_full_data().real
    t_true_data = t_true.val.get_full_data().real
    t_real_data = t_real.val.get_full_data().real
    pl.plot([go.Heatmap(z=m_data)], filename='map.html')
    pl.plot([go.Scatter(y=t_data), go.Scatter(y=t_true_data),
             go.Scatter(y=t_real_data), go.Scatter(y=t_d_data)], filename="t.html")

class AdjointFFTResponse(LinearOperator):
    def __init__(self, FFT, R, default_spaces=None):
        super(AdjointFFTResponse, self).__init__(default_spaces)
28
        self._domain = FFT.target
Jakob Knollmueller's avatar
Jakob Knollmueller committed
29
30
        self._target = R.target
        self.R = R
31
32
33
        self.FFT = FFT

    def _times(self, x, spaces=None):
Jakob Knollmueller's avatar
Jakob Knollmueller committed
34
        return self.R(self.FFT.adjoint_times(x))
35
36

    def _adjoint_times(self, x, spaces=None):
Jakob Knollmueller's avatar
Jakob Knollmueller committed
37
        return self.FFT(self.R.adjoint_times(x))
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
    @property
    def domain(self):
        return self._domain

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

    @property
    def unitary(self):
        return False

if __name__ == "__main__":

    distribution_strategy = 'not'

    # Set up position space
    s_space = RGSpace([128,128])
    # s_space = HPSpace(32)

    # Define harmonic transformation and associated harmonic space
    fft = FFTOperator(s_space)
    h_space = fft.target[0]

    # Setting up power space
Jakob Knollmueller's avatar
Jakob Knollmueller committed
63
64
    p_space = PowerSpace(h_space, logarithmic=False,
                         distribution_strategy=distribution_strategy, nbin=128)
65
66

    # Choosing the prior correlation structure and defining correlation operator
Jakob Knollmueller's avatar
Jakob Knollmueller committed
67
    pow_spec = (lambda k: (.05 / (k + 1) ** 2))
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
    S = create_power_operator(h_space, power_spectrum=pow_spec,
                              distribution_strategy=distribution_strategy)

    # Drawing a sample sh from the prior distribution in harmonic space
    sp = Field(p_space,  val=lambda z: pow_spec(z)**(1./2),
               distribution_strategy=distribution_strategy)
    sh = sp.power_synthesize(real_signal=True)


    # Choosing the measurement instrument
#    Instrument = SmoothingOperator(s_space, sigma=0.01)
    Instrument = DiagonalOperator(s_space, diagonal=1.)
#    Instrument._diagonal.val[200:400, 200:400] = 0


    #Adding a harmonic transformation to the instrument
Jakob Knollmueller's avatar
Jakob Knollmueller committed
84
85
86
    R = AdjointFFTResponse(fft, Instrument)

    noise = 1.
87
88
89
90
91
92
93
94
    N = DiagonalOperator(s_space, diagonal=noise, bare=True)
    n = Field.from_random(domain=s_space,
                          random_type='normal',
                          std=sqrt(noise),
                          mean=0)

    # Creating the mock data
    d = R(sh) + n
Jakob Knollmueller's avatar
Jakob Knollmueller committed
95
96
97
98
99

    realized_power = log(sh.power_analyze(logarithmic=p_space.config["logarithmic"],
                                          nbin=p_space.config["nbin"])**2)
    data_power  = log(fft(d).power_analyze(logarithmic=p_space.config["logarithmic"],
                                          nbin=p_space.config["nbin"])**2)
100
101
102
103
    d_data = d.val.get_full_data().real
    if rank == 0:
        pl.plot([go.Heatmap(z=d_data)], filename='data.html')

Jakob Knollmueller's avatar
Jakob Knollmueller committed
104
    #  minimization strategy
105
106
107
108
109
110
111
112

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


    minimizer1 = RelaxedNewton(convergence_tolerance=0,
                              convergence_level=1,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
113
                              iteration_limit=2,
114
                              callback=convergence_measure)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
115
116
117
118
    minimizer2 = VL_BFGS(convergence_tolerance=0,
                       iteration_limit=50,
                       callback=convergence_measure,
                       max_history_length=3)
119
120
121
122
123
124

    # Setting starting position
    flat_power = Field(p_space,val=10e-8)
    m0 = flat_power.power_synthesize(real_signal=True)

    t0 = Field(p_space, val=log(1./(1+p_space.kindex)**2))
Jakob Knollmueller's avatar
Jakob Knollmueller committed
125
    #t0 = data_power
126
127
128
129
130
131
132
133
134
    S0 = create_power_operator(h_space, power_spectrum=exp(t0),
                               distribution_strategy=distribution_strategy)


    for i in range(100):
        S0 = create_power_operator(h_space, power_spectrum=exp(t0),
                              distribution_strategy=distribution_strategy)

        # Initializing the  nonlinear Wiener Filter energy
Jakob Knollmueller's avatar
Jakob Knollmueller committed
135
        map_energy = WienerFilterEnergy(position=m0, d=d, R=R, N=N, S=S0)
136
        # Minimization with chosen minimizer
Jakob Knollmueller's avatar
Jakob Knollmueller committed
137
138
        map_energy = map_energy.analytic_solution()

139
140
141
142
143
        # Updating parameters for correlation structure reconstruction
        m0 = map_energy.position
        D0 = map_energy.curvature
        # Initializing the power energy with updated parameters
        power_energy = CriticalPowerEnergy(position=t0, m=m0, D=D0, sigma=10., samples=3)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
144
145
146
147
        if i > 0:
            (power_energy, convergence) = minimizer1(power_energy)
        else:
            (power_energy, convergence) = minimizer2(power_energy)
148
149
        # Setting new power spectrum
        t0 = power_energy.position
Jakob Knollmueller's avatar
Jakob Knollmueller committed
150
151
152
        t0.val[-1] = t0.val[-2]
        # Plotting current estimate
        plot_parameters(m0,t0,log(sp**2),realized_power, data_power)
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181

    # Transforming fields to position space for plotting

    ss = fft.adjoint_times(sh)
    m = fft.adjoint_times(map_energy.position)


    # Plotting

    d_data = d.val.get_full_data().real
    if rank == 0:
        pl.plot([go.Heatmap(z=d_data)], filename='data.html')

    tt_data = power_energy.position.val.get_full_data().real
    t_data = log(sp**2).val.get_full_data().real
    if rank == 0:
        pl.plot([go.Scatter(y=t_data),go.Scatter(y=tt_data)], filename="t.html")
    ss_data = ss.val.get_full_data().real
    if rank == 0:
        pl.plot([go.Heatmap(z=ss_data)], filename='ss.html')

    sh_data = sh.val.get_full_data().real
    if rank == 0:
        pl.plot([go.Heatmap(z=sh_data)], filename='sh.html')


    m_data = m.val.get_full_data().real
    if rank == 0:
        pl.plot([go.Heatmap(z=m_data)], filename='map.html')