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 24.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 104 105 106 107 108
def _structured_spaces(domain):
    if isinstance(domain[0], UnstructuredDomain):
        return np.arange(1, len(domain))
    return np.arange(len(domain))


Philipp Haim's avatar
Philipp Haim committed
109
def _total_fluctuation_realized(samples):
Philipp Haim's avatar
Philipp Haim committed
110 111 112
    spaces = _structured_spaces(samples[0].domain)
    co = ContractionOperator(samples[0].domain, spaces)
    size = co.domain.size/co.target.size
113 114
    res = 0.
    for s in samples:
Philipp Haim's avatar
Philipp Haim committed
115 116 117 118 119
        res = res + (s - co.adjoint(co(s)/size))**2
    res = res.mean(spaces)/len(samples)
    if np.isscalar(res):
        return np.sqrt(res)
    return np.sqrt(res.val)
120 121


Philipp Arras's avatar
Philipp Arras committed
122
class _LognormalMomentMatching(Operator):
Philipp Haim's avatar
Philipp Haim committed
123
    def __init__(self, mean, sig, key, N_copies):
Philipp Arras's avatar
Philipp Arras committed
124
        key = str(key)
Philipp Haim's avatar
Philipp Haim committed
125
        logmean, logsig = _lognormal_moments(mean, sig, N_copies)
Philipp Arras's avatar
Philipp Arras committed
126 127
        self._mean = mean
        self._sig = sig
Philipp Haim's avatar
Philipp Haim committed
128
        op = _normal(logmean, logsig, key, N_copies).exp()
Philipp Arras's avatar
Philipp Arras committed
129 130 131 132 133 134 135 136 137 138
        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
139 140


