getting_started_2.py 5.36 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2018 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.

Natalia Porqueres's avatar
Natalia Porqueres committed
19
20
21
import nifty5 as ift
import numpy as np

Martin Reinecke's avatar
Martin Reinecke committed
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class GaussianEnergy2(ift.Operator):
    def __init__(self, mean=None, covariance=None):
        super(GaussianEnergy2, self).__init__()
        self._mean = mean
        self._icov = None if covariance is None else covariance.inverse

    def __call__(self, x):
        residual = x if self._mean is None else x-self._mean
        icovres = residual if self._icov is None else self._icov(residual)
        res = .5 * (residual*icovres).sum()
        metric = ift.SandwichOperator.make(x.jac, self._icov)
        return res.add_metric(metric)

class PoissonianEnergy2(ift.Operator):
    def __init__(self, op, d):
        super(PoissonianEnergy2, self).__init__()
        self._op = op
        self._d = d

    def __call__(self, x):
        x = self._op(x)
        res = (x - self._d*x.log()).sum()
        metric = ift.SandwichOperator.make(x.jac, ift.makeOp(1./x.val))
        return res.add_metric(metric)

class MyHamiltonian(ift.Operator):
    def __init__(self, lh):
        super(MyHamiltonian, self).__init__()
        self._lh = lh
        self._prior = GaussianEnergy2()

    def __call__(self, x):
        return self._lh(x) + self._prior(x)

class EnergyAdapter(ift.Energy):
    def __init__(self, position, op):
        super(EnergyAdapter, self).__init__(position)
        self._op = op
        pvar = ift.Linearization.make_var(position)
        self._res = op(pvar)

    def at(self, position):
        return EnergyAdapter(position, self._op)

    @property
    def value(self):
        return self._res.val.local_data[()]

    @property
    def gradient(self):
        return self._res.gradient

    @property
    def metric(self):
        return self._res.metric

Natalia Porqueres's avatar
Natalia Porqueres committed
78

Jakob Knollmueller's avatar
Jakob Knollmueller committed
79
80
def get_2D_exposure():
    x_shape, y_shape = position_space.shape
Natalia Porqueres's avatar
Natalia Porqueres committed
81

Jakob Knollmueller's avatar
Jakob Knollmueller committed
82
    exposure = np.ones(position_space.shape)
