critical_filtering.py 5.92 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

11
np.random.seed(42)
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
63 64
    p_space = PowerSpace(h_space, logarithmic=True,
                         distribution_strategy=distribution_strategy, nbin=120)
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
    R = AdjointFFTResponse(fft, Instrument)

86
    noise = 5.
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

    realized_power = log(sh.power_analyze(logarithmic=p_space.config["logarithmic"],
                                          nbin=p_space.config["nbin"])**2)
98
    data_power = log(fft(d).power_analyze(logarithmic=p_space.config["logarithmic"],
Jakob Knollmueller's avatar
Jakob Knollmueller committed
99
                                          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,
113
                              iteration_limit=3,
114
                              callback=convergence_measure)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
115 116 117
    minimizer2 = VL_BFGS(convergence_tolerance=0,
                       iteration_limit=50,
                       callback=convergence_measure,
118
                       max_history_length=10)
119 120 121 122 123

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

124 125
    # t0 = Field(p_space, val=log(1./(1+p_space.kindex)**2))
    t0 = data_power- 1.
126 127 128 129
    S0 = create_power_operator(h_space, power_spectrum=exp(t0),
                               distribution_strategy=distribution_strategy)


130
    for i in range(500):
131 132 133 134
        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
        # Updating parameters for correlation structure reconstruction
        m0 = map_energy.position
        D0 = map_energy.curvature
        # Initializing the power energy with updated parameters
143 144 145 146
        power_energy = CriticalPowerEnergy(position=t0, m=m0, D=D0, sigma=1000., samples=1)

        (power_energy, convergence) = minimizer1(power_energy)

147
        # Setting new power spectrum
148 149
        t0.val  = power_energy.position.val.real
        # t0.val[-1] = t0.val[-2]
Jakob Knollmueller's avatar
Jakob Knollmueller committed
150 151
        # Plotting current estimate
        plot_parameters(m0,t0,log(sp**2),realized_power, data_power)
152 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

    # 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')