Philipp Frank's avatar
Philipp Frank committed
141
class _SlopeRemover(EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
142
    def __init__(self, domain, space=0):
Philipp Frank's avatar
Philipp Frank committed
143
        self._domain = makeDomain(domain)
144 145
        assert isinstance(self._domain[space], PowerSpace)
        logkl = _relative_log_k_lengths(self._domain[space])
146
        self._sc = logkl/float(logkl[-1])
Philipp Arras's avatar
Philipp Arras committed
147

148
        self._space = space
Philipp Haim's avatar
Philipp Haim committed
149 150 151
        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
152
        self._capability = self.TIMES | self.ADJOINT_TIMES
Philipp Arras's avatar
Philipp Arras committed
153

154 155
    def apply(self, x, mode):
        self._check_input(x, mode)
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
156
        x = x.val
Philipp Frank's avatar
Philipp Frank committed
157
        if mode == self.TIMES:
Philipp Haim's avatar
Philipp Haim committed
158
            res = x - x[self._last]*self._sc[self._extender]
Philipp Frank's avatar
Philipp Frank committed
159
        else:
160 161
            res = x.copy()
            res[self._last] -= (x*self._sc[self._extender]).sum(
Philipp Arras's avatar
Philipp Arras committed
162
                axis=self._space, keepdims=True)
Martin Reinecke's avatar
Martin Reinecke committed
163
        return makeField(self._tgt(mode), res)
Philipp Frank's avatar
Philipp Frank committed
164

Philipp Arras's avatar
Philipp Arras committed
165 166

class _TwoLogIntegrations(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
167
    def __init__(self, target, space=0):
Philipp Arras's avatar
Philipp Arras committed
168
        self._target = makeDomain(target)
169 170 171 172 173
        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
174
        self._log_vol = _log_vol(self._target[space])
Philipp Arras's avatar
Philipp Arras committed
175 176 177 178
        self._capability = self.TIMES | self.ADJOINT_TIMES

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

Martin Reinecke's avatar
Martin Reinecke committed
180
        # Maybe make class properties
181 182
        axis = self._target.axes[self._space][0]
        sl = (slice(None),)*axis
Philipp Haim's avatar
Fixes  
Philipp Haim committed
183
        extender_sl = (None,)*axis + (slice(None),) + (None,)*(self._target.axes[-1][-1] - axis)
184 185
        first = sl + (0,)
        second = sl + (1,)
Martin Reinecke's avatar
Martin Reinecke committed
186 187 188
        from_third = sl + (slice(2, None),)
        no_border = sl + (slice(1, -1),)
        reverse = sl + (slice(None, None, -1),)
189

Philipp Arras's avatar
Philipp Arras committed
190
        if mode == self.TIMES:
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
191
            x = x.val
Philipp Arras's avatar
Philipp Arras committed
192
            res = np.empty(self._target.shape)
193
            res[first] = res[second] = 0
Martin Reinecke's avatar
Martin Reinecke committed
194
            res[from_third] = np.cumsum(x[second], axis=axis)
195
            res[from_third] = (res[from_third] + res[no_border])/2*self._log_vol[extender_sl] + x[first]
Martin Reinecke's avatar
Martin Reinecke committed
196
            res[from_third] = np.cumsum(res[from_third], axis=axis)
Philipp Arras's avatar
Philipp Arras committed
197
        else:
Martin Reinecke's avatar
Martin Reinecke committed
198
            x = x.val_rw()
Philipp Arras's avatar
Philipp Arras committed
199
            res = np.zeros(self._domain.shape)
Martin Reinecke's avatar
Martin Reinecke committed
200
            x[from_third] = np.cumsum(x[from_third][reverse], axis=axis)[reverse]
201
            res[first] += x[from_third]
202
            x[from_third] *= (self._log_vol/2.)[extender_sl]
203
            x[no_border] += x[from_third]
Martin Reinecke's avatar
Martin Reinecke committed
204
            res[second] += np.cumsum(x[from_third][reverse], axis=axis)[reverse]
Martin Reinecke's avatar
Martin Reinecke committed
205
        return makeField(self._tgt(mode), res)
Philipp Arras's avatar
Philipp Arras committed
206 207 208


class _Normalization(Operator):
Martin Reinecke's avatar
Martin Reinecke committed
209
    def __init__(self, domain, space=0):
Philipp Arras's avatar
Philipp Arras committed
210
        self._domain = self._target = makeDomain(domain)
211
        assert isinstance(self._domain[space], PowerSpace)
212 213 214
        hspace = list(self._domain)
        hspace[space] = hspace[space].harmonic_partner
        hspace = makeDomain(hspace)
Philipp Arras's avatar
Philipp Arras committed
215 216 217
        pd = PowerDistributor(hspace,
                              power_space=self._domain[space],
                              space=space)
Martin Reinecke's avatar
Martin Reinecke committed
218
        mode_multiplicity = pd.adjoint(full(pd.target, 1.)).val_rw()
219
        zero_mode = (slice(None),)*self._domain.axes[space][0] + (0,)
Philipp Haim's avatar
Philipp Haim committed
220
        mode_multiplicity[zero_mode] = 0
Philipp Arras's avatar
Philipp Arras committed
221
        self._mode_multiplicity = makeField(self._domain, mode_multiplicity)
222
        self._specsum = _SpecialSum(self._domain, space)
Philipp Arras's avatar
Philipp Arras committed
223 224 225 226 227 228 229

    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
230
        return self._specsum(self._mode_multiplicity*spec)**(-0.5)*amp
Philipp Arras's avatar
Philipp Arras committed
231 232 233


class _SpecialSum(EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
234
    def __init__(self, domain, space=0):
Philipp Arras's avatar
Philipp Arras committed
235 236
        self._domain = makeDomain(domain)
        self._capability = self.TIMES | self.ADJOINT_TIMES
237
        self._contractor = ContractionOperator(domain, space)
Philipp Arras's avatar
Philipp Arras committed
238 239 240

    def apply(self, x, mode):
        self._check_input(x, mode)
241
        return self._contractor.adjoint(self._contractor(x))
Philipp Arras's avatar
Philipp Arras committed
242 243


Philipp Haim's avatar
Philipp Haim committed
244
class _Distributor(LinearOperator):
Martin Reinecke's avatar
Martin Reinecke committed
245
    def __init__(self, dofdex, domain, target, space=0):
Philipp Haim's avatar
Philipp Haim committed
246 247 248 249 250 251 252 253 254
        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
255
        x = x.val
Philipp Haim's avatar
Philipp Haim committed
256 257 258 259 260
        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
261
        return makeField(self._tgt(mode), res)
Martin Reinecke's avatar
Martin Reinecke committed
262

263

264 265
class _Amplitude(Operator):
    def __init__(self, target, fluctuations, flexibility, asperity,
266
                 loglogavgslope, azm, totvol, key, dofdex):
Philipp Arras's avatar
Philipp Arras committed
267 268 269 270 271 272 273 274 275 276 277
        """
        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
278 279
        if len(dofdex) > 0:
            N_copies = max(dofdex) + 1
Philipp Haim's avatar
Philipp Haim committed
280
            space = 1
Philipp Frank's avatar
cleanup  
Philipp Frank committed
281 282
            distributed_tgt = makeDomain((UnstructuredDomain(len(dofdex)),
                                          target))
Philipp Haim's avatar
Philipp Haim committed
283 284 285
            target = makeDomain((UnstructuredDomain(N_copies), target))
            Distributor = _Distributor(dofdex, target, distributed_tgt, 0)
        else:
Philipp Haim's avatar
Philipp Haim committed
286
            N_copies = 0
Philipp Haim's avatar
Philipp Haim committed
287
            space = 0
Philipp Haim's avatar
Philipp Haim committed
288
            distributed_tgt = target = makeDomain(target)
Martin Reinecke's avatar
Martin Reinecke committed
289
        azm_expander = ContractionOperator(distributed_tgt, spaces=space).adjoint
Philipp Haim's avatar
Philipp Haim committed
290
        assert isinstance(target[space], PowerSpace)
Martin Reinecke's avatar
Martin Reinecke committed
291

292
        twolog = _TwoLogIntegrations(target, space)
Philipp Arras's avatar
Philipp Arras committed
293
        dom = twolog.domain
294

295
        shp = dom[space].shape
Martin Reinecke's avatar
Martin Reinecke committed
296 297
        expander = ContractionOperator(dom, spaces=space).adjoint
        ps_expander = ContractionOperator(twolog.target, spaces=space).adjoint
Philipp Arras's avatar
Philipp Arras committed
298 299 300

        # Prepare constant fields
        foo = np.zeros(shp)
301
        foo[0] = foo[1] = np.sqrt(_log_vol(target[space]))
Martin Reinecke's avatar
Martin Reinecke committed
302
        vflex = DiagonalOperator(makeField(dom[space], foo), dom, space)
Philipp Arras's avatar
Philipp Arras committed
303 304 305

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

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

312
        vslope = DiagonalOperator(
Philipp Arras's avatar
Philipp Arras committed
313
            makeField(target[space], _relative_log_k_lengths(target[space])),
Martin Reinecke's avatar
Martin Reinecke committed
314
            target, space)
315 316

        foo, bar = [np.zeros(target[space].shape) for _ in range(2)]
Philipp Arras's avatar
Philipp Arras committed
317
        bar[1:] = foo[0] = totvol
Philipp Arras's avatar
Philipp Arras committed
318 319 320 321
        vol0, vol1 = [
            DiagonalOperator(makeField(target[space], aa), target, space)
            for aa in (foo, bar)
        ]
322

Martin Reinecke's avatar
Martin Reinecke committed
323
        # Prepare fields for Adder
324
        shift, vol0 = [op(full(op.domain, 1)) for op in (shift, vol0)]
Philipp Arras's avatar
Philipp Arras committed
325 326
        # End prepare constant fields

327 328 329 330
        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
331
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Arras's avatar
Philipp Arras committed
332 333

        xi = ducktape(dom, None, key)
Philipp Arras's avatar
Philipp Arras committed
334
        sigma = sig_flex*(Adder(shift) @ sig_asp).sqrt()
335 336
        smooth = _SlopeRemover(target, space) @ twolog @ (sigma*xi)
        op = _Normalization(target, space) @ (slope + smooth)
Philipp Haim's avatar
Philipp Haim committed
337
        if N_copies > 0:
Philipp Haim's avatar
Philipp Haim committed
338 339
            op = Distributor @ op
            sig_fluc = Distributor @ sig_fluc
Philipp Haim's avatar
Philipp Haim committed
340
            op = Adder(Distributor(vol0)) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Philipp Arras's avatar
Philipp Arras committed
341 342
            self._fluc = (_Distributor(dofdex, fluctuations.target,
                                       distributed_tgt[0]) @ fluctuations)
Philipp Haim's avatar
Philipp Haim committed
343
        else:
Philipp Frank's avatar
cleanup  
Philipp Frank committed
344
            op = Adder(vol0) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Philipp Frank's avatar
fixup  
Philipp Frank committed
345
            self._fluc = fluctuations
Philipp Arras's avatar
Philipp Arras committed
346

Philipp Arras's avatar
Philipp Arras committed
347 348
        self.apply = op.apply
        self._domain, self._target = op.domain, op.target
349
        self._space = space
Philipp Arras's avatar
Philipp Arras committed
350

Philipp Arras's avatar
Philipp Arras committed
351 352 353 354
    @property
    def fluctuation_amplitude(self):
        return self._fluc

355 356

class CorrelatedFieldMaker:
Philipp Haim's avatar
Philipp Haim committed
357
    def __init__(self, amplitude_offset, prefix, total_N):
358 359
        if not isinstance(amplitude_offset, Operator):
            raise TypeError("amplitude_offset needs to be an operator")
360
        self._a = []
361
        self._position_spaces = []
Philipp Arras's avatar
Formats  
Philipp Arras committed
362

363 364
        self._azm = amplitude_offset
        self._prefix = prefix
Philipp Haim's avatar
Philipp Haim committed
365
        self._total_N = total_N
Philipp Arras's avatar
Formats  
Philipp Arras committed
366

367
    @staticmethod
Philipp Frank's avatar
Philipp Frank committed
368
    def make(offset_amplitude_mean, offset_amplitude_stddev, prefix,
Martin Reinecke's avatar
Martin Reinecke committed
369 370
             total_N=0,
             dofdex=None):
Philipp Frank's avatar
Philipp Frank committed
371 372
        if dofdex is None:
            dofdex = np.full(total_N, 0)
373 374
        elif len(dofdex) != total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Frank's avatar
Philipp Frank committed
375
        N = max(dofdex) + 1 if total_N > 0 else 0
376 377
        zm = _LognormalMomentMatching(offset_amplitude_mean,
                                      offset_amplitude_stddev,
Philipp Haim's avatar
Philipp Haim committed
378
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
379
                                      N)
Philipp Frank's avatar
fixup  
Philipp Frank committed
380
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
381
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
Philipp Haim's avatar
Philipp Haim committed
382
        return CorrelatedFieldMaker(zm, prefix, total_N)
383 384

    def add_fluctuations(self,
385
                         position_space,
386 387 388 389 390 391 392 393
                         fluctuations_mean,
                         fluctuations_stddev,
                         flexibility_mean,
                         flexibility_stddev,
                         asperity_mean,
                         asperity_stddev,
                         loglogavgslope_mean,
                         loglogavgslope_stddev,
Martin Reinecke's avatar
Martin Reinecke committed
394 395 396 397
                         prefix='',
                         index=None,
                         dofdex=None,
                         harmonic_partner=None):
Philipp Frank's avatar
Philipp Frank committed
398 399
        if harmonic_partner is None:
            harmonic_partner = position_space.get_default_codomain()
Philipp Frank's avatar
Fixup  
Philipp Frank committed
400 401 402
        else:
            position_space.check_codomain(harmonic_partner)
            harmonic_partner.check_codomain(position_space)
403

Philipp Haim's avatar
Philipp Haim committed
404 405
        if dofdex is None:
            dofdex = np.full(self._total_N, 0)
406 407
        elif len(dofdex) != self._total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Haim's avatar
Philipp Haim committed
408

Philipp Haim's avatar
Philipp Haim committed
409
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
410
            space = 1
Philipp Haim's avatar
Philipp Haim committed
411 412
            N = max(dofdex) + 1
            position_space = makeDomain((UnstructuredDomain(N), position_space))
Philipp Haim's avatar
Philipp Haim committed
413 414
        else:
            space = 0
Philipp Haim's avatar
Philipp Haim committed
415
            N = 0
Philipp Haim's avatar
Philipp Haim committed
416
            position_space = makeDomain(position_space)
Philipp Arras's avatar
Philipp Arras committed
417
        prefix = str(prefix)
Martin Reinecke's avatar
Martin Reinecke committed
418
        # assert isinstance(position_space[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
419

Philipp Arras's avatar
Philipp Arras committed
420 421
        fluct = _LognormalMomentMatching(fluctuations_mean,
                                         fluctuations_stddev,
422
                                         self._prefix + prefix + 'fluctuations',
Philipp Haim's avatar
Philipp Haim committed
423
                                         N)
Philipp Arras's avatar
Philipp Arras committed
424
        flex = _LognormalMomentMatching(flexibility_mean, flexibility_stddev,
425
                                        self._prefix + prefix + 'flexibility',
Philipp Haim's avatar
Philipp Haim committed
426
                                        N)
Philipp Arras's avatar
Philipp Arras committed
427
        asp = _LognormalMomentMatching(asperity_mean, asperity_stddev,
Philipp Arras's avatar
Philipp Arras committed
428
                                       self._prefix + prefix + 'asperity', N)
429
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
430
                        self._prefix + prefix + 'loglogavgslope', N)
Philipp Arras's avatar
Philipp Arras committed
431 432
        amp = _Amplitude(PowerSpace(harmonic_partner), fluct, flex, asp, avgsl,
                         self._azm, position_space[-1].total_volume,
433
                         self._prefix + prefix + 'spectrum', dofdex)
Philipp Haim's avatar
Philipp Haim committed
434

435 436
        if index is not None:
            self._a.insert(index, amp)
437
            self._position_spaces.insert(index, position_space)
438 439
        else:
            self._a.append(amp)
440
            self._position_spaces.append(position_space)
441

Philipp Frank's avatar
fixup  
Philipp Frank committed
442
    def _finalize_from_op(self):
Philipp Haim's avatar
Philipp Haim committed
443
        n_amplitudes = len(self._a)
Philipp Haim's avatar
Philipp Haim committed
444
        if self._total_N > 0:
Philipp Arras's avatar
Philipp Arras committed
445 446 447
            hspace = makeDomain(
                [UnstructuredDomain(self._total_N)] +
                [dd.target[-1].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
448 449
            spaces = tuple(range(1, n_amplitudes + 1))
            amp_space = 1
Philipp Haim's avatar
Philipp Haim committed
450 451
        else:
            hspace = makeDomain(
Philipp Arras's avatar
Philipp Arras committed
452
                [dd.target[0].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
453
            spaces = tuple(range(n_amplitudes))
Philipp Haim's avatar
Philipp Haim committed
454
            amp_space = 0
455

Martin Reinecke's avatar
Martin Reinecke committed
456
        expander = ContractionOperator(hspace, spaces=spaces).adjoint
Philipp Frank's avatar
fixup  
Philipp Frank committed
457
        azm = expander @ self._azm
458

459
        ht = HarmonicTransformOperator(hspace,
Philipp Haim's avatar
Philipp Haim committed
460
                                       self._position_spaces[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
461
                                       space=spaces[0])
462
        for i in range(1, n_amplitudes):
463 464 465 466 467 468 469
            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]
Philipp Haim's avatar
Philipp Haim committed
470
            pd = PowerDistributor(co.target, pp, amp_space)
471 472 473
            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
474

Philipp Arras's avatar
Formats  
Philipp Arras committed
475
    def finalize(self, offset=None, prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
476 477 478
        """
        offset vs zeromode: volume factor
        """
Philipp Frank's avatar
fixup  
Philipp Frank committed
479
        op = self._finalize_from_op()
Philipp Arras's avatar
Philipp Arras committed
480
        if offset is not None:
481 482 483 484 485 486 487
            # 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
488
        self.statistics_summary(prior_info)
489 490
        return op

491 492 493 494 495 496
    def statistics_summary(self, prior_info):
        from ..sugar import from_random

        if prior_info == 0:
            return

497 498
        lst = [('Offset amplitude', self.amplitude_total_offset),
               ('Total fluctuation amplitude', self.total_fluctuation)]
499
        namps = len(self._a)
500 501 502 503 504 505 506 507
        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:
508 509 510
            sc = StatCalculator()
            for _ in range(prior_info):
                sc.add(op(from_random('normal', op.domain)))
Martin Reinecke's avatar
merge  
Martin Reinecke committed
511 512
            mean = sc.mean.val
            stddev = sc.var.sqrt().val
513
            for m, s in zip(mean.flatten(), stddev.flatten()):
514
                logger.info('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
515 516 517

    def moment_slice_to_average(self, fluctuations_slice_mean, nsamples=1000):
        fluctuations_slice_mean = float(fluctuations_slice_mean)
518 519 520
        if not fluctuations_slice_mean > 0:
            msg = "fluctuations_slice_mean must be greater zero; got {!r}"
            raise ValueError(msg.format(fluctuations_slice_mean))
521
        from ..sugar import from_random
522 523
        scm = 1.
        for a in self._a:
Philipp Haim's avatar
Philipp Haim committed
524
            op = a.fluctuation_amplitude*self._azm.one_over()
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
525
            res = np.array([op(from_random('normal', op.domain)).val
526 527
                            for _ in range(nsamples)])
            scm *= res**2 + 1.
528
        return fluctuations_slice_mean/np.mean(np.sqrt(scm))
529

Philipp Arras's avatar
Philipp Arras committed
530
    @property
Philipp Haim's avatar
Philipp Haim committed
531
    def normalized_amplitudes(self):
532
        return self._a
Philipp Arras's avatar
Philipp Arras committed
533

Philipp Haim's avatar
Philipp Haim committed
534 535 536 537 538 539 540
    @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
541 542
        dom = self._a[0].target
        expand = ContractionOperator(dom, len(dom)-1).adjoint
Philipp Haim's avatar
Philipp Haim committed
543 544
        return self._a[0]*(expand @ self.amplitude_total_offset)

545 546 547
    @property
    def amplitude_total_offset(self):
        return self._azm
Philipp Arras's avatar
Philipp Arras committed
548 549

    @property
550
    def total_fluctuation(self):
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 len(self._a) == 1:
555
            return self.average_fluctuation(0)
556 557
        q = 1.
        for a in self._a:
Philipp Haim's avatar
Philipp Haim committed
558
            fl = a.fluctuation_amplitude*self._azm.one_over()
Philipp Arras's avatar
Philipp Arras committed
559
            q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Philipp Arras's avatar
Formats  
Philipp Arras committed
560
        return (Adder(full(q.target, -1.)) @ q).sqrt()*self._azm
561

Philipp Arras's avatar
Philipp Arras committed
562
    def slice_fluctuation(self, space):
563
        """Returns operator which acts on prior or posterior samples"""
564
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
565
            raise NotImplementedError
566
        if space >= len(self._a):
567
            raise ValueError("invalid space specified; got {!r}".format(space))
568
        if len(self._a) == 1:
569
            return self.average_fluctuation(0)
570 571
        q = 1.
        for j in range(len(self._a)):
Philipp Haim's avatar
Philipp Haim committed
572
            fl = self._a[j].fluctuation_amplitude*self._azm.one_over()
573
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
574
                q = q*fl**2
575
            else:
Philipp Arras's avatar
Philipp Arras committed
576
                q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Philipp Arras's avatar
Formats  
Philipp Arras committed
577
        return q.sqrt()*self._azm
Philipp Arras's avatar
Philipp Arras committed
578 579

    def average_fluctuation(self, space):
580
        """Returns operator which acts on prior or posterior samples"""
581
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
582
            raise NotImplementedError
583
        if space >= len(self._a):
584
            raise ValueError("invalid space specified; got {!r}".format(space))
585
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
586 587
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
588

589 590
    @staticmethod
    def offset_amplitude_realized(samples):
Philipp Haim's avatar
Philipp Haim committed
591
        spaces = _structured_spaces(samples[0].domain)
592 593
        res = 0.
        for s in samples:
Philipp Haim's avatar
Philipp Haim committed
594 595 596 597
            res = res + s.mean(spaces)**2
        if np.isscalar(res):
            return np.sqrt(res/len(samples))
        return np.sqrt(res.val/len(samples))
Philipp Arras's avatar
Philipp Arras committed
598

599 600 601 602 603 604 605 606
    @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."""
Philipp Haim's avatar
Philipp Haim committed
607 608
        spaces = _structured_spaces(samples[0].domain)
        if space >= len(spaces):
609
            raise ValueError("invalid space specified; got {!r}".format(space))
Philipp Haim's avatar
Philipp Haim committed
610
        if len(spaces) == 1:
611
            return _total_fluctuation_realized(samples)
Philipp Haim's avatar
Philipp Haim committed
612
        space = space + spaces[0]
Philipp Arras's avatar
Philipp Arras committed
613
        res1, res2 = 0., 0.
614
        for s in samples:
Philipp Frank's avatar
fixes  
Philipp Frank committed
615 616 617 618
            res1 = res1 + s**2
            res2 = res2 + s.mean(space)**2
        res1 = res1/len(samples)
        res2 = res2/len(samples)
Philipp Haim's avatar
Philipp Haim committed
619 620 621 622
        res = res1.mean(spaces) - res2.mean(spaces[:-1])
        if np.isscalar(res):
            return np.sqrt(res)
        return np.sqrt(res.val)
Philipp Frank's avatar
fixes  
Philipp Frank committed
623

Philipp Arras's avatar
Philipp Arras committed
624
    @staticmethod
625 626 627
    def average_fluctuation_realized(samples, space):
        """Computes average fluctuations from collection of field (defined in signal
        space) realizations."""
Philipp Haim's avatar
Philipp Haim committed
628 629
        spaces = _structured_spaces(samples[0].domain)
        if space >= len(spaces):
630
            raise ValueError("invalid space specified; got {!r}".format(space))
Philipp Haim's avatar
Philipp Haim committed
631
        if len(spaces) == 1:
632
            return _total_fluctuation_realized(samples)
Philipp Haim's avatar
Philipp Haim committed
633 634 635 636 637 638
        space = space + spaces[0]
        sub_spaces = set(spaces)
        sub_spaces.remove(space)
        sub_dom = makeDomain([samples[0].domain[ind]
                              for ind in set([0,]) | set([space,])])
        co = ContractionOperator(sub_dom, len(sub_dom)-1)
Philipp Arras's avatar
Philipp Arras committed
639 640
        res = 0.
        for s in samples:
Philipp Haim's avatar
Philipp Haim committed
641 642 643 644 645 646 647 648 649
            r = s.mean(sub_spaces)
            if min(spaces) == 0:
                res = res + (r - r.mean(spaces[:-1]))**2
            else:
                res = res + (r - co.adjoint(r.mean(spaces[:-1])))**2
        res = res.mean(spaces[0])/len(samples)
        if np.isscalar(res):
            return np.sqrt(res)
        return np.sqrt(res.val)