Martin Reinecke's avatar
Martin Reinecke committed
83
84
85
86
87
88
    exposure[x_shape//3:x_shape//2, :] *= 2.
    exposure[x_shape*4//5:x_shape, :] *= .1
    exposure[x_shape//2:x_shape*3//2, :] *= 3.
    exposure[:, x_shape//3:x_shape//2] *= 2.
    exposure[:, x_shape*4//5:x_shape] *= .1
    exposure[:, x_shape//2:x_shape*3//2] *= 3.
89
90

    exposure = ift.Field.from_global_data(position_space, exposure)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
91
    return exposure
Natalia Porqueres's avatar
Natalia Porqueres committed
92

Natalia Porqueres's avatar
Natalia Porqueres committed
93

Jakob Knollmueller's avatar
Jakob Knollmueller committed
94
if __name__ == '__main__':
Philipp Arras's avatar
Philipp Arras committed
95
    # FIXME description of the tutorial
96
    np.random.seed(41)
Natalia Porqueres's avatar
Natalia Porqueres committed
97

Jakob Knollmueller's avatar
Jakob Knollmueller committed
98
    # Set up the position space of the signal
99
    #
Jakob Knollmueller's avatar
Jakob Knollmueller committed
100
101
102
    # # One dimensional regular grid with uniform exposure
    # position_space = ift.RGSpace(1024)
    # exposure = np.ones(position_space.shape)
Natalia Porqueres's avatar
Natalia Porqueres committed
103

104
    # Two-dimensional regular grid with inhomogeneous exposure
Jakob Knollmueller's avatar
Jakob Knollmueller committed
105
106
    position_space = ift.RGSpace([512, 512])
    exposure = get_2D_exposure()
Natalia Porqueres's avatar
Natalia Porqueres committed
107

Jakob Knollmueller's avatar
Jakob Knollmueller committed
108
109
    # # Sphere with with uniform exposure
    # position_space = ift.HPSpace(128)
110
    # exposure = ift.Field.full(position_space, 1.)
Natalia Porqueres's avatar
Natalia Porqueres committed
111

Jakob Knollmueller's avatar
Jakob Knollmueller committed
112
113
114
    # Defining harmonic space and transform
    harmonic_space = position_space.get_default_codomain()
    HT = ift.HarmonicTransformOperator(harmonic_space, position_space)
Natalia Porqueres's avatar
Natalia Porqueres committed
115

Martin Reinecke's avatar
Martin Reinecke committed
116
    domain = ift.DomainTuple.make(harmonic_space)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
117
    position = ift.from_random('normal', domain)
Natalia Porqueres's avatar
Natalia Porqueres committed
118

Jakob Knollmueller's avatar
Jakob Knollmueller committed
119
120
121
    # Define power spectrum and amplitudes
    def sqrtpspec(k):
        return 1. / (20. + k**2)
Natalia Porqueres's avatar
Natalia Porqueres committed
122

Jakob Knollmueller's avatar
Jakob Knollmueller committed
123
124
125
126
    p_space = ift.PowerSpace(harmonic_space)
    pd = ift.PowerDistributor(harmonic_space, p_space)
    a = ift.PS_field(p_space, sqrtpspec)
    A = pd(a)
Natalia Porqueres's avatar
Natalia Porqueres committed
127

Jakob Knollmueller's avatar
Jakob Knollmueller committed
128
    # Set up a sky model
Martin Reinecke's avatar
Martin Reinecke committed
129
    sky = lambda inp: (HT(inp*A)).exp()
Natalia Porqueres's avatar
Natalia Porqueres committed
130

Jakob Knollmueller's avatar
Jakob Knollmueller committed
131
132
133
134
    M = ift.DiagonalOperator(exposure)
    GR = ift.GeometryRemover(position_space)
    # Set up instrumental response
    R = GR * M
Natalia Porqueres's avatar
Natalia Porqueres committed
135

Jakob Knollmueller's avatar
Jakob Knollmueller committed
136
137
    # Generate mock data
    d_space = R.target[0]
Martin Reinecke's avatar
Martin Reinecke committed
138
139
    lamb = lambda inp: R(sky(inp))
    mock_position = ift.from_random('normal', domain)
Martin Reinecke's avatar
Martin Reinecke committed
140
141
142
    #ift.extra.check_value_gradient_consistency2(lamb, mock_position)
    #testl = GaussianEnergy2(None, M)
    #ift.extra.check_value_gradient_metric_consistency2(testl, sky(mock_position))
Martin Reinecke's avatar
Martin Reinecke committed
143
    data = lamb(mock_position)
144
145
    data = np.random.poisson(data.to_global_data().astype(np.float64))
    data = ift.Field.from_global_data(d_space, data)
Natalia Porqueres's avatar
Natalia Porqueres committed
146

Jakob Knollmueller's avatar
Jakob Knollmueller committed
147
    # Compute likelihood and Hamiltonian
Martin Reinecke's avatar
Martin Reinecke committed
148
    likelihood = PoissonianEnergy2(lamb, data)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
149
    ic_cg = ift.GradientNormController(iteration_limit=50)
150
151
    ic_newton = ift.GradientNormController(name='Newton', iteration_limit=50,
                                           tol_abs_gradnorm=1e-3)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
152
    minimizer = ift.RelaxedNewton(ic_newton)
Natalia Porqueres's avatar
Natalia Porqueres committed
153

Jakob Knollmueller's avatar
Jakob Knollmueller committed
154
    # Minimize the Hamiltonian
Martin Reinecke's avatar
Martin Reinecke committed
155
156
157
    H = MyHamiltonian(likelihood)
    H = EnergyAdapter(position, H)
    #ift.extra.check_value_gradient_consistency(H)
158
    H = H.make_invertible(ic_cg)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
159
    H, convergence = minimizer(H)
Natalia Porqueres's avatar
Natalia Porqueres committed
160

Jakob Knollmueller's avatar
Jakob Knollmueller committed
161
    # Plot results
Martin Reinecke's avatar
Martin Reinecke committed
162
    result_sky = sky(H.position)
Martin Reinecke's avatar
Martin Reinecke committed
163
164
    ift.plot(result_sky)
    ift.plot_finish()
Martin Reinecke's avatar
Martin Reinecke committed
165
    # FIXME PLOTTING