energy_operators.py 11.2 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
    Examples
    --------
     - Information Hamiltonian, i.e. negative-log-probabilities.
Martin Reinecke's avatar
Martin Reinecke committed
40
     - Gibbs free energy, i.e. an averaged Hamiltonian, aka Kullback-Leibler
Philipp Arras's avatar
Philipp Arras committed
41
       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
    Parameters
    ----------
    domain : Domain, DomainTuple or tuple of Domain
52
        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
    """Computes the L2-norm of a Field or MultiField with respect to a
69
    specific kernel given by `endo`.
Philipp Arras's avatar
Philipp Arras committed
70
71
72

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

    Parameters
    ----------
Philipp Arras's avatar
Philipp Arras committed
76
    endo : EndomorphicOperator
77
         Kernel of the 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
    Parameters
    ----------
    mean : Field
        Mean of the Gaussian. Default is 0.
    covariance : LinearOperator
        Covariance of the Gaussian. Default is the identity operator.
Philipp Arras's avatar
Fixup    
Philipp Arras committed
114
    domain : Domain, DomainTuple, tuple of Domain or MultiDomain
Philipp Arras's avatar
Philipp Arras committed
115
116
117
118
119
120
        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
        if mean is not None and not isinstance(mean, Field):
            raise TypeError
        if covariance is not None and not isinstance(covariance,
                                                     LinearOperator):
            raise TypeError

Martin Reinecke's avatar
Martin Reinecke committed
130
131
132
133
134
135
136
137
138
139
        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
140
141
142
143
        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
144
145
146
        self._icov = None if covariance is None else covariance.inverse

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

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


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

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

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

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

    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
181
    """
Philipp Arras's avatar
Philipp Arras committed
182

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

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

201

202
class InverseGammaLikelihood(EnergyOperator):
Philipp Arras's avatar
Philipp Arras committed
203
204
    """
    FIXME
205
    """
Philipp Arras's avatar
Philipp Arras committed
206

207
    def __init__(self, d):
Philipp Arras's avatar
Philipp Arras committed
208
209
        if not isinstance(d, Field):
            raise TypeError
210
211
        self._d = d
        self._domain = DomainTuple.make(d.domain)
212
213

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


Martin Reinecke's avatar
Martin Reinecke committed
224
class BernoulliEnergy(EnergyOperator):
Philipp Arras's avatar
Philipp Arras committed
225
    """Computes likelihood energy of expected event frequency constrained by
226
227
    event data.

Philipp Arras's avatar
Philipp Arras committed
228
229
230
231
232
233
234
    .. math ::
        E(f) = -\\log \\text{Bernoulli}(d|f)
             = -d^\\dagger \\log f  - (1-d)^\\dagger \\log(1-f),

    where f is a field defined on `d.domain` with the expected
    frequencies of events.

235
236
    Parameters
    ----------
Martin Reinecke's avatar
Martin Reinecke committed
237
    d : Field
Philipp Arras's avatar
Philipp Arras committed
238
        Data field with events (1) or non-events (0).
Martin Reinecke's avatar
Martin Reinecke committed
239
    """
Philipp Arras's avatar
Philipp Arras committed
240

241
    def __init__(self, d):
Philipp Arras's avatar
Philipp Arras committed
242
243
244
245
246
        print(d.dtype)
        if not isinstance(d, Field) or not np.issubdtype(d.dtype, np.integer):
            raise TypeError
        if not np.all(np.logical_or(d.local_data == 0, d.local_data == 1)):
            raise ValueError
Martin Reinecke's avatar
Martin Reinecke committed
247
        self._d = d
248
        self._domain = DomainTuple.make(d.domain)
Martin Reinecke's avatar
Martin Reinecke committed
249
250

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


262
class StandardHamiltonian(EnergyOperator):
Philipp Arras's avatar
Philipp Arras committed
263
264
    """Computes an information Hamiltonian in its standard form, i.e. with the
    prior being a Gaussian with unit covariance.
265

Philipp Arras's avatar
Philipp Arras committed
266
    Let the likelihood energy be :math:`E_{lh}`. Then this operator computes:
267

Philipp Arras's avatar
Philipp Arras committed
268
269
    .. math ::
         H(f) = 0.5 f^\\dagger f + E_{lh}(f):
270

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

Martin Reinecke's avatar
Martin Reinecke committed
274
    By implementing prior information this way, the field prior is represented
275
276
277
    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.

Philipp Arras's avatar
Philipp Arras committed
278
279
280
281
282
283
284
285
    The metric of this operator can be used as covariance for drawing Gaussian
    samples.

    Parameters
    ----------
    lh : EnergyOperator
        The likelihood energy.
    ic_samp : IterationController
286
        Tells an internal :class:`SamplingEnabler` which convergence criterion
Philipp Arras's avatar
Philipp Arras committed
287
288
289
290
291
292
293
        to use to draw Gaussian samples.


    See also
    --------
    `Encoding prior knowledge in the structure of the likelihood`,
    Jakob Knollmüller, Torsten A. Ensslin,
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
    def __repr__(self):
        subs = 'Likelihood:\n{}'.format(utilities.indent(self._lh.__repr__()))
        subs += '\nPrior: Quadratic{}'.format(self._lh.domain.keys())
Martin Reinecke's avatar
Martin Reinecke committed
317
        return 'StandardHamiltonian:\n' + utilities.indent(subs)
Philipp Arras's avatar
Philipp Arras committed
318

Martin Reinecke's avatar
Martin Reinecke committed
319

Martin Reinecke's avatar
Martin Reinecke committed
320
class AveragedEnergy(EnergyOperator):
Torsten Ensslin's avatar
Torsten Ensslin committed
321
    """Averages an energy over samples
Martin Reinecke's avatar
Martin Reinecke committed
322

323
324
325
    Parameters
    ----------
    h: Hamiltonian
Philipp Arras's avatar
Philipp Arras committed
326
       The energy to be averaged.
Martin Reinecke's avatar
Martin Reinecke committed
327
    res_samples : iterable of Fields
Torsten Ensslin's avatar
Torsten Ensslin committed
328
329
       Set of residual sample points to be added to mean field for
       approximate estimation of the KL.
330

Philipp Arras's avatar
Philipp Arras committed
331
332
    Note
    ----
Martin Reinecke's avatar
Martin Reinecke committed
333
334
    Having symmetrized residual samples, with both v_i and -v_i being
    present ensures that the distribution mean is exactly represented.
Torsten Ensslin's avatar
Fix te    
Torsten Ensslin committed
335

Philipp Arras's avatar
Philipp Arras committed
336
337
338
339
    :class:`AveragedEnergy(h)` approximates
    :math:`\\left< H(f) \\right>_{G(f-m,D)}` if the residuals
    :math:`f-m` are drawn from a Gaussian distribution with covariance
    :math:`D`.
Martin Reinecke's avatar
Martin Reinecke committed
340
    """
Martin Reinecke's avatar
Martin Reinecke committed
341
342
343

    def __init__(self, h, res_samples):
        self._h = h
Martin Reinecke's avatar
Martin Reinecke committed
344
        self._domain = h.domain
Martin Reinecke's avatar
Martin Reinecke committed
345
346
347
        self._res_samples = tuple(res_samples)

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