correlated_fields.py 23.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
Philipp Arras's avatar
Philipp Arras committed
15
# Authors: Philipp Frank, Philipp Arras, Philipp Haim
Martin Reinecke's avatar
Martin Reinecke committed
16
#
17
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
Philipp Arras's avatar
Philipp Arras committed
18

19
20
21
from functools import reduce
from operator import mul

Philipp Arras's avatar
Philipp Arras committed
22
import numpy as np
23

Philipp Arras's avatar
Philipp Arras committed
24
from ..domain_tuple import DomainTuple
Philipp Arras's avatar
Philipp Arras committed
25
26
from ..domains.power_space import PowerSpace
from ..domains.unstructured_domain import UnstructuredDomain
Philipp Arras's avatar
Philipp Arras committed
27
from ..field import Field
28
from ..logger import logger
Philipp Arras's avatar
Philipp Arras committed
29
from ..multi_field import MultiField
Philipp Arras's avatar
Philipp Arras committed
30
from ..operators.adder import Adder
31
from ..operators.contraction_operator import ContractionOperator
Philipp Arras's avatar
Philipp Arras committed
32
from ..operators.diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
33
from ..operators.distributors import PowerDistributor
Philipp Arras's avatar
Philipp Arras committed
34
from ..operators.endomorphic_operator import EndomorphicOperator
Martin Reinecke's avatar
Martin Reinecke committed
35
from ..operators.harmonic_operators import HarmonicTransformOperator
Philipp Arras's avatar
Philipp Arras committed
36
from ..operators.linear_operator import LinearOperator
Philipp Arras's avatar
Philipp Arras committed
37
from ..operators.operator import Operator
Philipp Arras's avatar
Philipp Arras committed
38
from ..operators.simple_linear_operators import ducktape
39
from ..probing import StatCalculator
Martin Reinecke's avatar
Martin Reinecke committed
40
from ..sugar import makeField, full, makeDomain
41

42

Philipp Haim's avatar
Philipp Haim committed
43
44
45
def _reshaper(x, N):
    x = np.asfarray(x)
    if x.shape in [(), (1,)]:
Philipp Haim's avatar
Philipp Haim committed
46
        return np.full(N, x) if N != 0 else x.reshape(())
Philipp Haim's avatar
Philipp Haim committed
47
48
    elif x.shape == (N,):
        return x
49
50
51
    else:
        raise TypeError("Shape of parameters cannot be interpreted")

52

Martin Reinecke's avatar
Martin Reinecke committed
53
def _lognormal_moments(mean, sig, N=0):
Philipp Haim's avatar
Philipp Haim committed
54
55
56
57
    if N == 0:
        mean, sig = np.asfarray(mean), np.asfarray(sig)
    else:
        mean, sig = (_reshaper(param, N) for param in (mean, sig))
58
    if not np.all(mean > 0):
59
        raise ValueError("mean must be greater 0; got {!r}".format(mean))
60
    if not np.all(sig > 0):
61
        raise ValueError("sig must be greater 0; got {!r}".format(sig))
62

Philipp Arras's avatar
Philipp Arras committed
63
64
    logsig = np.sqrt(np.log((sig/mean)**2 + 1))
    logmean = np.log(mean) - logsig**2/2
65
    return logmean, logsig
Philipp Arras's avatar
Philipp Arras committed
66
67


Martin Reinecke's avatar
Martin Reinecke committed
68
def _normal(mean, sig, key, N=0):
Philipp Haim's avatar
Philipp Haim committed
69
    if N == 0:
Philipp Haim's avatar
Philipp Haim committed
70
        domain = DomainTuple.scalar_domain()
Philipp Haim's avatar
Philipp Haim committed
71
        mean, sig = np.asfarray(mean), np.asfarray(sig)
Philipp Haim's avatar
Philipp Haim committed
72
73
    else:
        domain = UnstructuredDomain(N)
Philipp Haim's avatar
Philipp Haim committed
74
        mean, sig = (_reshaper(param, N) for param in (mean, sig))
Martin Reinecke's avatar
Martin Reinecke committed
75
76
    return Adder(makeField(domain, mean)) @ (DiagonalOperator(
        makeField(domain, sig)) @ ducktape(domain, None, key))
Philipp Arras's avatar
Philipp Arras committed
77
78


Philipp Arras's avatar
Philipp Arras committed
79
def _log_k_lengths(pspace):
Philipp Arras's avatar
Philipp Arras committed
80
    """Log(k_lengths) without zeromode"""
