correlated_fields.py 23.1 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
104
105
106
107
108
109
110


def _stats(op, samples):
    sc = StatCalculator()
    for s in samples:
        sc.add(op(s.extract(op.domain)))
    return sc.mean.to_global_data(), sc.var.sqrt().to_global_data()


Philipp Arras's avatar
Philipp Arras committed
111
class _LognormalMomentMatching(Operator):
Philipp Haim's avatar
Philipp Haim committed
112
    def __init__(self, mean, sig, key, N_copies):
Philipp Arras's avatar
Philipp Arras committed
113
        key = str(key)
Philipp Haim's avatar
Philipp Haim committed
114
        logmean, logsig = _lognormal_moments(mean, sig, N_copies)
Philipp Arras's avatar
Philipp Arras committed
115
116
        self._mean = mean
        self._sig = sig
Philipp Haim's avatar
Philipp Haim committed
117
        op = _normal(logmean, logsig, key, N_copies).exp()
Philipp Arras's avatar
Philipp Arras committed
118
119
120
121
122
123
124
125
126
127
        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
128
129


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

137
        self._space = space
Philipp Haim's avatar
Philipp Haim committed
138
139
140
        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
141
        self._capability = self.TIMES | self.ADJOINT_TIMES
Philipp Arras's avatar
Philipp Arras committed
142

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

Philipp Arras's avatar
Philipp Arras committed
154
155

