critical_filtering.py 4.91 KB
Newer Older
1
from nifty import *
2 3
from nifty.library.wiener_filter import WienerFilterEnergy
from nifty.library.critical_filter import CriticalPowerEnergy
4 5 6 7 8 9 10
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
def plot_parameters(m,t,p, p_d):
Jakob Knollmueller's avatar
Jakob Knollmueller committed
15 16

    x = log(t.domain[0].kindex)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
17
    m = fft.adjoint_times(m)
18 19 20 21 22 23
    m = m.val.get_full_data().real
    t = t.val.get_full_data().real
    p = p.val.get_full_data().real
    p_d = p_d.val.get_full_data().real
    pl.plot([go.Heatmap(z=m)], filename='map.html')
    pl.plot([go.Scatter(x=x,y=t), go.Scatter(x=x ,y=p), go.Scatter(x=x, y=p_d)], filename="t.html")
Jakob Knollmueller's avatar
Jakob Knollmueller committed
24

Jakob Knollmueller's avatar
Jakob Knollmueller committed
25 26 27 28

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

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

    def _adjoint_times(self, x, spaces=None):
Jakob Knollmueller's avatar
Jakob Knollmueller committed
38
        return self.FFT(self.R.adjoint_times(x))
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63
    @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
64
    p_space = PowerSpace(h_space, logarithmic=True,
65
                         distribution_strategy=distribution_strategy)
66 67

    # Choosing the prior correlation structure and defining correlation operator
68
    p_spec = (lambda k: (.5 / (k + 1) ** 3))
Jakob Knollmueller's avatar
Jakob Knollmueller committed
69
    S = create_power_operator(h_space, power_spectrum=p_spec,
70 71 72
                              distribution_strategy=distribution_strategy)

    # Drawing a sample sh from the prior distribution in harmonic space
Jakob Knollmueller's avatar
Jakob Knollmueller committed
73
    sp = Field(p_space,  val=p_spec,
74 75 76 77 78
               distribution_strategy=distribution_strategy)
    sh = sp.power_synthesize(real_signal=True)


    # Choosing the measurement instrument
Martin Reinecke's avatar
Martin Reinecke committed
79 80
    Instrument = SmoothingOperator(s_space, sigma=0.01)
    # Instrument = DiagonalOperator(s_space, diagonal=1.)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
81
    # Instrument._diagonal.val[200:400, 200:400] = 0
Jakob Knollmueller's avatar
Jakob Knollmueller committed
82
    #Instrument._diagonal.val[64:512-64, 64:512-64] = 0
83 84 85


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

88
    noise = 1.
89 90 91 92 93 94 95 96
    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
97

98 99
    # The information source
    j = R.adjoint_times(N.inverse_times(d))
Martin Reinecke's avatar
Martin Reinecke committed
100 101
    realized_power = log(sh.power_analyze(binbounds=p_space.binbounds))
    data_power = log(fft(d).power_analyze(binbounds=p_space.binbounds))
102
    d_data = d.val.get_full_data().real
Martin Reinecke's avatar
Martin Reinecke committed
103 104
    #if rank == 0:
    #    pl.plot([go.Heatmap(z=d_data)], filename='data.html')
105

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

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


Martin Reinecke's avatar
Martin Reinecke committed
113
    minimizer1 = RelaxedNewton(convergence_tolerance=1e-2,
114
                              convergence_level=2,
Martin Reinecke's avatar
Martin Reinecke committed
115
                              iteration_limit=5,
116
                              callback=convergence_measure)
117

Martin Reinecke's avatar
Martin Reinecke committed
118
    minimizer2 = VL_BFGS(convergence_tolerance=1e-3,
Martin Reinecke's avatar
Martin Reinecke committed
119
                       iteration_limit=20,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
120
                       callback=convergence_measure,
Martin Reinecke's avatar
Martin Reinecke committed
121 122 123 124
                       max_history_length=10)
    minimizer3 = SteepestDescent(convergence_tolerance=1e-3,
                       iteration_limit=70,
                       callback=convergence_measure)
125 126

    # Setting starting position
Martin Reinecke's avatar
Martin Reinecke committed
127
    flat_power = Field(p_space,val=1e-8)
128 129
    m0 = flat_power.power_synthesize(real_signal=True)

130 131
    t0 = Field(p_space, val=log(1./(1+p_space.kindex)**2))

Jakob Knollmueller's avatar
Jakob Knollmueller committed
132

133

134
    for i in range(500):
Jakob Knollmueller's avatar
Jakob Knollmueller committed
135
        S0 = create_power_operator(h_space, power_spectrum=exp(t0),
136 137 138
                              distribution_strategy=distribution_strategy)

        # Initializing the  nonlinear Wiener Filter energy
Martin Reinecke's avatar
Martin Reinecke committed
139
        map_energy = WienerFilterEnergy(position=m0, d=d, R=R, N=N, S=S0)
140
        # Solving the Wiener Filter analytically
141
        D0 = map_energy.curvature
142
        m0 = D0.inverse_times(j)
143
        # Initializing the power energy with updated parameters
144
        power_energy = CriticalPowerEnergy(position=t0, m=m0, D=D0, smoothness_prior=10., samples=3)
145

Martin Reinecke's avatar
Martin Reinecke committed
146
        (power_energy, convergence) = minimizer3(power_energy)
147

148

149
        # Setting new power spectrum
150
        t0.val  = power_energy.position.val.real
Jakob Knollmueller's avatar
Jakob Knollmueller committed
151

Jakob Knollmueller's avatar
Jakob Knollmueller committed
152
        # Plotting current estimate
153
        print i
Martin Reinecke's avatar
Martin Reinecke committed
154 155
        #if i%50 == 0:
        #    plot_parameters(m0,t0,log(sp), data_power)
156

Jakob Knollmueller's avatar
Jakob Knollmueller committed
157