energy_adapter.py 8.45 KB
Newer Older
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-2020 Max-Planck-Society
15
16
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
Martin Reinecke's avatar
Martin Reinecke committed
17

18
19
import numpy as np

Philipp Frank's avatar
Philipp Frank committed
20
from .. import random
Martin Reinecke's avatar
Martin Reinecke committed
21
from ..linearization import Linearization
Philipp Arras's avatar
Philipp Arras committed
22
from ..minimization.energy import Energy
Philipp Frank's avatar
Philipp Frank committed
23
24
25
from ..utilities import myassert, allreduce_sum
from ..multi_domain import MultiDomain
from ..sugar import from_random
Philipp Frank's avatar
Philipp Frank committed
26
27
from ..domain_tuple import DomainTuple

Martin Reinecke's avatar
Martin Reinecke committed
28
29

class EnergyAdapter(Energy):
Martin Reinecke's avatar
Martin Reinecke committed
30
31
32
33
34
    """Helper class which provides the traditional Nifty Energy interface to
    Nifty operators with a scalar target domain.

    Parameters
    -----------
Philipp Arras's avatar
Philipp Arras committed
35
36
    position: Field or MultiField
        The position where the minimization process is started.
Philipp Arras's avatar
Philipp Arras committed
37
38
39
    op: EnergyOperator
        The expression computing the energy from the input data.
    constants: list of strings
Martin Reinecke's avatar
Martin Reinecke committed
40
41
42
        The component names of the operator's input domain which are assumed
        to be constant during the minimization process.
        If the operator's input domain is not a MultiField, this must be empty.
Philipp Arras's avatar
Philipp Arras committed
43
44
45
        Default: [].
    want_metric: bool
        If True, the class will provide a `metric` property. This should only
Martin Reinecke's avatar
Martin Reinecke committed
46
        be enabled if it is required, because it will most likely consume
Philipp Arras's avatar
Philipp Arras committed
47
        additional resources. Default: False.
48
49
50
51
52
    nanisinf : bool
        If true, nan energies which can happen due to overflows in the forward
        model are interpreted as inf. Thereby, the code does not crash on
        these occaisions but rather the minimizer is told that the position it
        has tried is not sensible.
Martin Reinecke's avatar
Martin Reinecke committed
53
54
    """

55
    def __init__(self, position, op, constants=[], want_metric=False,
56
57
                 nanisinf=False):
        if len(constants) > 0:
58
            cstpos = position.extract_by_keys(constants)
59
            _, op = op.simplify_for_constant_input(cstpos)
60
61
            varkeys = set(op.domain.keys()) - set(constants)
            position = position.extract_by_keys(varkeys)
62
63
        super(EnergyAdapter, self).__init__(position)
        self._op = op
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
64
        self._want_metric = want_metric