class _TwoLogIntegrations(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
156
    def __init__(self, target, space=0):
Philipp Arras's avatar
Philipp Arras committed
157
        self._target = makeDomain(target)
158
159
160
161
162
        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
163
        self._log_vol = _log_vol(self._target[space])
Philipp Arras's avatar
Philipp Arras committed
164
165
166
167
        self._capability = self.TIMES | self.ADJOINT_TIMES

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

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

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


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

    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
217
        return self._specsum(self._mode_multiplicity*spec)**(-0.5)*amp
Philipp Arras's avatar
Philipp Arras committed
218
219
220


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

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


Philipp Haim's avatar
Philipp Haim committed
231
class _Distributor(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
232
    def __init__(self, dofdex, domain, target, space=0):
Philipp Haim's avatar
Philipp Haim committed
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
        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
249

250

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

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

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

        # Prepare constant fields
        foo = np.zeros(shp)
288
289
        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
290
291
292

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

        foo = np.ones(shp)
296
297
        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
298

299
        vslope = DiagonalOperator(
Martin Reinecke's avatar
Martin Reinecke committed
300
301
302
            from_global_data(target[space],
                             _relative_log_k_lengths(target[space])),
            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
Martin Reinecke's avatar
Martin Reinecke committed
306
        vol0, vol1 = [DiagonalOperator(from_global_data(target[space], aa),
Philipp Frank's avatar
cleanup    
Philipp Frank committed
307
                                       target, space) for aa in (foo, bar)]
308

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

313
314
315
316
        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
317
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Arras's avatar
Philipp Arras committed
318
319

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

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

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

341
342

class CorrelatedFieldMaker:
Philipp Haim's avatar
Philipp Haim committed
343
    def __init__(self, amplitude_offset, prefix, total_N):
Philipp Frank's avatar
fixup    
Philipp Frank committed
344
        assert isinstance(amplitude_offset, Operator)
345
        self._a = []
346
        self._position_spaces = []
Philipp Arras's avatar
Formats    
Philipp Arras committed
347

348
349
        self._azm = amplitude_offset
        self._prefix = prefix
Philipp Haim's avatar
Philipp Haim committed
350
        self._total_N = total_N
Philipp Arras's avatar
Formats    
Philipp Arras committed
351

352
    @staticmethod
Philipp Frank's avatar
Philipp Frank committed
353
    def make(offset_amplitude_mean, offset_amplitude_stddev, prefix,
Martin Reinecke's avatar
Martin Reinecke committed
354
355
             total_N=0,
             dofdex=None):
Philipp Frank's avatar
Philipp Frank committed
356
357
358
359
360
        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
361
362
        zm = _LognormalMomentMatching(offset_amplitude_mean,
                                      offset_amplitude_stddev,
Philipp Haim's avatar
Philipp Haim committed
363
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
364
                                      N)
Philipp Frank's avatar
fixup    
Philipp Frank committed
365
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
366
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
Philipp Haim's avatar
Philipp Haim committed
367
        return CorrelatedFieldMaker(zm, prefix, total_N)
368
369

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

Philipp Haim's avatar
Philipp Haim committed
389
390
391
392
393
        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
394
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
395
            space = 1
Philipp Haim's avatar
Philipp Haim committed
396
397
            N = max(dofdex) + 1
            position_space = makeDomain((UnstructuredDomain(N), position_space))
Philipp Haim's avatar
Philipp Haim committed
398
399
        else:
            space = 0
Philipp Haim's avatar
Philipp Haim committed
400
            N = 0
Philipp Haim's avatar
Philipp Haim committed
401
            position_space = makeDomain(position_space)
Philipp Arras's avatar
Philipp Arras committed
402
        prefix = str(prefix)
Martin Reinecke's avatar
Martin Reinecke committed
403
        # assert isinstance(position_space[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
404

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

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

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

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

446
        ht = HarmonicTransformOperator(hspace,
Philipp Haim's avatar
Philipp Haim committed
447
                                       self._position_spaces[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
448
                                       space=spaces[0])
449
        for i in range(1, n_amplitudes):
450
            ht = (HarmonicTransformOperator(ht.target,
Philipp Haim's avatar
Philipp Haim committed
451
                                            self._position_spaces[i][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
452
                                            space=spaces[i]) @ ht)
453

Philipp Haim's avatar
Philipp Haim committed
454
        pd = PowerDistributor(hspace, self._a[0].target[amp_space], amp_space)
455
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
456
            pd = (pd @ PowerDistributor(pd.domain,
Philipp Haim's avatar
Philipp Haim committed
457
                                        self._a[i].target[amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
458
                                        space=spaces[i]))
Philipp Arras's avatar
Philipp Arras committed
459

460
461
        a = ContractionOperator(pd.domain, spaces[1:]).adjoint @ self._a[0]
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
462
            co = ContractionOperator(pd.domain,
Martin Reinecke's avatar
Martin Reinecke committed
463
                                     spaces[:i] + spaces[i+1:])
464
            a = a*(co.adjoint @ self._a[i])
Philipp Arras's avatar
Philipp Arras committed
465

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

Philipp Arras's avatar
Formats    
Philipp Arras committed
468
    def finalize(self, offset=None, prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
469
470
471
        """
        offset vs zeromode: volume factor
        """
472
        op = self._finalize_from_op()
Philipp Arras's avatar
Philipp Arras committed
473
        if offset is not None:
474
475
476
477
478
479
480
            # 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
481

482
483
484
485
486
487
488
489
490
491
492
493
        if prior_info > 0:
            from ..sugar import from_random
            samps = [
                from_random('normal', op.domain) for _ in range(prior_info)
            ]
            self.statistics_summary(samps)
        return op

    def statistics_summary(self, samples):
        lst = [('Offset amplitude', self.amplitude_total_offset),
               ('Total fluctuation amplitude', self.total_fluctuation)]

494
        namps = len(self._a)
495
496
497
498
499
500
501
502
503
        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:
            mean, stddev = _stats(op, samples)
504
505
            for m, s in zip(mean.flatten(), stddev.flatten()):
                print('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
506
507
508
509

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

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

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

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

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

Philipp Arras's avatar
Philipp Arras committed
551
    def slice_fluctuation(self, space):
552
        """Returns operator which acts on prior or posterior samples"""
553
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
554
            raise NotImplementedError
555
556
        assert space < len(self._a)
        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
572
        assert space < len(self._a)
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
573
574
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
575

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

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