energy_operators.py 12.3 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
# 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/>.
#
14
# Copyright(C) 2013-2019 Max-Planck-Society
Martin Reinecke's avatar
Martin Reinecke committed
15
#
16
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
Martin Reinecke's avatar
Martin Reinecke committed
17

Philipp Arras's avatar
Philipp Arras committed
18
19
import numpy as np

Philipp Arras's avatar
Philipp Arras committed
20
from .. import utilities
Martin Reinecke's avatar
Martin Reinecke committed
21
from ..domain_tuple import DomainTuple
Philipp Arras's avatar
Philipp Arras committed
22
23
from ..field import Field
from ..linearization import Linearization
Philipp Arras's avatar
Philipp Arras committed
24
25
from ..sugar import makeDomain, makeOp
from .linear_operator import LinearOperator
Martin Reinecke's avatar
Martin Reinecke committed
26
from .operator import Operator
Martin Reinecke's avatar
fix    
Martin Reinecke committed
27
from .sampling_enabler import SamplingEnabler
Philipp Arras's avatar
Philipp Arras committed
28
from .sandwich_operator import SandwichOperator
Martin Reinecke's avatar
Martin Reinecke committed
29
from .simple_linear_operators import VdotOperator
Martin Reinecke's avatar
Martin Reinecke committed
30
31
32


class EnergyOperator(Operator):
Philipp Arras's avatar
Philipp Arras committed
33
    """Operator which has a scalar domain as target domain.
34

Martin Reinecke's avatar
Martin Reinecke committed
35
    It is intended as an objective function for field inference.
36

Philipp Arras's avatar
Philipp Arras committed
37
38
39
40
41
    Examples
    --------
     - Information Hamiltonian, i.e. negative-log-probabilities.
     - Gibbs free energy, i.e. an averaged Hamiltonian, aka Kullbach-Leibler
       divergence.
42
    """
Martin Reinecke's avatar
Martin Reinecke committed
43
44
45
46
    _target = DomainTuple.scalar_domain()


class SquaredNormOperator(EnergyOperator):
Philipp Arras's avatar
Philipp Arras committed
47
    """Computes the L2-norm of the output of an operator.
48

Philipp Arras's avatar
Philipp Arras committed
49
50
51
52
    Parameters
    ----------
    domain : Domain, DomainTuple or tuple of Domain
        Target domain of the operator in which the L2-norm shall be computed.
Martin Reinecke's avatar
Martin Reinecke committed
53
    """
Philipp Arras's avatar
Philipp Arras committed
54

Martin Reinecke's avatar
Martin Reinecke committed
55
56
57
58
    def __init__(self, domain):
        self._domain = domain

    def apply(self, x):
59
        self._check_input(x)
Martin Reinecke's avatar
Martin Reinecke committed
60
        if isinstance(x, Linearization):
Martin Reinecke's avatar
Martin Reinecke committed
61
            val = Field.scalar(x.val.vdot(x.val))
Martin Reinecke's avatar
Martin Reinecke committed
62
            jac = VdotOperator(2*x.val)(x.jac)
63
            return x.new(val, jac)
Martin Reinecke's avatar
Martin Reinecke committed
64
        return Field.scalar(x.vdot(x))
Martin Reinecke's avatar
Martin Reinecke committed
65
66
67


class QuadraticFormOperator(EnergyOperator):
Philipp Arras's avatar
Philipp Arras committed
68
69
70
71
72
    """Computes the L2-norm of a Field or MultiField with respect to a
    specific metric `endo`.

    .. math ::
        E(f) = \\frac12 f^\\dagger \\text{endo}(f)
73
74
75

    Parameters
    ----------
Philipp Arras's avatar
Philipp Arras committed
76
77
    endo : EndomorphicOperator
         Kernel of quadratic form.
Martin Reinecke's avatar
Martin Reinecke committed
78
    """
Philipp Arras's avatar
Philipp Arras committed
79
80

    def __init__(self, endo):
Martin Reinecke's avatar
Martin Reinecke committed
81
        from .endomorphic_operator import EndomorphicOperator