Philipp Arras's avatar
Philipp Arras committed
81
82
83
    return np.log(pspace.k_lengths[1:])


Philipp Arras's avatar
Philipp Arras committed
84
def _relative_log_k_lengths(power_space):
Philipp Arras's avatar
Philipp Arras committed
85
86
    """Log-distance to first bin
    logkl.shape==power_space.shape, logkl[0]=logkl[1]=0"""
Philipp Arras's avatar
Philipp Arras committed
87
88
89
90
91
92
    power_space = DomainTuple.make(power_space)
    assert isinstance(power_space[0], PowerSpace)
    assert len(power_space) == 1
    logkl = _log_k_lengths(power_space[0])
    assert logkl.shape[0] == power_space[0].shape[0] - 1
    logkl -= logkl[0]
Philipp Arras's avatar
Philipp Arras committed
93
    return np.insert(logkl, 0, 0)
Philipp Arras's avatar
Philipp Arras committed
94
95


Philipp Arras's avatar
Philipp Arras committed
96
def _log_vol(power_space):
97
    power_space = makeDomain(power_space)
Philipp Arras's avatar
Philipp Arras committed
98
99
100
101
102
    assert isinstance(power_space[0], PowerSpace)
    logk_lengths = _log_k_lengths(power_space[0])
    return logk_lengths[1:] - logk_lengths[:-1]


Philipp Haim's avatar
Philipp Haim committed
103
def _total_fluctuation_realized(samples):
104
105
    res = 0.
    for s in samples:
Philipp Haim's avatar
Fixes  
Philipp Haim committed
106
        res = res + (s - s.mean())**2
Philipp Haim's avatar
Philipp Haim committed
107
    return np.sqrt((res/len(samples)).mean())
108
109


Philipp Arras's avatar
Philipp Arras committed
110
class _LognormalMomentMatching(Operator):
Philipp Haim's avatar
Philipp Haim committed
111
    def __init__(self, mean, sig, key, N_copies):
Philipp Arras's avatar
Philipp Arras committed
112
        key = str(key)
Philipp Haim's avatar
Philipp Haim committed
113
        logmean, logsig = _lognormal_moments(mean, sig, N_copies)
Philipp Arras's avatar
Philipp Arras committed
114
115
        self._mean = mean
        self._sig = sig
Philipp Haim's avatar
Philipp Haim committed
116
        op = _normal(logmean, logsig, key, N_copies).exp()
Philipp Arras's avatar
Philipp Arras committed
117
118
119
120
121
122
123
124
125
126
        self._domain, self._target = op.domain, op.target
        self.apply = op.apply

    @property
    def mean(self):
        return self._mean

    @property
    def std(self):
        return self._sig
Philipp Arras's avatar
Philipp Arras committed
127
128


