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