Philipp Arras's avatar
Philipp Arras committed
82
        if not isinstance(endo, EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
83
            raise TypeError("op must be an EndomorphicOperator")
Philipp Arras's avatar
Philipp Arras committed
84
85
        self._op = endo
        self._domain = endo.domain
Martin Reinecke's avatar
Martin Reinecke committed
86
87

    def apply(self, x):
88
        self._check_input(x)
Martin Reinecke's avatar
Martin Reinecke committed
89
        if isinstance(x, Linearization):
Martin Reinecke's avatar
Martin Reinecke committed
90
91
            t1 = self._op(x.val)
            jac = VdotOperator(t1)(x.jac)
Martin Reinecke's avatar
Martin Reinecke committed
92
            val = Field.scalar(0.5*x.val.vdot(t1))
93
            return x.new(val, jac)
Martin Reinecke's avatar
Martin Reinecke committed
94
        return Field.scalar(0.5*x.vdot(self._op(x)))
Martin Reinecke's avatar
Martin Reinecke committed
95
96
97


class GaussianEnergy(EnergyOperator):
Martin Reinecke's avatar
Martin Reinecke committed
98
    """Class for energies of fields with Gaussian probability distribution.
99

Philipp Arras's avatar
Philipp Arras committed
100
    Represents up to constants in :math:`m`:
Martin Reinecke's avatar
Martin Reinecke committed
101

Philipp Arras's avatar
Philipp Arras committed
102
103
    .. math ::
        E(f) = - \\log G(f-m, D) = 0.5 (f-m)^\\dagger D^{-1} (f-m),
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
104

Philipp Arras's avatar
Philipp Arras committed
105
106
    an information energy for a Gaussian distribution with mean m and
    covariance D.
107

Philipp Arras's avatar
Philipp Arras committed
108
109
110
111
112
113
114
115
116
117
118
119
120
    Parameters
    ----------
    mean : Field
        Mean of the Gaussian. Default is 0.
    covariance : LinearOperator
        Covariance of the Gaussian. Default is the identity operator.
    domain : Domain, DomainTuple of tuple of Domain
        Operator domain. By default it is inferred from `mean` or
        `covariance` if specified

    Note
    ----
    At least one of the arguments has to be provided.
Martin Reinecke's avatar
Martin Reinecke committed
121
    """
Martin Reinecke's avatar
Martin Reinecke committed
122

Martin Reinecke's avatar
Martin Reinecke committed
123
    def __init__(self, mean=None, covariance=None, domain=None):
Philipp Arras's avatar
Philipp Arras committed
124
125
126
127
128
129
130
131
        if mean is not None and not isinstance(mean, Field):
            raise TypeError
        if covariance is not None and not isinstance(covariance,
                                                     LinearOperator):
            raise TypeError
        if domain is not None:
            domain = DomainTuple.make(domain)

Martin Reinecke's avatar
Martin Reinecke committed
132
133
134
135
136
137
138
139
140
141
        self._domain = None
        if mean is not None:
            self._checkEquivalence(mean.domain)
        if covariance is not None:
            self._checkEquivalence(covariance.domain)
        if domain is not None:
            self._checkEquivalence(domain)
        if self._domain is None:
            raise ValueError("no domain given")
        self._mean = mean
Martin Reinecke's avatar
Martin Reinecke committed
142
143
144
145
        if covariance is None:
            self._op = SquaredNormOperator(self._domain).scale(0.5)
        else:
            self._op = QuadraticFormOperator(covariance.inverse)
Martin Reinecke's avatar
Martin Reinecke committed
146
147
148
        self._icov = None if covariance is None else covariance.inverse

    def _checkEquivalence(self, newdom):
Martin Reinecke's avatar
fix    
Martin Reinecke committed
149
        newdom = makeDomain(newdom)
Martin Reinecke's avatar
Martin Reinecke committed
150
        if self._domain is None:
Philipp Arras's avatar
Philipp Arras committed
151
            self._domain = newdom
Martin Reinecke's avatar
Martin Reinecke committed
152
        else:
Philipp Arras's avatar
Philipp Arras committed
153
            if self._domain != newdom:
Martin Reinecke's avatar
Martin Reinecke committed
154
155
156
                raise ValueError("domain mismatch")

    def apply(self, x):
157
        self._check_input(x)
Philipp Arras's avatar
Philipp Arras committed
158
        residual = x if self._mean is None else x - self._mean
Philipp Arras's avatar
Changes    
Philipp Arras committed
159
        res = self._op(residual).real
160
        if not isinstance(x, Linearization) or not x.want_metric:
Martin Reinecke's avatar
Martin Reinecke committed
161
162
163
164
165
166
            return res
        metric = SandwichOperator.make(x.jac, self._icov)
        return res.add_metric(metric)


class PoissonianEnergy(EnergyOperator):
Philipp Arras's avatar
Philipp Arras committed
167
168
    """Class for likelihood Hamiltonians of expected count field constrained
    by Poissonian count data.
169

Philipp Arras's avatar
Philipp Arras committed
170
    Represents up to an f-independent term :math:`log(d!)`:
171

Philipp Arras's avatar
Philipp Arras committed
172
173
    .. math ::
        E(f) = -\\log \\text{Poisson}(d|f) = \\sum f - d^\\dagger \\log(f),
174

Philipp Arras's avatar
Philipp Arras committed
175
    where f is a :class:`Field` in data space with the expectation values for
Martin Reinecke's avatar
Martin Reinecke committed
176
    the counts.
Philipp Arras's avatar
Philipp Arras committed
177
178
179
180
181
182

    Parameters
    ----------
    d : Field
        Data field with counts. Needs to have integer dtype and all field
        values need to be non-negative.
Martin Reinecke's avatar
Martin Reinecke committed
183
    """
Philipp Arras's avatar
Philipp Arras committed
184

185
    def __init__(self, d):
Philipp Arras's avatar
Philipp Arras committed
186
187
188
189
        if not isinstance(d, Field) or not np.issubdtype(d.dtype, np.integer):
            raise TypeError
        if np.any(d.local_data < 0):
            raise ValueError
190
191
        self._d = d
        self._domain = DomainTuple.make(d.domain)
Martin Reinecke's avatar
Martin Reinecke committed
192
193

    def apply(self, x):
194
        self._check_input(x)
Martin Reinecke's avatar
Martin Reinecke committed
195
196
        res = x.sum() - x.log().vdot(self._d)
        if not isinstance(x, Linearization):
Martin Reinecke's avatar
Martin Reinecke committed
197
            return Field.scalar(res)
198
199
        if not x.want_metric:
            return res
Martin Reinecke's avatar
Martin Reinecke committed
200
201
202
        metric = SandwichOperator.make(x.jac, makeOp(1./x.val))
        return res.add_metric(metric)

203

204
class InverseGammaLikelihood(EnergyOperator):
Martin Reinecke's avatar
Martin Reinecke committed
205
    """Special class for inverse Gamma distributed covariances.
206
207
208

    RL FIXME: To be documented.
    """
Philipp Arras's avatar
Philipp Arras committed
209

Martin Reinecke's avatar
Martin Reinecke committed
210
    def __init__(self, d):
Philipp Arras's avatar
Philipp Arras committed
211
212
        if not isinstance(d, Field):
            raise TypeError
213
214
        self._d = d
        self._domain = DomainTuple.make(d.domain)
215
216

    def apply(self, x):
217
        self._check_input(x)
Philipp Frank's avatar
Philipp Frank committed
218
        res = 0.5*(x.log().sum() + (1./x).vdot(self._d))
219
220
        if not isinstance(x, Linearization):
            return Field.scalar(res)
221
222
        if not x.want_metric:
            return res
223
224
225
226
        metric = SandwichOperator.make(x.jac, makeOp(0.5/(x.val**2)))
        return res.add_metric(metric)


Martin Reinecke's avatar
Martin Reinecke committed
227
class BernoulliEnergy(EnergyOperator):
Martin Reinecke's avatar
Martin Reinecke committed
228
    """Class for likelihood-energies of expected event frequency constrained by
229
230
231
232
    event data.

    Parameters
    ----------
Martin Reinecke's avatar
Martin Reinecke committed
233
    d : Field
234
235
        data field with events (=1) or non-events (=0)

Martin Reinecke's avatar
Martin Reinecke committed
236
    Notes
237
    -----
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
238
    ``E = BernoulliEnergy(d)`` represents
239

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
240
241
    :math:`E(f) = -\\log \\text{Bernoulli}(d|f) =
    -d^\\dagger \\log f  - (1-d)^\\dagger \\log(1-f)`,
242

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
243
244
    where f is a field in data space (d.domain) with the expected
    frequencies of events.
Martin Reinecke's avatar
Martin Reinecke committed
245
    """
Philipp Arras's avatar
Philipp Arras committed
246

247
    def __init__(self, d):
Martin Reinecke's avatar
Martin Reinecke committed
248
        self._d = d
249
        self._domain = DomainTuple.make(d.domain)
Martin Reinecke's avatar
Martin Reinecke committed
250
251

    def apply(self, x):
252
        self._check_input(x)
Philipp Arras's avatar
Philipp Arras committed
253
        v = x.log().vdot(-self._d) - (1. - x).log().vdot(1. - self._d)
Martin Reinecke's avatar
Martin Reinecke committed
254
        if not isinstance(x, Linearization):
Martin Reinecke's avatar
Martin Reinecke committed
255
            return Field.scalar(v)
256
257
        if not x.want_metric:
            return v
Philipp Arras's avatar
Philipp Arras committed
258
        met = makeOp(1./(x.val*(1. - x.val)))
Martin Reinecke's avatar
Martin Reinecke committed
259
260
261
262
263
        met = SandwichOperator.make(x.jac, met)
        return v.add_metric(met)


class Hamiltonian(EnergyOperator):
264
265
266
267
268
    """Class for information Hamiltonians.

    Parameters
    ----------
    lh : EnergyOperator
Martin Reinecke's avatar
Martin Reinecke committed
269
        a likelihood energy
Martin Reinecke's avatar
Martin Reinecke committed
270
    ic_samp : IterationController
Martin Reinecke's avatar
Martin Reinecke committed
271
272
273
        is passed to SamplingEnabler to draw Gaussian distributed samples
        with covariance = metric of Hamiltonian
        (= Hessian without terms that generate negative eigenvalues)
274

Martin Reinecke's avatar
Martin Reinecke committed
275
    Notes
276
    -----
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
277
    ``H = Hamiltonian(E_lh)`` represents
278

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
279
    :math:`H(f) = 0.5 f^\\dagger f + E_{lh}(f)`
280

Martin Reinecke's avatar
Martin Reinecke committed
281
282
    an information Hamiltonian for a field f with a white Gaussian prior
    (unit covariance) and the likelihood energy :math:`E_{lh}`.
283

Martin Reinecke's avatar
Martin Reinecke committed
284
    Other field priors can be represented via transformations of a white
285
286
    Gaussian field into a field with the desired prior probability structure.

Martin Reinecke's avatar
Martin Reinecke committed
287
    By implementing prior information this way, the field prior is represented
288
289
290
    by a generative model, from which NIFTy can draw samples and infer a field
    using the Maximum a Posteriori (MAP) or the Variational Bayes (VB) method.

Martin Reinecke's avatar
Martin Reinecke committed
291
    For more details see:
292
293
    "Encoding prior knowledge in the structure of the likelihood"
    Jakob Knollmüller, Torsten A. Ensslin, submitted, arXiv:1812.04403
Martin Reinecke's avatar
Martin Reinecke committed
294
    `<https://arxiv.org/abs/1812.04403>`_
Martin Reinecke's avatar
Martin Reinecke committed
295
    """
Philipp Arras's avatar
Philipp Arras committed
296

Martin Reinecke's avatar
Martin Reinecke committed
297
298
299
300
    def __init__(self, lh, ic_samp=None):
        self._lh = lh
        self._prior = GaussianEnergy(domain=lh.domain)
        self._ic_samp = ic_samp
Martin Reinecke's avatar
Martin Reinecke committed
301
        self._domain = lh.domain
Martin Reinecke's avatar
Martin Reinecke committed
302
303

    def apply(self, x):
304
        self._check_input(x)
Philipp Arras's avatar
Philipp Arras committed
305
306
307
        if (self._ic_samp is None or not isinstance(x, Linearization)
                or not x.want_metric):
            return self._lh(x) + self._prior(x)
Martin Reinecke's avatar
Martin Reinecke committed
308
        else:
309
            lhx, prx = self._lh(x), self._prior(x)
Martin Reinecke's avatar
Martin Reinecke committed
310
311
            mtr = SamplingEnabler(lhx.metric, prx.metric.inverse,
                                  self._ic_samp, prx.metric.inverse)
Philipp Arras's avatar
Philipp Arras committed
312
            return (lhx + prx).add_metric(mtr)
Martin Reinecke's avatar
Martin Reinecke committed
313

Philipp Arras's avatar
Philipp Arras committed
314
315
316
317
318
    def __repr__(self):
        subs = 'Likelihood:\n{}'.format(utilities.indent(self._lh.__repr__()))
        subs += '\nPrior: Quadratic{}'.format(self._lh.domain.keys())
        return 'Hamiltonian:\n' + utilities.indent(subs)

Martin Reinecke's avatar
Martin Reinecke committed
319

Martin Reinecke's avatar
Martin Reinecke committed
320
321
class AveragedEnergy(EnergyOperator):
    """Class for Kullbach-Leibler (KL) Divergence or Gibbs free energies
Martin Reinecke's avatar
Martin Reinecke committed
322
323

    Precisely a sample averaged Hamiltonian (or other energy) that represents
324
325
326
    approximatively the relevant part of a KL to be used in Variational Bayes
    inference if the samples are drawn from the approximating Gaussian.

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
327
328
    Let :math:`Q(f) = G(f-m,D)` Gaussian used to approximate
    :math:`P(f|d)`, the correct posterior with information Hamiltonian
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
329
    :math:`H(d,f) = -\\log P(d,f) = -\\log P(f|d) + \\text{const.}`
330
331
332

    The KL divergence between those should then be optimized for m. It is

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
333
334
335
    :math:`KL(Q,P) = \\int Df Q(f) \\log Q(f)/P(f)\\\\
    = \\left< \\log Q(f) \\right>_Q(f) - \\left< \\log P(f) \\right>_Q(f)\\\\
    = \\text{const} + \\left< H(f) \\right>_G(f-m,D)`
Martin Reinecke's avatar
Martin Reinecke committed
336

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
337
338
    in essence the information Hamiltonian averaged over a Gaussian
    distribution centered on the mean m.
339

Martin Reinecke's avatar
Martin Reinecke committed
340
    ``AveragedEnergy(H)`` approximates
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
341
    :math:`\\left< H(f) \\right>_{G(f-m,D)}` if the residuals
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
342
    :math:`f-m` are drawn from covariance :math:`D`.
Martin Reinecke's avatar
Martin Reinecke committed
343

344
345
346
347
    Parameters
    ----------
    h: Hamiltonian
       the Hamiltonian/energy to be averaged
Martin Reinecke's avatar
Martin Reinecke committed
348
349
350
    res_samples : iterable of Fields
       set of residual sample points to be added to mean field
       for approximate estimation of the KL
351

Martin Reinecke's avatar
Martin Reinecke committed
352
353
    Notes
    -----
Martin Reinecke's avatar
Martin Reinecke committed
354
    ``KL = AveragedEnergy(H, samples)`` represents
Martin Reinecke's avatar
Martin Reinecke committed
355

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
356
    :math:`\\text{KL}(m) = \\sum_i H(m+v_i) / N`,
357

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
358
359
    where :math:`v_i` are the residual samples, :math:`N` is their number,
    and :math:`m` is the mean field around which the samples are drawn.
360

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
361
362
363
364
    Having symmetrized residual samples, with both v_i and -v_i being present
    ensures that the distribution mean is exactly represented. This reduces
    sampling noise and helps the numerics of the KL minimization process in the
    variational Bayes inference.
Martin Reinecke's avatar
Martin Reinecke committed
365
    """
Philipp Arras's avatar
Philipp Arras committed
366

Martin Reinecke's avatar
Martin Reinecke committed
367
368
    def __init__(self, h, res_samples):
        self._h = h
Martin Reinecke's avatar
Martin Reinecke committed
369
        self._domain = h.domain
Martin Reinecke's avatar
Martin Reinecke committed
370
371
372
        self._res_samples = tuple(res_samples)

    def apply(self, x):
373
        self._check_input(x)
Philipp Arras's avatar
Philipp Arras committed
374
375
        mymap = map(lambda v: self._h(x + v), self._res_samples)
        return utilities.my_sum(mymap)*(1./len(self._res_samples))