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

Philipp Arras's avatar
Philipp Arras committed
19
import numpy as np
20

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

38

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

48

Martin Reinecke's avatar
Martin Reinecke committed
49
def _lognormal_moments(mean, sig, N=0):
Philipp Haim's avatar
Philipp Haim committed
50 51 52 53
    if N == 0:
        mean, sig = np.asfarray(mean), np.asfarray(sig)
    else:
        mean, sig = (_reshaper(param, N) for param in (mean, sig))
Martin Reinecke's avatar
Martin Reinecke committed
54
    assert np.all(mean > 0)
55
    assert np.all(sig > 0)
Philipp Arras's avatar
Philipp Arras committed
56 57
    logsig = np.sqrt(np.log((sig/mean)**2 + 1))
    logmean = np.log(mean) - logsig**2/2
58
    return logmean, logsig
Philipp Arras's avatar
Philipp Arras committed
59 60


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


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


Philipp Arras's avatar
Philipp Arras committed
78
def _relative_log_k_lengths(power_space):
Philipp Arras's avatar
Philipp Arras committed
79 80
    """Log-distance to first bin
    logkl.shape==power_space.shape, logkl[0]=logkl[1]=0"""
Philipp Arras's avatar
Philipp Arras committed
81 82 83 84 85 86
    power_space = DomainTuple.make(power_space)
    assert isinstance(power_space[0], PowerSpace)
    assert len(power_space) == 1
    logkl = _log_k_lengths(power_space[0])
    assert logkl.shape[0] == power_space[0].shape[0] - 1
    logkl -= logkl[0]
Philipp Arras's avatar
Philipp Arras committed
87
    return np.insert(logkl, 0, 0)
Philipp Arras's avatar
Philipp Arras committed
88 89


Philipp Arras's avatar
Philipp Arras committed
90
def _log_vol(power_space):
91
    power_space = makeDomain(power_space)
Philipp Arras's avatar
Philipp Arras committed
92 93 94 95 96
    assert isinstance(power_space[0], PowerSpace)
    logk_lengths = _log_k_lengths(power_space[0])
    return logk_lengths[1:] - logk_lengths[:-1]


Philipp Haim's avatar
Philipp Haim committed
97
def _total_fluctuation_realized(samples):
98 99
    res = 0.
    for s in samples:
Philipp Haim's avatar
Fixes  
Philipp Haim committed
100
        res = res + (s - s.mean())**2
Philipp Haim's avatar
Philipp Haim committed
101
    return np.sqrt((res/len(samples)).mean())
102 103


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


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

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

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

Philipp Arras's avatar
Philipp Arras committed
147 148

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

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

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

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


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

    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
210
        return self._specsum(self._mode_multiplicity*spec)**(-0.5)*amp
Philipp Arras's avatar
Philipp Arras committed
211 212 213


class _SpecialSum(EndomorphicOperator):
Martin Reinecke's avatar
Martin Reinecke committed
214
    def __init__(self, domain, space=0):
Philipp Arras's avatar
Philipp Arras committed
215 216
        self._domain = makeDomain(domain)
        self._capability = self.TIMES | self.ADJOINT_TIMES
217
        self._contractor = ContractionOperator(domain, space)
Philipp Arras's avatar
Philipp Arras committed
218 219 220

    def apply(self, x, mode):
        self._check_input(x, mode)
221
        return self._contractor.adjoint(self._contractor(x))
Philipp Arras's avatar
Philipp Arras committed
222 223


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

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

    def apply(self, x, mode):
        self._check_input(x, mode)
        x = x.to_global_data()
        if mode == self.TIMES:
            res = x[self._dofdex]
        else:
            res = np.empty(self._tgt(mode).shape)
            res[self._dofdex] = x
        return from_global_data(self._tgt(mode), res)
Martin Reinecke's avatar
Martin Reinecke committed
242

243

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

272
        twolog = _TwoLogIntegrations(target, space)
Philipp Arras's avatar
Philipp Arras committed
273
        dom = twolog.domain
274

275
        shp = dom[space].shape
Martin Reinecke's avatar
Martin Reinecke committed
276 277
        expander = ContractionOperator(dom, spaces=space).adjoint
        ps_expander = ContractionOperator(twolog.target, spaces=space).adjoint
Philipp Arras's avatar
Philipp Arras committed
278 279 280

        # Prepare constant fields
        foo = np.zeros(shp)