Philipp Frank's avatar
Philipp Frank committed
129
class _SlopeRemover(EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
130
    def __init__(self, domain, space=0):
Philipp Frank's avatar
Philipp Frank committed
131
        self._domain = makeDomain(domain)
132
133
        assert isinstance(self._domain[space], PowerSpace)
        logkl = _relative_log_k_lengths(self._domain[space])
134
        self._sc = logkl/float(logkl[-1])
Philipp Arras's avatar
Philipp Arras committed
135

136
        self._space = space
Philipp Haim's avatar
Philipp Haim committed
137
138
139
        axis = self._domain.axes[space][0]
        self._last = (slice(None),)*axis + (-1,) + (None,)
        self._extender = (None,)*(axis) + (slice(None),) + (None,)*(self._domain.axes[-1][-1]-axis)
Philipp Frank's avatar
Philipp Frank committed
140
        self._capability = self.TIMES | self.ADJOINT_TIMES
Philipp Arras's avatar
Philipp Arras committed
141

142
143
    def apply(self, x, mode):
        self._check_input(x, mode)
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
144
        x = x.val
Philipp Frank's avatar
Philipp Frank committed
145
        if mode == self.TIMES:
Philipp Haim's avatar
Philipp Haim committed
146
            res = x - x[self._last]*self._sc[self._extender]
Philipp Frank's avatar
Philipp Frank committed
147
        else:
148
149
            res = x.copy()
            res[self._last] -= (x*self._sc[self._extender]).sum(
Philipp Arras's avatar
Philipp Arras committed
150
                axis=self._space, keepdims=True)
Martin Reinecke's avatar
Martin Reinecke committed
151
        return makeField(self._tgt(mode), res)
Philipp Frank's avatar
Philipp Frank committed
152

Philipp Arras's avatar
Philipp Arras committed
153
154

class _TwoLogIntegrations(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
155
    def __init__(self, target, space=0):
Philipp Arras's avatar
Philipp Arras committed
156
        self._target = makeDomain(target)
157
158
159
160
161
        assert isinstance(self.target[space], PowerSpace)
        dom = list(self._target)
        dom[space] = UnstructuredDomain((2, self.target[space].shape[0]-2))
        self._domain = makeDomain(dom)
        self._space = space
162
        self._log_vol = _log_vol(self._target[space])
Philipp Arras's avatar
Philipp Arras committed
163
164
165
166
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
167

Martin Reinecke's avatar
Martin Reinecke committed
168
        # Maybe make class properties
169
170
        axis = self._target.axes[self._space][0]
        sl = (slice(None),)*axis
Philipp Haim's avatar
Fixes  
Philipp Haim committed
171
        extender_sl = (None,)*axis + (slice(None),) + (None,)*(self._target.axes[-1][-1] - axis)
172
173
        first = sl + (0,)
        second = sl + (1,)
Martin Reinecke's avatar
Martin Reinecke committed
174
175
176
        from_third = sl + (slice(2, None),)
        no_border = sl + (slice(1, -1),)
        reverse = sl + (slice(None, None, -1),)
177

Philipp Arras's avatar
Philipp Arras committed
178
        if mode == self.TIMES:
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
179
            x = x.val
Philipp Arras's avatar
Philipp Arras committed
180
            res = np.empty(self._target.shape)
181
            res[first] = res[second] = 0
Martin Reinecke's avatar
Martin Reinecke committed
182
            res[from_third] = np.cumsum(x[second], axis=axis)
183
            res[from_third] = (res[from_third] + res[no_border])/2*self._log_vol[extender_sl] + x[first]
Martin Reinecke's avatar
Martin Reinecke committed
184
            res[from_third] = np.cumsum(res[from_third], axis=axis)
Philipp Arras's avatar
Philipp Arras committed
185
        else:
Martin Reinecke's avatar
Martin Reinecke committed
186
            x = x.val_rw()
Philipp Arras's avatar
Philipp Arras committed
187
            res = np.zeros(self._domain.shape)
Martin Reinecke's avatar
Martin Reinecke committed
188
            x[from_third] = np.cumsum(x[from_third][reverse], axis=axis)[reverse]
189
            res[first] += x[from_third]
190
            x[from_third] *= (self._log_vol/2.)[extender_sl]
191
            x[no_border] += x[from_third]
Martin Reinecke's avatar
Martin Reinecke committed
192
            res[second] += np.cumsum(x[from_third][reverse], axis=axis)[reverse]
Martin Reinecke's avatar
Martin Reinecke committed
193
        return makeField(self._tgt(mode), res)
Philipp Arras's avatar
Philipp Arras committed
194
195
196


class _Normalization(Operator):
Martin Reinecke's avatar
Martin Reinecke committed
197
    def __init__(self, domain, space=0):
Philipp Arras's avatar
Philipp Arras committed
198
        self._domain = self._target = makeDomain(domain)
199
        assert isinstance(self._domain[space], PowerSpace)
200
201
202
        hspace = list(self._domain)
        hspace[space] = hspace[space].harmonic_partner
        hspace = makeDomain(hspace)
Philipp Arras's avatar
Philipp Arras committed
203
204
205
        pd = PowerDistributor(hspace,
                              power_space=self._domain[space],
                              space=space)
Martin Reinecke's avatar
Martin Reinecke committed
206
        mode_multiplicity = pd.adjoint(full(pd.target, 1.)).val_rw()
207
        zero_mode = (slice(None),)*self._domain.axes[space][0] + (0,)
Philipp Haim's avatar
Philipp Haim committed
208
        mode_multiplicity[zero_mode] = 0
Philipp Arras's avatar
Philipp Arras committed
209
        self._mode_multiplicity = makeField(self._domain, mode_multiplicity)
210
        self._specsum = _SpecialSum(self._domain, space)
Philipp Arras's avatar
Philipp Arras committed
211
212
213
214
215
216
217

    def apply(self, x):
        self._check_input(x)
        amp = x.exp()
        spec = (2*x).exp()
        # FIXME This normalizes also the zeromode which is supposed to be left
        # untouched by this operator
218
        return self._specsum(self._mode_multiplicity*spec)**(-0.5)*amp
Philipp Arras's avatar
Philipp Arras committed
219
220
221


class _SpecialSum(EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
222
    def __init__(self, domain, space=0):
Philipp Arras's avatar
Philipp Arras committed
223
224
        self._domain = makeDomain(domain)
        self._capability = self.TIMES | self.ADJOINT_TIMES
225
        self._contractor = ContractionOperator(domain, space)
Philipp Arras's avatar
Philipp Arras committed
226
227
228

    def apply(self, x, mode):
        self._check_input(x, mode)
229
        return self._contractor.adjoint(self._contractor(x))
Philipp Arras's avatar
Philipp Arras committed
230
231


Philipp Haim's avatar
Philipp Haim committed
232
class _Distributor(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
233
    def __init__(self, dofdex, domain, target, space=0):
Philipp Haim's avatar
Philipp Haim committed
234
235
236
237
238
239
240
241
242
        self._dofdex = dofdex

        self._target = makeDomain(target)
        self._domain = makeDomain(domain)
        self._sl = (slice(None),)*space
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
243
        x = x.val
Philipp Haim's avatar
Philipp Haim committed
244
245
246
247
248
        if mode == self.TIMES:
            res = x[self._dofdex]
        else:
            res = np.empty(self._tgt(mode).shape)
            res[self._dofdex] = x
Martin Reinecke's avatar
Martin Reinecke committed
249
        return makeField(self._tgt(mode), res)
Martin Reinecke's avatar
Martin Reinecke committed
250

251

252
253
class _Amplitude(Operator):
    def __init__(self, target, fluctuations, flexibility, asperity,
254
                 loglogavgslope, azm, totvol, key, dofdex):
Philipp Arras's avatar
Philipp Arras committed
255
256
257
258
259
260
261
262
263
264
265
        """
        fluctuations > 0
        flexibility > 0
        asperity > 0
        loglogavgslope probably negative
        """
        assert isinstance(fluctuations, Operator)
        assert isinstance(flexibility, Operator)
        assert isinstance(asperity, Operator)
        assert isinstance(loglogavgslope, Operator)

Philipp Haim's avatar
Philipp Haim committed
266
267
        if len(dofdex) > 0:
            N_copies = max(dofdex) + 1
Philipp Haim's avatar
Philipp Haim committed
268
            space = 1
Philipp Frank's avatar
cleanup  
Philipp Frank committed
269
270
            distributed_tgt = makeDomain((UnstructuredDomain(len(dofdex)),
                                          target))
Philipp Haim's avatar
Philipp Haim committed
271
272
273
            target = makeDomain((UnstructuredDomain(N_copies), target))
            Distributor = _Distributor(dofdex, target, distributed_tgt, 0)
        else:
Philipp Haim's avatar
Philipp Haim committed
274
            N_copies = 0
Philipp Haim's avatar
Philipp Haim committed
275
            space = 0
Philipp Haim's avatar
Philipp Haim committed
276
            distributed_tgt = target = makeDomain(target)
Martin Reinecke's avatar
Martin Reinecke committed
277
        azm_expander = ContractionOperator(distributed_tgt, spaces=space).adjoint
Philipp Haim's avatar
Philipp Haim committed
278
        assert isinstance(target[space], PowerSpace)
Martin Reinecke's avatar
Martin Reinecke committed
279

280
        twolog = _TwoLogIntegrations(target, space)
Philipp Arras's avatar
Philipp Arras committed
281
        dom = twolog.domain
282

283
        shp = dom[space].shape
Martin Reinecke's avatar
Martin Reinecke committed
284
285
        expander = ContractionOperator(dom, spaces=space).adjoint
        ps_expander = ContractionOperator(twolog.target, spaces=space).adjoint
Philipp Arras's avatar
Philipp Arras committed
286
287
288

        # Prepare constant fields
        foo = np.zeros(shp)
289
        foo[0] = foo[1] = np.sqrt(_log_vol(target[space]))
Martin Reinecke's avatar
Martin Reinecke committed
290
        vflex = DiagonalOperator(makeField(dom[space], foo), dom, space)
Philipp Arras's avatar
Philipp Arras committed
291
292
293

        foo = np.zeros(shp, dtype=np.float64)
        foo[0] += 1
Martin Reinecke's avatar
Martin Reinecke committed
294
        vasp = DiagonalOperator(makeField(dom[space], foo), dom, space)
Philipp Arras's avatar
Philipp Arras committed
295
296

        foo = np.ones(shp)
297
        foo[0] = _log_vol(target[space])**2/12.
Martin Reinecke's avatar
Martin Reinecke committed
298
        shift = DiagonalOperator(makeField(dom[space], foo), dom, space)
Martin Reinecke's avatar
Martin Reinecke committed
299

300
        vslope = DiagonalOperator(
Philipp Arras's avatar
Philipp Arras committed
301
            makeField(target[space], _relative_log_k_lengths(target[space])),
Martin Reinecke's avatar
Martin Reinecke committed
302
            target, space)
303
304

        foo, bar = [np.zeros(target[space].shape) for _ in range(2)]
Philipp Arras's avatar
Philipp Arras committed
305
        bar[1:] = foo[0] = totvol
Philipp Arras's avatar
Philipp Arras committed
306
307
308
309
        vol0, vol1 = [
            DiagonalOperator(makeField(target[space], aa), target, space)
            for aa in (foo, bar)
        ]
310

Martin Reinecke's avatar
Martin Reinecke committed
311
        # Prepare fields for Adder
312
        shift, vol0 = [op(full(op.domain, 1)) for op in (shift, vol0)]
Philipp Arras's avatar
Philipp Arras committed
313
314
        # End prepare constant fields

315
316
317
318
        slope = vslope @ ps_expander @ loglogavgslope
        sig_flex = vflex @ expander @ flexibility
        sig_asp = vasp @ expander @ asperity
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Haim's avatar
Philipp Haim committed
319
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Arras's avatar
Philipp Arras committed
320
321

        xi = ducktape(dom, None, key)
Philipp Arras's avatar
Philipp Arras committed
322
        sigma = sig_flex*(Adder(shift) @ sig_asp).sqrt()
323
324
        smooth = _SlopeRemover(target, space) @ twolog @ (sigma*xi)
        op = _Normalization(target, space) @ (slope + smooth)
Philipp Haim's avatar
Philipp Haim committed
325
        if N_copies > 0:
Philipp Haim's avatar
Philipp Haim committed
326
327
            op = Distributor @ op
            sig_fluc = Distributor @ sig_fluc
Philipp Haim's avatar
Philipp Haim committed
328
            op = Adder(Distributor(vol0)) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Philipp Arras's avatar
Philipp Arras committed
329
330
            self._fluc = (_Distributor(dofdex, fluctuations.target,
                                       distributed_tgt[0]) @ fluctuations)
Philipp Haim's avatar
Philipp Haim committed
331
        else:
Philipp Frank's avatar
cleanup  
Philipp Frank committed
332
            op = Adder(vol0) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Philipp Frank's avatar
fixup  
Philipp Frank committed
333
            self._fluc = fluctuations
Philipp Arras's avatar
Philipp Arras committed
334

Philipp Arras's avatar
Philipp Arras committed
335
336
        self.apply = op.apply
        self._domain, self._target = op.domain, op.target
337
        self._space = space
Philipp Arras's avatar
Philipp Arras committed
338

Philipp Arras's avatar
Philipp Arras committed
339
340
341
342
    @property
    def fluctuation_amplitude(self):
        return self._fluc

343
344

class CorrelatedFieldMaker:
Philipp Haim's avatar
Philipp Haim committed
345
    def __init__(self, amplitude_offset, prefix, total_N):
346
347
        if not isinstance(amplitude_offset, Operator):
            raise TypeError("amplitude_offset needs to be an operator")
348
        self._a = []
349
        self._position_spaces = []
Philipp Arras's avatar
Formats  
Philipp Arras committed
350

351
352
        self._azm = amplitude_offset
        self._prefix = prefix
Philipp Haim's avatar
Philipp Haim committed
353
        self._total_N = total_N
Philipp Arras's avatar
Formats  
Philipp Arras committed
354

355
    @staticmethod
Philipp Frank's avatar
Philipp Frank committed
356
    def make(offset_amplitude_mean, offset_amplitude_stddev, prefix,
Martin Reinecke's avatar
Martin Reinecke committed
357
358
             total_N=0,
             dofdex=None):
Philipp Frank's avatar
Philipp Frank committed
359
360
        if dofdex is None:
            dofdex = np.full(total_N, 0)
361
362
        elif len(dofdex) != total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Frank's avatar
Philipp Frank committed
363
        N = max(dofdex) + 1 if total_N > 0 else 0
364
365
        zm = _LognormalMomentMatching(offset_amplitude_mean,
                                      offset_amplitude_stddev,
Philipp Haim's avatar
Philipp Haim committed
366
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
367
                                      N)
Philipp Frank's avatar
fixup  
Philipp Frank committed
368
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
369
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
Philipp Haim's avatar
Philipp Haim committed
370
        return CorrelatedFieldMaker(zm, prefix, total_N)
371
372

    def add_fluctuations(self,
373
                         position_space,
374
375
376
377
378
379
380
381
                         fluctuations_mean,
                         fluctuations_stddev,
                         flexibility_mean,
                         flexibility_stddev,
                         asperity_mean,
                         asperity_stddev,
                         loglogavgslope_mean,
                         loglogavgslope_stddev,
Martin Reinecke's avatar
Martin Reinecke committed
382
383
384
385
                         prefix='',
                         index=None,
                         dofdex=None,
                         harmonic_partner=None):
Philipp Frank's avatar
Philipp Frank committed
386
387
        if harmonic_partner is None:
            harmonic_partner = position_space.get_default_codomain()
Philipp Frank's avatar
Fixup  
Philipp Frank committed
388
389
390
        else:
            position_space.check_codomain(harmonic_partner)
            harmonic_partner.check_codomain(position_space)
391

Philipp Haim's avatar
Philipp Haim committed
392
393
        if dofdex is None:
            dofdex = np.full(self._total_N, 0)
394
395
        elif len(dofdex) != self._total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Haim's avatar
Philipp Haim committed
396

Philipp Haim's avatar
Philipp Haim committed
397
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
398
            space = 1
Philipp Haim's avatar
Philipp Haim committed
399
400
            N = max(dofdex) + 1
            position_space = makeDomain((UnstructuredDomain(N), position_space))
Philipp Haim's avatar
Philipp Haim committed
401
402
        else:
            space = 0
Philipp Haim's avatar
Philipp Haim committed
403
            N = 0
Philipp Haim's avatar
Philipp Haim committed
404
            position_space = makeDomain(position_space)
Philipp Arras's avatar
Philipp Arras committed
405
        prefix = str(prefix)
Martin Reinecke's avatar
Martin Reinecke committed
406
        # assert isinstance(position_space[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
407

Philipp Arras's avatar
Philipp Arras committed
408
409
        fluct = _LognormalMomentMatching(fluctuations_mean,
                                         fluctuations_stddev,
410
                                         self._prefix + prefix + 'fluctuations',
Philipp Haim's avatar
Philipp Haim committed
411
                                         N)
Philipp Arras's avatar
Philipp Arras committed
412
        flex = _LognormalMomentMatching(flexibility_mean, flexibility_stddev,
413
                                        self._prefix + prefix + 'flexibility',
Philipp Haim's avatar
Philipp Haim committed
414
                                        N)
Philipp Arras's avatar
Philipp Arras committed
415
        asp = _LognormalMomentMatching(asperity_mean, asperity_stddev,
Philipp Arras's avatar
Philipp Arras committed
416
                                       self._prefix + prefix + 'asperity', N)
417
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
418
                        self._prefix + prefix + 'loglogavgslope', N)
Philipp Arras's avatar
Philipp Arras committed
419
420
        amp = _Amplitude(PowerSpace(harmonic_partner), fluct, flex, asp, avgsl,
                         self._azm, position_space[-1].total_volume,
421
                         self._prefix + prefix + 'spectrum', dofdex)
Philipp Haim's avatar
Philipp Haim committed
422

423
424
        if index is not None:
            self._a.insert(index, amp)
425
            self._position_spaces.insert(index, position_space)
426
427
        else:
            self._a.append(amp)
428
            self._position_spaces.append(position_space)
429

Philipp Frank's avatar
fixup  
Philipp Frank committed
430
    def _finalize_from_op(self):
Philipp Haim's avatar
Philipp Haim committed
431
        n_amplitudes = len(self._a)
Philipp Haim's avatar
Philipp Haim committed
432
        if self._total_N > 0:
Philipp Arras's avatar
Philipp Arras committed
433
434
435
            hspace = makeDomain(
                [UnstructuredDomain(self._total_N)] +
                [dd.target[-1].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
436
437
            spaces = tuple(range(1, n_amplitudes + 1))
            amp_space = 1
Philipp Haim's avatar
Philipp Haim committed
438
439
        else:
            hspace = makeDomain(
Philipp Arras's avatar
Philipp Arras committed
440
                [dd.target[0].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
441
            spaces = tuple(range(n_amplitudes))
Philipp Haim's avatar
Philipp Haim committed
442
            amp_space = 0
443

Martin Reinecke's avatar
Martin Reinecke committed
444
        expander = ContractionOperator(hspace, spaces=spaces).adjoint
Philipp Frank's avatar
fixup  
Philipp Frank committed
445
        azm = expander @ self._azm
446

447
        ht = HarmonicTransformOperator(hspace,
Philipp Haim's avatar
Philipp Haim committed
448
                                       self._position_spaces[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
449
                                       space=spaces[0])
450
        for i in range(1, n_amplitudes):
451
452
453
454
455
456
457
            ht = HarmonicTransformOperator(ht.target,
                                           self._position_spaces[i][amp_space],
                                           space=spaces[i]) @ ht
        a = []
        for ii in range(n_amplitudes):
            co = ContractionOperator(hspace, spaces[:ii] + spaces[ii + 1:])
            pp = self._a[ii].target[amp_space]
Philipp Haim's avatar
Philipp Haim committed
458
            pd = PowerDistributor(co.target, pp, amp_space)
459
460
461
            a.append(co.adjoint @ pd @ self._a[ii])
        corr = reduce(mul, a)
        return ht(azm*corr*ducktape(hspace, None, self._prefix + 'xi'))
Philipp Arras's avatar
Philipp Arras committed
462

Philipp Arras's avatar
Formats  
Philipp Arras committed
463
    def finalize(self, offset=None, prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
464
465
466
        """
        offset vs zeromode: volume factor
        """
Philipp Frank's avatar
fixup  
Philipp Frank committed
467
        op = self._finalize_from_op()
Philipp Arras's avatar
Philipp Arras committed
468
        if offset is not None:
469
470
471
472
473
474
475
            # Deviations from this offset must not be considered here as they
            # are learned by the zeromode
            if isinstance(offset, (Field, MultiField)):
                op = Adder(offset) @ op
            else:
                offset = float(offset)
                op = Adder(full(op.target, offset)) @ op
476
        self.statistics_summary(prior_info)
477
478
        return op

479
480
481
482
483
484
    def statistics_summary(self, prior_info):
        from ..sugar import from_random

        if prior_info == 0:
            return

485
486
        lst = [('Offset amplitude', self.amplitude_total_offset),
               ('Total fluctuation amplitude', self.total_fluctuation)]
487
        namps = len(self._a)
488
489
490
491
492
493
494
495
        if namps > 1:
            for ii in range(namps):
                lst.append(('Slice fluctuation (space {})'.format(ii),
                            self.slice_fluctuation(ii)))
                lst.append(('Average fluctuation (space {})'.format(ii),
                            self.average_fluctuation(ii)))

        for kk, op in lst:
496
497
498
            sc = StatCalculator()
            for _ in range(prior_info):
                sc.add(op(from_random('normal', op.domain)))
Martin Reinecke's avatar
merge  
Martin Reinecke committed
499
500
            mean = sc.mean.val
            stddev = sc.var.sqrt().val
501
            for m, s in zip(mean.flatten(), stddev.flatten()):
502
                logger.info('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
503
504
505

    def moment_slice_to_average(self, fluctuations_slice_mean, nsamples=1000):
        fluctuations_slice_mean = float(fluctuations_slice_mean)
506
507
508
        if not fluctuations_slice_mean > 0:
            msg = "fluctuations_slice_mean must be greater zero; got {!r}"
            raise ValueError(msg.format(fluctuations_slice_mean))
509
        from ..sugar import from_random
510
511
        scm = 1.
        for a in self._a:
Philipp Haim's avatar
Philipp Haim committed
512
            op = a.fluctuation_amplitude*self._azm.one_over()
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
513
            res = np.array([op(from_random('normal', op.domain)).val
514
515
                            for _ in range(nsamples)])
            scm *= res**2 + 1.
516
        return fluctuations_slice_mean/np.mean(np.sqrt(scm))
517

Philipp Arras's avatar
Philipp Arras committed
518
    @property
Philipp Haim's avatar
Philipp Haim committed
519
    def normalized_amplitudes(self):
520
        return self._a
Philipp Arras's avatar
Philipp Arras committed
521

Philipp Haim's avatar
Philipp Haim committed
522
523
524
525
526
527
528
    @property
    def amplitude(self):
        if len(self._a) > 1:
            s = ('If more than one spectrum is present in the model,',
                 ' no unique set of amplitudes exist because only the',
                 ' relative scale is determined.')
            raise NotImplementedError(s)
Philipp Haim's avatar
Fix  
Philipp Haim committed
529
530
        dom = self._a[0].target
        expand = ContractionOperator(dom, len(dom)-1).adjoint
Philipp Haim's avatar
Philipp Haim committed
531
532
        return self._a[0]*(expand @ self.amplitude_total_offset)

533
534
535
    @property
    def amplitude_total_offset(self):
        return self._azm
Philipp Arras's avatar
Philipp Arras committed
536
537

    @property
538
    def total_fluctuation(self):
539
        """Returns operator which acts on prior or posterior samples"""
540
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
541
            raise NotImplementedError
542
        if len(self._a) == 1:
543
            return self.average_fluctuation(0)
544
545
        q = 1.
        for a in self._a:
Philipp Haim's avatar
Philipp Haim committed
546
            fl = a.fluctuation_amplitude*self._azm.one_over()
Philipp Arras's avatar
Philipp Arras committed
547
            q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Philipp Arras's avatar
Formats  
Philipp Arras committed
548
        return (Adder(full(q.target, -1.)) @ q).sqrt()*self._azm
549

Philipp Arras's avatar
Philipp Arras committed
550
    def slice_fluctuation(self, space):
551
        """Returns operator which acts on prior or posterior samples"""
552
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
553
            raise NotImplementedError
554
        if space >= len(self._a):
555
            raise ValueError("invalid space specified; got {!r}".format(space))
556
        if len(self._a) == 1:
557
            return self.average_fluctuation(0)
558
559
        q = 1.
        for j in range(len(self._a)):
Philipp Haim's avatar
Philipp Haim committed
560
            fl = self._a[j].fluctuation_amplitude*self._azm.one_over()
561
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
562
                q = q*fl**2
563
            else:
Philipp Arras's avatar
Philipp Arras committed
564
                q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Philipp Arras's avatar
Formats  
Philipp Arras committed
565
        return q.sqrt()*self._azm
Philipp Arras's avatar
Philipp Arras committed
566
567

    def average_fluctuation(self, space):
568
        """Returns operator which acts on prior or posterior samples"""
569
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
570
            raise NotImplementedError
571
        if space >= len(self._a):
572
            raise ValueError("invalid space specified; got {!r}".format(space))
573
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
574
575
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
576

577
578
    @staticmethod
    def offset_amplitude_realized(samples):
579
580
        res = 0.
        for s in samples:
Philipp Frank's avatar
fixes  
Philipp Frank committed
581
            res = res + s.mean()**2
582
        return np.sqrt(res/len(samples))
Philipp Arras's avatar
Philipp Arras committed
583

584
585
586
587
588
589
590
591
    @staticmethod
    def total_fluctuation_realized(samples):
        return _total_fluctuation_realized(samples)

    @staticmethod
    def slice_fluctuation_realized(samples, space):
        """Computes slice fluctuations from collection of field (defined in signal
        space) realizations."""
592
        ldom = len(samples[0].domain)
593
        if space >= ldom:
594
            raise ValueError("invalid space specified; got {!r}".format(space))
595
        if ldom == 1:
596
            return _total_fluctuation_realized(samples)
Philipp Arras's avatar
Philipp Arras committed
597
        res1, res2 = 0., 0.
598
        for s in samples:
Philipp Frank's avatar
fixes  
Philipp Frank committed
599
600
601
602
603
604
605
            res1 = res1 + s**2
            res2 = res2 + s.mean(space)**2
        res1 = res1/len(samples)
        res2 = res2/len(samples)
        res = res1.mean() - res2.mean()
        return np.sqrt(res)

Philipp Arras's avatar
Philipp Arras committed
606
    @staticmethod
607
608
609
610
    def average_fluctuation_realized(samples, space):
        """Computes average fluctuations from collection of field (defined in signal
        space) realizations."""
        ldom = len(samples[0].domain)
611
        if space >= ldom:
612
            raise ValueError("invalid space specified; got {!r}".format(space))
613
614
615
616
617
618
        if ldom == 1:
            return _total_fluctuation_realized(samples)
        spaces = ()
        for i in range(ldom):
            if i != space:
                spaces += (i,)
Philipp Arras's avatar
Philipp Arras committed
619
620
        res = 0.
        for s in samples:
621
622
623
624
            r = s.mean(spaces)
            res = res + (r - r.mean())**2
        res = res/len(samples)
        return np.sqrt(res.mean())