There is a maintenance of MPCDF Gitlab on Thursday, April 22st 2020, 9:00 am CEST - Expect some service interruptions during this time

correlated_fields.py 23.3 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1 2 3 4 5 6 7 8 9 10 11 12 13
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
14
# Copyright(C) 2013-2019 Max-Planck-Society
Philipp Arras's avatar
Philipp Arras committed
15
# Authors: Philipp Frank, Philipp Arras, Philipp Haim
Martin Reinecke's avatar
Martin Reinecke committed
16
#
17
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
Philipp Arras's avatar
Philipp Arras committed
18

19 20 21
from functools import reduce
from operator import mul

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

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

42

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

52

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

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


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


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


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


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


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


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

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

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


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

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

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

Philipp Arras's avatar
Philipp Arras committed
153 154

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

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

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

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


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

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


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

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


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

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

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

251

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

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

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

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

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

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

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

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

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

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

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

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

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

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

343 344

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

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

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

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

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

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

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

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

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

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

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

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

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

        if prior_info == 0:
            return

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

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

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

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

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

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

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

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

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

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

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

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

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