281 282
        foo[0] = foo[1] = np.sqrt(_log_vol(target[space]))
        vflex = DiagonalOperator(from_global_data(dom[space], foo), dom, space)
Philipp Arras's avatar
Philipp Arras committed
283 284 285

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

        foo = np.ones(shp)
289 290
        foo[0] = _log_vol(target[space])**2/12.
        shift = DiagonalOperator(from_global_data(dom[space], foo), dom, space)
Martin Reinecke's avatar
Martin Reinecke committed
291

292
        vslope = DiagonalOperator(
Martin Reinecke's avatar
Martin Reinecke committed
293 294 295
            from_global_data(target[space],
                             _relative_log_k_lengths(target[space])),
            target, space)
296 297

        foo, bar = [np.zeros(target[space].shape) for _ in range(2)]
Philipp Arras's avatar
Philipp Arras committed
298
        bar[1:] = foo[0] = totvol
Martin Reinecke's avatar
Martin Reinecke committed
299
        vol0, vol1 = [DiagonalOperator(from_global_data(target[space], aa),
Philipp Frank's avatar
cleanup  
Philipp Frank committed
300
                                       target, space) for aa in (foo, bar)]
301

Martin Reinecke's avatar
Martin Reinecke committed
302
        # Prepare fields for Adder
303
        shift, vol0 = [op(full(op.domain, 1)) for op in (shift, vol0)]
Philipp Arras's avatar
Philipp Arras committed
304 305
        # End prepare constant fields

306 307 308 309
        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
310
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Arras's avatar
Philipp Arras committed
311 312

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

Philipp Arras's avatar
Philipp Arras committed
326 327
        self.apply = op.apply
        self._domain, self._target = op.domain, op.target
328
        self._space = space
Philipp Arras's avatar
Philipp Arras committed
329

Philipp Arras's avatar
Philipp Arras committed
330 331 332 333
    @property
    def fluctuation_amplitude(self):
        return self._fluc

334 335

class CorrelatedFieldMaker:
Philipp Haim's avatar
Philipp Haim committed
336
    def __init__(self, amplitude_offset, prefix, total_N):
Philipp Frank's avatar
fixup  
Philipp Frank committed
337
        assert isinstance(amplitude_offset, Operator)
338
        self._a = []
339
        self._position_spaces = []
Philipp Arras's avatar
Formats  
Philipp Arras committed
340

341 342
        self._azm = amplitude_offset
        self._prefix = prefix
Philipp Haim's avatar
Philipp Haim committed
343
        self._total_N = total_N
Philipp Arras's avatar
Formats  
Philipp Arras committed
344

345
    @staticmethod
Philipp Frank's avatar
Philipp Frank committed
346
    def make(offset_amplitude_mean, offset_amplitude_stddev, prefix,
Martin Reinecke's avatar
Martin Reinecke committed
347 348
             total_N=0,
             dofdex=None):
Philipp Frank's avatar
Philipp Frank committed
349 350 351 352 353
        if dofdex is None:
            dofdex = np.full(total_N, 0)
        else:
            assert len(dofdex) == total_N
        N = max(dofdex) + 1 if total_N > 0 else 0
354 355
        zm = _LognormalMomentMatching(offset_amplitude_mean,
                                      offset_amplitude_stddev,
Philipp Haim's avatar
Philipp Haim committed
356
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
357
                                      N)
Philipp Frank's avatar
fixup  
Philipp Frank committed
358
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
359
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
Philipp Haim's avatar
Philipp Haim committed
360
        return CorrelatedFieldMaker(zm, prefix, total_N)
361 362

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

Philipp Haim's avatar
Philipp Haim committed
382 383 384 385 386
        if dofdex is None:
            dofdex = np.full(self._total_N, 0)
        else:
            assert len(dofdex) == self._total_N

Philipp Haim's avatar
Philipp Haim committed
387
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
388
            space = 1
Philipp Haim's avatar
Philipp Haim committed
389 390
            N = max(dofdex) + 1
            position_space = makeDomain((UnstructuredDomain(N), position_space))
Philipp Haim's avatar
Philipp Haim committed
391 392
        else:
            space = 0
Philipp Haim's avatar
Philipp Haim committed
393
            N = 0
Philipp Haim's avatar
Philipp Haim committed
394
            position_space = makeDomain(position_space)