65
        lin = Linearization.make_var(position, want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
66
        tmp = self._op(lin)
Martin Reinecke's avatar
stage2    
Martin Reinecke committed
67
        self._val = tmp.val.val[()]
Martin Reinecke's avatar
Martin Reinecke committed
68
69
        self._grad = tmp.gradient
        self._metric = tmp._metric
70
71
72
        self._nanisinf = bool(nanisinf)
        if self._nanisinf and np.isnan(self._val):
            self._val = np.inf
Martin Reinecke's avatar
Martin Reinecke committed
73
74

    def at(self, position):
75
76
        return EnergyAdapter(position, self._op, want_metric=self._want_metric,
                             nanisinf=self._nanisinf)
Martin Reinecke's avatar
Martin Reinecke committed
77
78
79

    @property
    def value(self):
Martin Reinecke's avatar
Martin Reinecke committed
80
        return self._val
Martin Reinecke's avatar
Martin Reinecke committed
81
82
83

    @property
    def gradient(self):
Martin Reinecke's avatar
Martin Reinecke committed
84
        return self._grad
Martin Reinecke's avatar
Martin Reinecke committed
85

Martin Reinecke's avatar
Martin Reinecke committed
86
87
88
89
    @property
    def metric(self):
        return self._metric

Martin Reinecke's avatar
Martin Reinecke committed
90
91
    def apply_metric(self, x):
        return self._metric(x)
Philipp Frank's avatar
Philipp Frank committed
92
93
94


class StochasticEnergyAdapter(Energy):
Philipp Arras's avatar
Docs    
Philipp Arras committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108
    """Provide the energy interface for an energy operator where parts of the
    input are averaged instead of optimized.

    Specifically, a set of standard normal distributed samples are drawn for
    the input corresponding to `keys` and each sample is inserted partially
    into `op`. The resulting operators are then averaged.  The subdomain that
    is not sampled is left a stochastic average of an energy with the remaining
    subdomain being the DOFs that are considered to be optimization parameters.

    Notes
    -----
    `StochasticEnergyAdapter` should never be created using the constructor,
    but rather via the factory function :attr:`make`.
    """
Philipp Frank's avatar
Philipp Frank committed
109
    def __init__(self, position, op, keys, local_ops, n_samples, comm, nanisinf,
Philipp Frank's avatar
Philipp Frank committed
110
111
112
113
                 _callingfrommake=False):
        if not _callingfrommake:
            raise NotImplementedError
        super(StochasticEnergyAdapter, self).__init__(position)
Philipp Frank's avatar
Philipp Frank committed
114
115
        for lop in local_ops:
            myassert(position.domain == lop.domain)
Philipp Frank's avatar
Philipp Frank committed
116
117
118
119
120
        self._comm = comm
        self._local_ops = local_ops
        self._n_samples = n_samples
        lin = Linearization.make_var(position)
        v, g = [], []
Philipp Frank's avatar
Philipp Frank committed
121
122
        for lop in self._local_ops:
            tmp = lop(lin)
Philipp Frank's avatar
Philipp Frank committed
123
124
125
126
127
128
129
            v.append(tmp.val.val)
            g.append(tmp.gradient)
        self._val = allreduce_sum(v, self._comm)[()]/self._n_samples
        if np.isnan(self._val) and self._nanisinf:
            self._val = np.inf
        self._grad = allreduce_sum(g, self._comm)/self._n_samples

Philipp Frank's avatar
Philipp Frank committed
130
        self._op = op
Philipp Frank's avatar
Philipp Frank committed
131
132
133
134
135
136
137
138
139
140
141
        self._keys = keys

    @property
    def value(self):
        return self._val

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

    def at(self, position):
Philipp Frank's avatar
Philipp Frank committed
142
143
144
        return StochasticEnergyAdapter(position, self._op, self._keys,
                    self._local_ops, self._n_samples, self._comm, self._nanisinf,
                    _callingfrommake=True)
Philipp Frank's avatar
Philipp Frank committed
145
146
147
148
149
150
151
152
153
154

    def apply_metric(self, x):
        lin = Linearization.make_var(self.position, want_metric=True)
        res = []
        for op in self._local_ops:
            res.append(op(lin).metric(x))
        return allreduce_sum(res, self._comm)/self._n_samples

    @property
    def metric(self):
Philipp Frank's avatar
Philipp Frank committed
155
        from .kl_energies import _SelfAdjointOperatorWrapper
Philipp Frank's avatar
Philipp Frank committed
156
157
158
159
160
        return _SelfAdjointOperatorWrapper(self.position.domain,
                                           self.apply_metric)

    def resample_at(self, position):
        return StochasticEnergyAdapter.make(position, self._op, self._keys,
Philipp Frank's avatar
Philipp Frank committed
161
                                            self._n_samples, self._comm)
Philipp Frank's avatar
Philipp Frank committed
162
163

    @staticmethod
Philipp Frank's avatar
Philipp Frank committed
164
    def make(position, op, sampling_keys, n_samples, mirror_samples,
Philipp Arras's avatar
Docs    
Philipp Arras committed
165
166
167
             comm=None, nanisinf=False):
        """Factory function for StochasticEnergyAdapter.

Philipp Frank's avatar
Philipp Frank committed
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
        Parameters
        ----------
        position : MultiField
            Values of the optimization parameters
        op : Operator
            The objective function of the optimization problem. Must have a
            scalar target. The domain must be a `MultiDomain` with its keys
            being the union of `sampling_keys` and `position.domain.keys()`.
        sampling_keys : iterable of String
            The keys of the subdomain over which the stochastic average of `op`
            should be performed.
        n_samples : int
            Number of samples used for the stochastic estimate.
        mirror_samples : boolean
            Whether the negative of the drawn samples are also used, as they are
            equally legitimate samples. If true, the number of used samples
            doubles.
        comm : MPI communicator or None
            If not None, samples will be distributed as evenly as possible
            across this communicator. If `mirror_samples` is set, then a sample
            and its mirror image will always reside on the same task.
        nanisinf : bool
Philipp Arras's avatar
Docs    
Philipp Arras committed
190
191
192
            If true, nan energies, which can occur due to overflows in the
            forward model, are interpreted as inf which can be interpreted by
            optimizers.
Philipp Frank's avatar
Philipp Frank committed
193
        """
Philipp Frank's avatar
Philipp Frank committed
194
        myassert(op.target == DomainTuple.scalar_domain())
Philipp Frank's avatar
Philipp Frank committed
195
        samdom = {}
Philipp Frank's avatar
Philipp Frank committed
196
197
198
199
        if not isinstance(n_samples, int):
            raise TypeError
        for k in sampling_keys:
            if (k in position.domain.keys()) or (k not in op.domain.keys()):
Philipp Frank's avatar
Philipp Frank committed
200
                raise ValueError
Philipp Frank's avatar
Philipp Frank committed
201
            samdom[k] = op.domain[k]
Philipp Frank's avatar
Philipp Frank committed
202
203
204
        samdom = MultiDomain.make(samdom)
        local_ops = []
        sseq = random.spawn_sseq(n_samples)
Philipp Frank's avatar
Philipp Frank committed
205
        from .kl_energies import _get_lo_hi
Philipp Frank's avatar
Philipp Frank committed
206
207
208
209
210
211
212
213
214
        for i in range(*_get_lo_hi(comm, n_samples)):
            with random.Context(sseq[i]):
                rnd = from_random(samdom)
                _, tmp = op.simplify_for_constant_input(rnd)
                myassert(tmp.domain == position.domain)
                local_ops.append(tmp)
                if mirror_samples:
                    local_ops.append(op.simplify_for_constant_input(-rnd)[1])
        n_samples = 2*n_samples if mirror_samples else n_samples
Philipp Frank's avatar
Philipp Frank committed
215
216
        return StochasticEnergyAdapter(position, op, sampling_keys, local_ops,
                              n_samples, comm, nanisinf, _callingfrommake=True)