correlated_fields.py 22.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

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

46

Martin Reinecke's avatar
Martin Reinecke committed
47
def _lognormal_moments(mean, sig, N=0):
Philipp Haim's avatar
Philipp Haim committed
48 49 50 51
    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
52
    assert np.all(mean > 0)
53
    assert np.all(sig > 0)
Philipp Arras's avatar
Philipp Arras committed
54 55
    logsig = np.sqrt(np.log((sig/mean)**2 + 1))
    logmean = np.log(mean) - logsig**2/2
56
    return logmean, logsig
Philipp Arras's avatar
Philipp Arras committed
57 58


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


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


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


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


def _stats(op, samples):
    sc = StatCalculator()
    for s in samples:
        sc.add(op(s.extract(op.domain)))
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
106
    return sc.mean.val, sc.var.sqrt().val
107 108


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


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

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

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

Philipp Arras's avatar
Philipp Arras committed
152 153

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

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

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

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


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

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


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

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


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

248

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

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

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

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

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

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

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

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

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

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

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

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

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

339 340

class CorrelatedFieldMaker:
Philipp Haim's avatar
Philipp Haim committed
341
    def __init__(self, amplitude_offset, prefix, total_N):
Philipp Frank's avatar
fixup  
Philipp Frank committed
342
        assert isinstance(amplitude_offset, Operator)
343
        self._a = []
344
        self._position_spaces = []
Philipp Arras's avatar
Formats  
Philipp Arras committed
345

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

350
    @staticmethod
Philipp Frank's avatar
Philipp Frank committed
351
    def make(offset_amplitude_mean, offset_amplitude_stddev, prefix,
Martin Reinecke's avatar
Martin Reinecke committed
352 353
             total_N=0,
             dofdex=None):
Philipp Frank's avatar
Philipp Frank committed
354 355 356 357 358
        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
359 360
        zm = _LognormalMomentMatching(offset_amplitude_mean,
                                      offset_amplitude_stddev,
Philipp Haim's avatar
Philipp Haim committed
361
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
362
                                      N)
Philipp Frank's avatar
fixup  
Philipp Frank committed
363
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
364
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
Philipp Haim's avatar
Philipp Haim committed
365
        return CorrelatedFieldMaker(zm, prefix, total_N)
366 367

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

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

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

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

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

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

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

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

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

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

Philipp Arras's avatar
Formats  
Philipp Arras committed
466
    def finalize(self, offset=None, prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
467 468 469 470
        """
        offset vs zeromode: volume factor
        """
        if offset is not None:
471
            raise NotImplementedError
Philipp Arras's avatar
Philipp Arras committed
472
            offset = float(offset)
473

Philipp Frank's avatar
fixup  
Philipp Frank committed
474
        op = self._finalize_from_op()
475 476 477 478 479 480 481 482 483 484 485 486
        if prior_info > 0:
            from ..sugar import from_random
            samps = [
                from_random('normal', op.domain) for _ in range(prior_info)
            ]
            self.statistics_summary(samps)
        return op

    def statistics_summary(self, samples):
        lst = [('Offset amplitude', self.amplitude_total_offset),
               ('Total fluctuation amplitude', self.total_fluctuation)]

487
        namps = len(self._a)
488 489 490 491 492 493 494 495 496
        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:
            mean, stddev = _stats(op, samples)
497 498
            for m, s in zip(mean.flatten(), stddev.flatten()):
                print('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
499 500 501 502

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

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

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

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

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

Philipp Arras's avatar
Philipp Arras committed
544
    def slice_fluctuation(self, space):
545
        """Returns operator which acts on prior or posterior samples"""
546
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
547
            raise NotImplementedError
548 549
        assert space < len(self._a)
        if len(self._a) == 1:
550
            return self.average_fluctuation(0)
551 552
        q = 1.
        for j in range(len(self._a)):
Philipp Haim's avatar
Philipp Haim committed
553
            fl = self._a[j].fluctuation_amplitude*self._azm.one_over()
554
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
555
                q = q*fl**2
556
            else:
Philipp Arras's avatar
Philipp Arras committed
557
                q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Philipp Arras's avatar
Formats  
Philipp Arras committed
558
        return q.sqrt()*self._azm
Philipp Arras's avatar
Philipp Arras committed
559 560

    def average_fluctuation(self, space):
561
        """Returns operator which acts on prior or posterior samples"""
562
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
563
            raise NotImplementedError
564 565
        assert space < len(self._a)
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
566 567
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
568

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

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