Philipp Arras's avatar
Philipp Arras committed
395
        prefix = str(prefix)
Martin Reinecke's avatar
Martin Reinecke committed
396
        # assert isinstance(position_space[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
397

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

415 416
        if index is not None:
            self._a.insert(index, amp)
417
            self._position_spaces.insert(index, position_space)
418 419
        else:
            self._a.append(amp)
420
            self._position_spaces.append(position_space)
421

Philipp Frank's avatar
fixup  
Philipp Frank committed
422
    def _finalize_from_op(self):
Philipp Haim's avatar
Philipp Haim committed
423
        n_amplitudes = len(self._a)
Philipp Haim's avatar
Philipp Haim committed
424
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
425
            hspace = makeDomain([UnstructuredDomain(self._total_N)] +
Martin Reinecke's avatar
Martin Reinecke committed
426 427
                                [dd.target[-1].harmonic_partner
                                    for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
428 429
            spaces = tuple(range(1, n_amplitudes + 1))
            amp_space = 1
Philipp Haim's avatar
Philipp Haim committed
430 431
        else:
            hspace = makeDomain(
432
                     [dd.target[0].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
433
            spaces = tuple(range(n_amplitudes))
Philipp Haim's avatar
Philipp Haim committed
434
            amp_space = 0
435

Martin Reinecke's avatar
Martin Reinecke committed
436
        expander = ContractionOperator(hspace, spaces=spaces).adjoint
Philipp Frank's avatar
fixup  
Philipp Frank committed
437
        azm = expander @ self._azm
438

439
        ht = HarmonicTransformOperator(hspace,
Philipp Haim's avatar
Philipp Haim committed
440
                                       self._position_spaces[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
441
                                       space=spaces[0])
442
        for i in range(1, n_amplitudes):
443
            ht = (HarmonicTransformOperator(ht.target,
Philipp Haim's avatar
Philipp Haim committed
444
                                            self._position_spaces[i][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
445
                                            space=spaces[i]) @ ht)
446

Philipp Haim's avatar
Philipp Haim committed
447
        pd = PowerDistributor(hspace, self._a[0].target[amp_space], amp_space)
448
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
449
            pd = (pd @ PowerDistributor(pd.domain,
Philipp Haim's avatar
Philipp Haim committed
450
                                        self._a[i].target[amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
451
                                        space=spaces[i]))
Philipp Arras's avatar
Philipp Arras committed
452

453 454
        a = ContractionOperator(pd.domain, spaces[1:]).adjoint @ self._a[0]
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
455
            co = ContractionOperator(pd.domain,
Martin Reinecke's avatar
Martin Reinecke committed
456
                                     spaces[:i] + spaces[i+1:])
457
            a = a*(co.adjoint @ self._a[i])
Philipp Arras's avatar
Philipp Arras committed
458

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

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

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

        if prior_info == 0:
            return

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

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

Philipp Arras's avatar
Philipp Arras committed
514
    @property
Philipp Haim's avatar
Philipp Haim committed
515
    def normalized_amplitudes(self):
516
        return self._a
Philipp Arras's avatar
Philipp Arras committed
517

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

529 530 531
    @property
    def amplitude_total_offset(self):
        return self._azm
Philipp Arras's avatar
Philipp Arras committed
532 533

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

Philipp Arras's avatar
Philipp Arras committed
546
    def slice_fluctuation(self, space):
547
        """Returns operator which acts on prior or posterior samples"""
548
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
549
            raise NotImplementedError
550 551
        assert space < len(self._a)
        if len(self._a) == 1:
552
            return self.average_fluctuation(0)
553 554
        q = 1.
        for j in range(len(self._a)):
Philipp Haim's avatar
Philipp Haim committed
555
            fl = self._a[j].fluctuation_amplitude*self._azm.one_over()
556
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
557
                q = q*fl**2
558
            else:
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 q.sqrt()*self._azm
Philipp Arras's avatar
Philipp Arras committed
561 562

    def average_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 567
        assert space < len(self._a)
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
568 569
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
570

571 572
    @staticmethod
    def offset_amplitude_realized(samples):
573 574
        res = 0.
        for s in samples:
Philipp Frank's avatar
fixes  
Philipp Frank committed
575
            res = res + s.mean()**2
576
        return np.sqrt(res/len(samples))
Philipp Arras's avatar
Philipp Arras committed
577

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