correlated_fields.py 23.5 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
25
from ..field import Field
from ..multi_field import MultiField
Philipp Arras's avatar
Philipp Arras committed
26
from ..operators.adder import Adder
27
from ..operators.contraction_operator import ContractionOperator
Philipp Arras's avatar
Philipp Arras committed
28
from ..operators.diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
29
from ..operators.distributors import PowerDistributor
Philipp Arras's avatar
Philipp Arras committed
30
from ..operators.endomorphic_operator import EndomorphicOperator
Martin Reinecke's avatar
Martin Reinecke committed
31
from ..operators.harmonic_operators import HarmonicTransformOperator
Philipp Arras's avatar
Philipp Arras committed
32
from ..operators.linear_operator import LinearOperator
Philipp Arras's avatar
Philipp Arras committed
33
from ..operators.operator import Operator
Philipp Arras's avatar
Philipp Arras committed
34
from ..operators.simple_linear_operators import ducktape
35
from ..probing import StatCalculator
Martin Reinecke's avatar
Martin Reinecke committed
36
from ..sugar import makeField, full, makeDomain
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))
54
55
56
57
58
    if not np.all(mean > 0):
        raise ValueError(f"mean must be greater 0; got {mean!r}")
    if not np.all(sig > 0):
        raise ValueError(f"sig must be greater 0; got {sig!r}")

Philipp Arras's avatar
Philipp Arras committed
59
60
    logsig = np.sqrt(np.log((sig/mean)**2 + 1))
    logmean = np.log(mean) - logsig**2/2
61
    return logmean, logsig
Philipp Arras's avatar
Philipp Arras committed
62
63


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


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


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


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


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


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

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

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

Philipp Arras's avatar
Philipp Arras committed
149
150

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

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

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

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


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

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


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

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


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

248

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

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

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

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

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

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

297
        vslope = DiagonalOperator(
Martin Reinecke's avatar
Martin Reinecke committed
298
            makeField(target[space],
Martin Reinecke's avatar
Martin Reinecke committed
299
300
                             _relative_log_k_lengths(target[space])),
            target, space)
301
302

        foo, bar = [np.zeros(target[space].shape) for _ in range(2)]
Philipp Arras's avatar
Philipp Arras committed
303
        bar[1:] = foo[0] = totvol
Martin Reinecke's avatar
Martin Reinecke committed
304
        vol0, vol1 = [DiagonalOperator(makeField(target[space], aa),
Philipp Frank's avatar
cleanup    
Philipp Frank committed
305
                                       target, space) for aa in (foo, bar)]
306

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

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

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

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

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

339
340

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        if prior_info == 0:
            return

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

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

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

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

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

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

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

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

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

586
587
588
589
590
591
592
593
    @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."""
594
        ldom = len(samples[0].domain)
595
596
        if space >= ldom:
            raise ValueError(f"invalid space specified; got {space!r}")
597
        if ldom == 1:
598
            return _total_fluctuation_realized(samples)
Philipp Arras's avatar
Philipp Arras committed
599
        res1, res2 = 0., 0.
600
        for s in samples:
Philipp Frank's avatar
fixes    
Philipp Frank committed
601
602
603
604
605
606
607
            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
608
    @staticmethod
609
610
611
612
    def average_fluctuation_realized(samples, space):
        """Computes average fluctuations from collection of field (defined in signal
        space) realizations."""
        ldom = len(samples[0].domain)
613
614
        if space >= ldom:
            raise ValueError(f"invalid space specified; got {space!r}")
615
616
617
618
619
620
        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
621
622
        res = 0.
        for s in samples:
623
624
625
626
            r = s.mean(spaces)
            res = res + (r - r.mean())**2
        res = res/len(samples)
        return np.sqrt(res.mean())