getting_started_2.py 5.18 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 140
    lamb = lambda inp: R(sky(inp))
    mock_position = ift.from_random('normal', domain)
    data = lamb(mock_position)
141 142
    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
143

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

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

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