correlated_fields.py 23.4 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
Martin Reinecke's avatar
Martin Reinecke committed
209
        self._mode_multiplicity = makeField(self._domain,
Philipp Arras's avatar
Philipp Arras committed
210
                                                   mode_multiplicity)
211
        self._specsum = _SpecialSum(self._domain, space)
Philipp Arras's avatar
Philipp Arras committed
212 213 214 215 216 217 218

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


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

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


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

252

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

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

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

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

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

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

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

        foo, bar = [np.zeros(target[space].shape) for _ in range(2)]
Philipp Arras's avatar
Philipp Arras committed
307
        bar[1:] = foo[0] = totvol
Martin Reinecke's avatar
Martin Reinecke committed
308
        vol0, vol1 = [DiagonalOperator(makeField(target[space], aa),
Philipp Frank's avatar
cleanup  
Philipp Frank committed
309
                                       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,
Martin Reinecke's avatar
Martin Reinecke committed
416
                                       self._prefix + prefix + 'asperity',
Philipp Haim's avatar
Philipp Haim committed
417
                                       N)
418
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
419
                        self._prefix + prefix + 'loglogavgslope', N)
Philipp Frank's avatar
Philipp Frank committed
420
        amp = _Amplitude(PowerSpace(harmonic_partner),
Martin Reinecke's avatar
Martin Reinecke committed
421
                         fluct, flex, asp, avgsl, self._azm,
Philipp Frank's avatar
fixup  
Philipp Frank committed
422
                         position_space[-1].total_volume,
423
                         self._prefix + prefix + 'spectrum', dofdex)
Philipp Haim's avatar
Philipp Haim committed
424

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

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

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

449
        ht = HarmonicTransformOperator(hspace,
Philipp Haim's avatar
Philipp Haim committed
450
                                       self._position_spaces[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
451
                                       space=spaces[0])
452
        for i in range(1, n_amplitudes):
453 454 455 456 457 458 459 460 461 462 463
            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
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
            for m, s in zip(mean.flatten(), stddev.flatten()):
504
                logger.info('{}: {:.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
        if space >= len(self._a):
557
            raise ValueError("invalid space specified; got {!r}".format(space))
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
        if space >= len(self._a):
574
            raise ValueError("invalid space specified; got {!r}".format(space))
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
        if space >= ldom:
596
            raise ValueError("invalid space specified; got {!r}".format(space))
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
        if space >= ldom:
614
            raise ValueError("invalid space specified; got {!r}".format(space))
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())