correlated_fields.py 23 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

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

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

38

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

48

Martin Reinecke's avatar
Martin Reinecke committed
49
def _lognormal_moments(mean, sig, N=0):
Philipp Haim's avatar
Philipp Haim committed
50
51
52
53
    if N == 0:
        mean, sig = np.asfarray(mean), np.asfarray(sig)
    else:
        mean, sig = (_reshaper(param, N) for param in (mean, sig))
Martin Reinecke's avatar
Martin Reinecke committed
54
    assert np.all(mean > 0)
55
    assert np.all(sig > 0)
Philipp Arras's avatar
Philipp Arras committed
56
57
    logsig = np.sqrt(np.log((sig/mean)**2 + 1))
    logmean = np.log(mean) - logsig**2/2
58
    return logmean, logsig
Philipp Arras's avatar
Philipp Arras committed
59
60


Martin Reinecke's avatar
Martin Reinecke committed
61
def _normal(mean, sig, key, N=0):
Philipp Haim's avatar
Philipp Haim committed
62
    if N == 0:
Philipp Haim's avatar
Philipp Haim committed
63
        domain = DomainTuple.scalar_domain()
Philipp Haim's avatar
Philipp Haim committed
64
        mean, sig = np.asfarray(mean), np.asfarray(sig)
Philipp Haim's avatar
Philipp Haim committed
65
66
    else:
        domain = UnstructuredDomain(N)
Philipp Haim's avatar
Philipp Haim committed
67
        mean, sig = (_reshaper(param, N) for param in (mean, sig))
68
    return Adder(from_global_data(domain, mean)) @ (
Martin Reinecke's avatar
Martin Reinecke committed
69
        DiagonalOperator(from_global_data(domain, sig))
70
        @ ducktape(domain, None, key))
Philipp Arras's avatar
Philipp Arras committed
71
72


Philipp Arras's avatar
Philipp Arras committed
73
def _log_k_lengths(pspace):
Philipp Arras's avatar
Philipp Arras committed
74
    """Log(k_lengths) without zeromode"""
Philipp Arras's avatar
Philipp Arras committed
75
76
77
    return np.log(pspace.k_lengths[1:])


Philipp Arras's avatar
Philipp Arras committed
78
def _relative_log_k_lengths(power_space):
Philipp Arras's avatar
Philipp Arras committed
79
80
    """Log-distance to first bin
    logkl.shape==power_space.shape, logkl[0]=logkl[1]=0"""
Philipp Arras's avatar
Philipp Arras committed
81
82
83
84
85
86
    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
87
    return np.insert(logkl, 0, 0)
Philipp Arras's avatar
Philipp Arras committed
88
89


Philipp Arras's avatar
Philipp Arras committed
90
def _log_vol(power_space):
91
    power_space = makeDomain(power_space)
Philipp Arras's avatar
Philipp Arras committed
92
93
94
95
96
    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
97
def _total_fluctuation_realized(samples):
98
99
    res = 0.
    for s in samples:
Philipp Haim's avatar
Fixes    
Philipp Haim committed
100
        res = res + (s - s.mean())**2
Philipp Haim's avatar
Philipp Haim committed
101
    return np.sqrt((res/len(samples)).mean())
102
103


Philipp Arras's avatar
Philipp Arras committed
104
class _LognormalMomentMatching(Operator):
Philipp Haim's avatar
Philipp Haim committed
105
    def __init__(self, mean, sig, key, N_copies):
Philipp Arras's avatar
Philipp Arras committed
106
        key = str(key)
Philipp Haim's avatar
Philipp Haim committed
107
        logmean, logsig = _lognormal_moments(mean, sig, N_copies)
Philipp Arras's avatar
Philipp Arras committed
108
109
        self._mean = mean
        self._sig = sig
Philipp Haim's avatar
Philipp Haim committed
110
        op = _normal(logmean, logsig, key, N_copies).exp()
Philipp Arras's avatar
Philipp Arras committed
111
112
113
114
115
116
117
118
119
120
        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
121
122


