wiener_filter_via_hamiltonian.py 3.89 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

Martin Reinecke's avatar
Martin Reinecke committed
13

14 15 16 17 18 19 20 21 22 23 24 25 26
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))
Martin Reinecke's avatar
Martin Reinecke committed
27

28 29 30 31 32 33 34 35 36 37 38 39
    @property
    def domain(self):
        return self._domain

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

    @property
    def unitary(self):
        return False

40

41 42
if __name__ == "__main__":

43
    distribution_strategy = 'not'
44

Jakob Knollmueller's avatar
Jakob Knollmueller committed
45
    # Set up position space
Martin Reinecke's avatar
Martin Reinecke committed
46
    s_space = RGSpace([128, 128])
Jakob Knollmueller's avatar
Jakob Knollmueller committed
47 48 49
    # s_space = HPSpace(32)

    # Define harmonic transformation and associated harmonic space
50 51
    fft = FFTOperator(s_space)
    h_space = fft.target[0]
Jakob Knollmueller's avatar
Jakob Knollmueller committed
52 53

    # Setting up power space
54
    p_space = PowerSpace(h_space, distribution_strategy=distribution_strategy)
55

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

60
    S = create_power_operator(h_space, power_spectrum=p_spec,
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(p_space, val=p_spec,
65 66
               distribution_strategy=distribution_strategy)
    sh = sp.power_synthesize(real_signal=True)
67 68
    ss = fft.adjoint_times(sh)

Jakob Knollmueller's avatar
Jakob Knollmueller committed
69
    # Choosing the measurement instrument
Jakob Knollmueller's avatar
Jakob Knollmueller committed
70
    # Instrument = SmoothingOperator(s_space, sigma=0.05)
71 72
    Instrument = DiagonalOperator(s_space, diagonal=1.)
#    Instrument._diagonal.val[200:400, 200:400] = 0
Jakob Knollmueller's avatar
Jakob Knollmueller committed
73

Martin Reinecke's avatar
Martin Reinecke committed
74
    # Adding a harmonic transformation to the instrument
75
    R = AdjointFFTResponse(fft, Instrument)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
76
    signal_to_noise = 1.
77 78
    N = DiagonalOperator(s_space, diagonal=ss.var()/signal_to_noise, bare=True)
    n = Field.from_random(domain=s_space,
79 80
                          random_type='normal',
                          std=ss.std()/np.sqrt(signal_to_noise),
81
                          mean=0)
82

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

Jakob Knollmueller's avatar
Jakob Knollmueller committed
87 88
    # Choosing the minimization strategy

Martin Reinecke's avatar
Martin Reinecke committed
89
    def convergence_measure(energy, iteration):  # returns current energy
90
        x = energy.value
Martin Reinecke's avatar
Martin Reinecke committed
91
        print(x, iteration)
92

93 94
#    minimizer = SteepestDescent(convergence_tolerance=0,
#                                iteration_limit=50,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
95
#                                callback=convergence_measure)
96 97

    minimizer = RelaxedNewton(convergence_tolerance=0,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
98
                              iteration_limit=1,
Jakob Knollmueller's avatar
Jakob Knollmueller committed
99
                              callback=convergence_measure)
100 101 102 103 104 105
    #
    # minimizer = VL_BFGS(convergence_tolerance=0,
    #                    iteration_limit=50,
    #                    callback=convergence_measure,
    #                    max_history_length=3)
    #
Jakob Knollmueller's avatar
Jakob Knollmueller committed
106

107
    inverter = ConjugateGradient(convergence_level=3,
Martin Reinecke's avatar
Martin Reinecke committed
108
                                 convergence_tolerance=1e-5,
109
                                 preconditioner=None)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
110
    # Setting starting position
111
    m0 = Field(h_space, val=.0)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
112 113

    # Initializing the Wiener Filter energy
Martin Reinecke's avatar
Martin Reinecke committed
114
    energy = WienerFilterEnergy(position=m0, d=d, R=R, N=N, S=S)
115
    D0 = energy.curvature
116

Jakob Knollmueller's avatar
Jakob Knollmueller committed
117
    # Solving the problem analytically
118
    m0 = D0.inverse_times(j)
119

120 121
    sample_variance = Field(sh.domain, val=0.)
    sample_mean = Field(sh.domain, val=0.)
122 123

    # sampling the uncertainty map
124
    n_samples = 10
125
    for i in range(n_samples):
126
        sample = fft(sugar.generate_posterior_sample(0., D0))
127 128
        sample_variance += sample**2
        sample_mean += sample
129
    variance = (sample_variance - sample_mean**2)/n_samples