correlated_fields.py 23.7 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))
54 55 56 57 58
    if not np.all(mean > 0):
        raise ValueError(f"mean must be greater 0; got {mean!r}")
    if not np.all(sig > 0):
        raise ValueError(f"sig must be greater 0; got {sig!r}")

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


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


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


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


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


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


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

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

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

Philipp Arras's avatar
Philipp Arras committed
150 151

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

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

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

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


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

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


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

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


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

246

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

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

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

        # Prepare constant fields
        foo = np.zeros(shp)
284 285
        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
286 287 288

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

        foo = np.ones(shp)
292 293
        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
294

295
        vslope = DiagonalOperator(
Martin Reinecke's avatar
Martin Reinecke committed
296 297 298
            from_global_data(target[space],
                             _relative_log_k_lengths(target[space])),
            target, space)
299 300

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

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

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

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

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

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

337 338

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Philipp Arras's avatar
Formats  
Philipp Arras committed
465
    def finalize(self, offset=None, prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
466 467 468
        """
        offset vs zeromode: volume factor
        """
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 501 502
            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()
503 504
            for m, s in zip(mean.flatten(), stddev.flatten()):
                print('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
505 506 507

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

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

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

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

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

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

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

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

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

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

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