Philipp Frank's avatar
Philipp Frank committed
123
class _SlopeRemover(EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
124
    def __init__(self, domain, space=0):
Philipp Frank's avatar
Philipp Frank committed
125
        self._domain = makeDomain(domain)
126
127
        assert isinstance(self._domain[space], PowerSpace)
        logkl = _relative_log_k_lengths(self._domain[space])
128
        self._sc = logkl/float(logkl[-1])
Philipp Arras's avatar
Philipp Arras committed
129

130
        self._space = space
Philipp Haim's avatar
Philipp Haim committed
131
132
133
        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
134
        self._capability = self.TIMES | self.ADJOINT_TIMES
Philipp Arras's avatar
Philipp Arras committed
135

136
137
    def apply(self, x, mode):
        self._check_input(x, mode)
Philipp Frank's avatar
Philipp Frank committed
138
139
        x = x.to_global_data()
        if mode == self.TIMES:
Philipp Haim's avatar
Philipp Haim committed
140
            res = x - x[self._last]*self._sc[self._extender]
Philipp Frank's avatar
Philipp Frank committed
141
        else:
142
143
            res = x.copy()
            res[self._last] -= (x*self._sc[self._extender]).sum(
Martin Reinecke's avatar
Martin Reinecke committed
144
                    axis=self._space, keepdims=True)
145
        return from_global_data(self._tgt(mode), res)
Philipp Frank's avatar
Philipp Frank committed
146

Philipp Arras's avatar
Philipp Arras committed
147
148

class _TwoLogIntegrations(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
149
    def __init__(self, target, space=0):
Philipp Arras's avatar
Philipp Arras committed
150
        self._target = makeDomain(target)
151
152
153
154
155
        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
156
        self._log_vol = _log_vol(self._target[space])
Philipp Arras's avatar
Philipp Arras committed
157
158
159
160
        self._capability = self.TIMES | self.ADJOINT_TIMES

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

Martin Reinecke's avatar
Martin Reinecke committed
162
        # Maybe make class properties
163
164
        axis = self._target.axes[self._space][0]
        sl = (slice(None),)*axis
Philipp Haim's avatar
Fixes    
Philipp Haim committed
165
        extender_sl = (None,)*axis + (slice(None),) + (None,)*(self._target.axes[-1][-1] - axis)
166
167
        first = sl + (0,)
        second = sl + (1,)
Martin Reinecke's avatar
Martin Reinecke committed
168
169
170
        from_third = sl + (slice(2, None),)
        no_border = sl + (slice(1, -1),)
        reverse = sl + (slice(None, None, -1),)
171

Philipp Arras's avatar
Philipp Arras committed
172
        if mode == self.TIMES:
Philipp Haim's avatar
Philipp Haim committed
173
            x = x.to_global_data()
Philipp Arras's avatar
Philipp Arras committed
174
            res = np.empty(self._target.shape)
175
            res[first] = res[second] = 0
Martin Reinecke's avatar
Martin Reinecke committed
176
            res[from_third] = np.cumsum(x[second], axis=axis)
177
            res[from_third] = (res[from_third] + res[no_border])/2*self._log_vol[extender_sl] + x[first]
Martin Reinecke's avatar
Martin Reinecke committed
178
            res[from_third] = np.cumsum(res[from_third], axis=axis)
Philipp Arras's avatar
Philipp Arras committed
179
        else:
Philipp Haim's avatar
Philipp Haim committed
180
            x = x.to_global_data_rw()
Philipp Arras's avatar
Philipp Arras committed
181
            res = np.zeros(self._domain.shape)
Martin Reinecke's avatar
Martin Reinecke committed
182
            x[from_third] = np.cumsum(x[from_third][reverse], axis=axis)[reverse]
183
            res[first] += x[from_third]
184
            x[from_third] *= (self._log_vol/2.)[extender_sl]
185
            x[no_border] += x[from_third]
Martin Reinecke's avatar
Martin Reinecke committed
186
            res[second] += np.cumsum(x[from_third][reverse], axis=axis)[reverse]
187
        return from_global_data(self._tgt(mode), res)
Philipp Arras's avatar
Philipp Arras committed
188
189
190


class _Normalization(Operator):
Martin Reinecke's avatar
Martin Reinecke committed
191
    def __init__(self, domain, space=0):
Philipp Arras's avatar
Philipp Arras committed
192
        self._domain = self._target = makeDomain(domain)
193
        assert isinstance(self._domain[space], PowerSpace)
194
195
196
        hspace = list(self._domain)
        hspace[space] = hspace[space].harmonic_partner
        hspace = makeDomain(hspace)
Martin Reinecke's avatar
Martin Reinecke committed
197
        pd = PowerDistributor(hspace, power_space=self._domain[space], space=space)
198
        mode_multiplicity = pd.adjoint(full(pd.target, 1.)).to_global_data_rw()
199
        zero_mode = (slice(None),)*self._domain.axes[space][0] + (0,)
Philipp Haim's avatar
Philipp Haim committed
200
        mode_multiplicity[zero_mode] = 0
201
202
        self._mode_multiplicity = from_global_data(self._domain, mode_multiplicity)
        self._specsum = _SpecialSum(self._domain, space)
Philipp Arras's avatar
Philipp Arras committed
203
204
205
206
207
208
209

    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
210
        return self._specsum(self._mode_multiplicity*spec)**(-0.5)*amp
Philipp Arras's avatar
Philipp Arras committed
211
212
213


class _SpecialSum(EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
214
    def __init__(self, domain, space=0):
Philipp Arras's avatar
Philipp Arras committed
215
216
        self._domain = makeDomain(domain)
        self._capability = self.TIMES | self.ADJOINT_TIMES
217
        self._contractor = ContractionOperator(domain, space)
Philipp Arras's avatar
Philipp Arras committed
218
219
220

    def apply(self, x, mode):
        self._check_input(x, mode)
221
        return self._contractor.adjoint(self._contractor(x))
Philipp Arras's avatar
Philipp Arras committed
222
223


Philipp Haim's avatar
Philipp Haim committed
224
class _Distributor(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
225
    def __init__(self, dofdex, domain, target, space=0):
Philipp Haim's avatar
Philipp Haim committed
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
        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)
        x = x.to_global_data()
        if mode == self.TIMES:
            res = x[self._dofdex]
        else:
            res = np.empty(self._tgt(mode).shape)
            res[self._dofdex] = x
        return from_global_data(self._tgt(mode), res)
Martin Reinecke's avatar
Martin Reinecke committed
242

243

244
245
class _Amplitude(Operator):
    def __init__(self, target, fluctuations, flexibility, asperity,
246
                 loglogavgslope, azm, totvol, key, dofdex):
Philipp Arras's avatar
Philipp Arras committed
247
248
249
250
251
252
253
254
255
256
257
        """
        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
258
259
        if len(dofdex) > 0:
            N_copies = max(dofdex) + 1
Philipp Haim's avatar
Philipp Haim committed
260
            space = 1
Philipp Frank's avatar
cleanup    
Philipp Frank committed
261
262
            distributed_tgt = makeDomain((UnstructuredDomain(len(dofdex)),
                                          target))
Philipp Haim's avatar
Philipp Haim committed
263
264
265
            target = makeDomain((UnstructuredDomain(N_copies), target))
            Distributor = _Distributor(dofdex, target, distributed_tgt, 0)
        else:
Philipp Haim's avatar
Philipp Haim committed
266
            N_copies = 0
Philipp Haim's avatar
Philipp Haim committed
267
            space = 0
Philipp Haim's avatar
Philipp Haim committed
268
            distributed_tgt = target = makeDomain(target)
Martin Reinecke's avatar
Martin Reinecke committed
269
        azm_expander = ContractionOperator(distributed_tgt, spaces=space).adjoint
Philipp Haim's avatar
Philipp Haim committed
270
        assert isinstance(target[space], PowerSpace)
Martin Reinecke's avatar
Martin Reinecke committed
271

272
        twolog = _TwoLogIntegrations(target, space)
Philipp Arras's avatar
Philipp Arras committed
273
        dom = twolog.domain
274

275
        shp = dom[space].shape
Martin Reinecke's avatar
Martin Reinecke committed
276
277
        expander = ContractionOperator(dom, spaces=space).adjoint
        ps_expander = ContractionOperator(twolog.target, spaces=space).adjoint
Philipp Arras's avatar
Philipp Arras committed
278
279
280

        # Prepare constant fields
        foo = np.zeros(shp)
281
282
        foo[0] = foo[1] = np.sqrt(_log_vol(target[space]))
        vflex = DiagonalOperator(from_global_data(dom[space], foo), dom, space)
Philipp Arras's avatar
Philipp Arras committed
283
284
285

        foo = np.zeros(shp, dtype=np.float64)
        foo[0] += 1
286
        vasp = DiagonalOperator(from_global_data(dom[space], foo), dom, space)
Philipp Arras's avatar
Philipp Arras committed
287
288

        foo = np.ones(shp)
289
290
        foo[0] = _log_vol(target[space])**2/12.
        shift = DiagonalOperator(from_global_data(dom[space], foo), dom, space)
Martin Reinecke's avatar
Martin Reinecke committed
291

292
        vslope = DiagonalOperator(
Martin Reinecke's avatar
Martin Reinecke committed
293
294
295
            from_global_data(target[space],
                             _relative_log_k_lengths(target[space])),
            target, space)
296
297

        foo, bar = [np.zeros(target[space].shape) for _ in range(2)]
Philipp Arras's avatar
Philipp Arras committed
298
        bar[1:] = foo[0] = totvol
Martin Reinecke's avatar
Martin Reinecke committed
299
        vol0, vol1 = [DiagonalOperator(from_global_data(target[space], aa),
Philipp Frank's avatar
cleanup    
Philipp Frank committed
300
                                       target, space) for aa in (foo, bar)]
301

Martin Reinecke's avatar
Martin Reinecke committed
302
        # Prepare fields for Adder
303
        shift, vol0 = [op(full(op.domain, 1)) for op in (shift, vol0)]
Philipp Arras's avatar
Philipp Arras committed
304
305
        # End prepare constant fields

306
307
308
309
        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
310
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Arras's avatar
Philipp Arras committed
311
312

        xi = ducktape(dom, None, key)
Philipp Arras's avatar
Philipp Arras committed
313
        sigma = sig_flex*(Adder(shift) @ sig_asp).sqrt()
314
315
        smooth = _SlopeRemover(target, space) @ twolog @ (sigma*xi)
        op = _Normalization(target, space) @ (slope + smooth)
Philipp Haim's avatar
Philipp Haim committed
316
        if N_copies > 0:
Philipp Haim's avatar
Philipp Haim committed
317
318
            op = Distributor @ op
            sig_fluc = Distributor @ sig_fluc
Philipp Haim's avatar
Philipp Haim committed
319
            op = Adder(Distributor(vol0)) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Martin Reinecke's avatar
Martin Reinecke committed
320
            self._fluc = (_Distributor(dofdex, fluctuations.target, distributed_tgt[0]) @
Philipp Frank's avatar
Philipp Frank committed
321
                          fluctuations)
Philipp Haim's avatar
Philipp Haim committed
322
        else:
Philipp Frank's avatar
cleanup    
Philipp Frank committed
323
            op = Adder(vol0) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Philipp Frank's avatar
fixup    
Philipp Frank committed
324
            self._fluc = fluctuations
Philipp Arras's avatar
Philipp Arras committed
325

Philipp Arras's avatar
Philipp Arras committed
326
327
        self.apply = op.apply
        self._domain, self._target = op.domain, op.target
328
        self._space = space
Philipp Arras's avatar
Philipp Arras committed
329

Philipp Arras's avatar
Philipp Arras committed
330
331
332
333
    @property
    def fluctuation_amplitude(self):
        return self._fluc

334
335

class CorrelatedFieldMaker:
Philipp Haim's avatar
Philipp Haim committed
336
    def __init__(self, amplitude_offset, prefix, total_N):
Philipp Frank's avatar
fixup    
Philipp Frank committed
337
        assert isinstance(amplitude_offset, Operator)
338
        self._a = []
339
        self._position_spaces = []
Philipp Arras's avatar
Formats    
Philipp Arras committed
340

341
342
        self._azm = amplitude_offset
        self._prefix = prefix
Philipp Haim's avatar
Philipp Haim committed
343
        self._total_N = total_N
Philipp Arras's avatar
Formats    
Philipp Arras committed
344

345
    @staticmethod
Philipp Frank's avatar
Philipp Frank committed
346
    def make(offset_amplitude_mean, offset_amplitude_stddev, prefix,
Martin Reinecke's avatar
Martin Reinecke committed
347
348
             total_N=0,
             dofdex=None):
Philipp Frank's avatar
Philipp Frank committed
349
350
351
352
353
        if dofdex is None:
            dofdex = np.full(total_N, 0)
        else:
            assert len(dofdex) == total_N
        N = max(dofdex) + 1 if total_N > 0 else 0
354
355
        zm = _LognormalMomentMatching(offset_amplitude_mean,
                                      offset_amplitude_stddev,
Philipp Haim's avatar
Philipp Haim committed
356
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
357
                                      N)
Philipp Frank's avatar
fixup    
Philipp Frank committed
358
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
359
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
Philipp Haim's avatar
Philipp Haim committed
360
        return CorrelatedFieldMaker(zm, prefix, total_N)
361
362

    def add_fluctuations(self,
363
                         position_space,
364
365
366
367
368
369
370
371
                         fluctuations_mean,
                         fluctuations_stddev,
                         flexibility_mean,
                         flexibility_stddev,
                         asperity_mean,
                         asperity_stddev,
                         loglogavgslope_mean,
                         loglogavgslope_stddev,
Martin Reinecke's avatar
Martin Reinecke committed
372
373
374
375
                         prefix='',
                         index=None,
                         dofdex=None,
                         harmonic_partner=None):
Philipp Frank's avatar
Philipp Frank committed
376
377
        if harmonic_partner is None:
            harmonic_partner = position_space.get_default_codomain()
Philipp Frank's avatar
Fixup    
Philipp Frank committed
378
379
380
        else:
            position_space.check_codomain(harmonic_partner)
            harmonic_partner.check_codomain(position_space)
381

Philipp Haim's avatar
Philipp Haim committed
382
383
384
385
386
        if dofdex is None:
            dofdex = np.full(self._total_N, 0)
        else:
            assert len(dofdex) == self._total_N

Philipp Haim's avatar
Philipp Haim committed
387
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
388
            space = 1
Philipp Haim's avatar
Philipp Haim committed
389
390
            N = max(dofdex) + 1
            position_space = makeDomain((UnstructuredDomain(N), position_space))
Philipp Haim's avatar
Philipp Haim committed
391
392
        else:
            space = 0
Philipp Haim's avatar
Philipp Haim committed
393
            N = 0
Philipp Haim's avatar
Philipp Haim committed
394
            position_space = makeDomain(position_space)
Philipp Arras's avatar
Philipp Arras committed
395
        prefix = str(prefix)
Martin Reinecke's avatar
Martin Reinecke committed
396
        # assert isinstance(position_space[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
397

Philipp Arras's avatar
Philipp Arras committed
398
399
        fluct = _LognormalMomentMatching(fluctuations_mean,
                                         fluctuations_stddev,
400
                                         self._prefix + prefix + 'fluctuations',
Philipp Haim's avatar
Philipp Haim committed
401
                                         N)
Philipp Arras's avatar
Philipp Arras committed
402
        flex = _LognormalMomentMatching(flexibility_mean, flexibility_stddev,
403
                                        self._prefix + prefix + 'flexibility',
Philipp Haim's avatar
Philipp Haim committed
404
                                        N)
Philipp Arras's avatar
Philipp Arras committed
405
        asp = _LognormalMomentMatching(asperity_mean, asperity_stddev,
Martin Reinecke's avatar
Martin Reinecke committed
406
                                       self._prefix + prefix + 'asperity',
Philipp Haim's avatar
Philipp Haim committed
407
                                       N)
408
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
409
                        self._prefix + prefix + 'loglogavgslope', N)
Philipp Frank's avatar
Philipp Frank committed
410
        amp = _Amplitude(PowerSpace(harmonic_partner),
Martin Reinecke's avatar
Martin Reinecke committed
411
                         fluct, flex, asp, avgsl, self._azm,
Philipp Frank's avatar
fixup    
Philipp Frank committed
412
                         position_space[-1].total_volume,
413
                         self._prefix + prefix + 'spectrum', dofdex)
Philipp Haim's avatar
Philipp Haim committed
414

415
416
        if index is not None:
            self._a.insert(index, amp)
417
            self._position_spaces.insert(index, position_space)
418
419
        else:
            self._a.append(amp)
420
            self._position_spaces.append(position_space)
421

Philipp Frank's avatar
fixup    
Philipp Frank committed
422
    def _finalize_from_op(self):
Philipp Haim's avatar
Philipp Haim committed
423
        n_amplitudes = len(self._a)
Philipp Haim's avatar
Philipp Haim committed
424
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
425
            hspace = makeDomain([UnstructuredDomain(self._total_N)] +
Martin Reinecke's avatar
Martin Reinecke committed
426
427
                                [dd.target[-1].harmonic_partner
                                    for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
428
429
            spaces = tuple(range(1, n_amplitudes + 1))
            amp_space = 1
Philipp Haim's avatar
Philipp Haim committed
430
431
        else:
            hspace = makeDomain(
432
                     [dd.target[0].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
433
            spaces = tuple(range(n_amplitudes))
Philipp Haim's avatar
Philipp Haim committed
434
            amp_space = 0
435

Martin Reinecke's avatar
Martin Reinecke committed
436
        expander = ContractionOperator(hspace, spaces=spaces).adjoint
Philipp Frank's avatar
fixup    
Philipp Frank committed
437
        azm = expander @ self._azm
438

439
        ht = HarmonicTransformOperator(hspace,
Philipp Haim's avatar
Philipp Haim committed
440
                                       self._position_spaces[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
441
                                       space=spaces[0])
442
        for i in range(1, n_amplitudes):
443
            ht = (HarmonicTransformOperator(ht.target,
Philipp Haim's avatar
Philipp Haim committed
444
                                            self._position_spaces[i][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
445
                                            space=spaces[i]) @ ht)
446

Philipp Haim's avatar
Philipp Haim committed
447
        pd = PowerDistributor(hspace, self._a[0].target[amp_space], amp_space)
448
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
449
            pd = (pd @ PowerDistributor(pd.domain,
Philipp Haim's avatar
Philipp Haim committed
450
                                        self._a[i].target[amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
451
                                        space=spaces[i]))
Philipp Arras's avatar
Philipp Arras committed
452

453
454
        a = ContractionOperator(pd.domain, spaces[1:]).adjoint @ self._a[0]
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
455
            co = ContractionOperator(pd.domain,
Martin Reinecke's avatar
Martin Reinecke committed
456
                                     spaces[:i] + spaces[i+1:])
457
            a = a*(co.adjoint @ self._a[i])
Philipp Arras's avatar
Philipp Arras committed
458

Philipp Frank's avatar
fixup    
Philipp Frank committed
459
        return ht(azm*(pd @ a)*ducktape(hspace, None, self._prefix + 'xi'))
Philipp Arras's avatar
Philipp Arras committed
460

Philipp Arras's avatar
Formats    
Philipp Arras committed
461
    def finalize(self, offset=None, prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
462
463
464
        """
        offset vs zeromode: volume factor
        """
465
        op = self._finalize_from_op()
Philipp Arras's avatar
Philipp Arras committed
466
        if offset is not None:
467
468
469
470
471
472
473
            # 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
474
        self.statistics_summary(prior_info)
475
476
        return op

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

        if prior_info == 0:
            return

483
484
        lst = [('Offset amplitude', self.amplitude_total_offset),
               ('Total fluctuation amplitude', self.total_fluctuation)]
485
        namps = len(self._a)
486
487
488
489
490
491
492
493
        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:
494
495
496
497
498
            sc = StatCalculator()
            for _ in range(prior_info):
                sc.add(op(from_random('normal', op.domain)))
            mean = sc.mean.to_global_data()
            stddev = sc.var.sqrt().to_global_data()
499
500
            for m, s in zip(mean.flatten(), stddev.flatten()):
                print('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
501
502
503
504

    def moment_slice_to_average(self, fluctuations_slice_mean, nsamples=1000):
        fluctuations_slice_mean = float(fluctuations_slice_mean)
        assert fluctuations_slice_mean > 0
505
        from ..sugar import from_random
506
507
        scm = 1.
        for a in self._a:
Philipp Haim's avatar
Philipp Haim committed
508
            op = a.fluctuation_amplitude*self._azm.one_over()
Martin Reinecke's avatar
Martin Reinecke committed
509
            res = np.array([op(from_random('normal', op.domain)).to_global_data()
510
511
                            for _ in range(nsamples)])
            scm *= res**2 + 1.
512
        return fluctuations_slice_mean/np.mean(np.sqrt(scm))
513

Philipp Arras's avatar
Philipp Arras committed
514
    @property
Philipp Haim's avatar
Philipp Haim committed
515
    def normalized_amplitudes(self):
516
        return self._a
Philipp Arras's avatar
Philipp Arras committed
517

Philipp Haim's avatar
Philipp Haim committed
518
519
520
521
522
523
524
    @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
525
526
        dom = self._a[0].target
        expand = ContractionOperator(dom, len(dom)-1).adjoint
Philipp Haim's avatar
Philipp Haim committed
527
528
        return self._a[0]*(expand @ self.amplitude_total_offset)

529
530
531
    @property
    def amplitude_total_offset(self):
        return self._azm
Philipp Arras's avatar
Philipp Arras committed
532
533

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

Philipp Arras's avatar
Philipp Arras committed
546
    def slice_fluctuation(self, space):
547
        """Returns operator which acts on prior or posterior samples"""
548
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
549
            raise NotImplementedError
550
551
        assert space < len(self._a)
        if len(self._a) == 1:
552
            return self.average_fluctuation(0)
553
554
        q = 1.
        for j in range(len(self._a)):
Philipp Haim's avatar
Philipp Haim committed
555
            fl = self._a[j].fluctuation_amplitude*self._azm.one_over()
556
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
557
                q = q*fl**2
558
            else:
Philipp Arras's avatar
Philipp Arras committed
559
                q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Philipp Arras's avatar
Formats    
Philipp Arras committed
560
        return q.sqrt()*self._azm
Philipp Arras's avatar
Philipp Arras committed
561
562

    def average_fluctuation(self, space):
563
        """Returns operator which acts on prior or posterior samples"""
564
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
565
            raise NotImplementedError
566
567
        assert space < len(self._a)
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
568
569
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
570

571
572
    @staticmethod
    def offset_amplitude_realized(samples):
573
574
        res = 0.
        for s in samples:
Philipp Frank's avatar
fixes    
Philipp Frank committed
575
            res = res + s.mean()**2
576
        return np.sqrt(res/len(samples))
Philipp Arras's avatar
Philipp Arras committed
577

578
579
580
581
582
583
584
585
    @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."""
586
587
588
        ldom = len(samples[0].domain)
        assert space < ldom
        if ldom == 1:
589
            return _total_fluctuation_realized(samples)
Philipp Arras's avatar
Philipp Arras committed
590
        res1, res2 = 0., 0.
591
        for s in samples:
Philipp Frank's avatar
fixes    
Philipp Frank committed
592
593
594
595
596
597
598
            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
599
    @staticmethod
600
601
602
603
604
605
606
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)
        assert space < ldom
        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
611
612
        res = 0.
        for s in samples:
613
614
615
616
            r = s.mean(spaces)
            res = res + (r - r.mean())**2
        res = res/len(samples)
        return np.sqrt(res.mean())