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

19 20 21
from functools import reduce
from operator import mul

Philipp Arras's avatar
Philipp Arras committed
22
import numpy as np
23

Philipp Arras's avatar
Philipp Arras committed
24
from ..domain_tuple import DomainTuple
Philipp Arras's avatar
Philipp Arras committed
25 26
from ..domains.power_space import PowerSpace
from ..domains.unstructured_domain import UnstructuredDomain
Philipp Arras's avatar
Philipp Arras committed
27
from ..field import Field
Philipp Arras's avatar
Philipp Arras committed
28
from ..linearization import Linearization
29
from ..logger import logger
Philipp Arras's avatar
Philipp Arras committed
30
from ..multi_field import MultiField
Philipp Arras's avatar
Philipp Arras committed
31
from ..operators.adder import Adder
32
from ..operators.contraction_operator import ContractionOperator
Philipp Arras's avatar
Philipp Arras committed
33
from ..operators.diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
34
from ..operators.distributors import PowerDistributor
Philipp Arras's avatar
Philipp Arras committed
35
from ..operators.endomorphic_operator import EndomorphicOperator
Martin Reinecke's avatar
Martin Reinecke committed
36
from ..operators.harmonic_operators import HarmonicTransformOperator
Philipp Arras's avatar
Philipp Arras committed
37
from ..operators.linear_operator import LinearOperator
Philipp Arras's avatar
Philipp Arras committed
38
from ..operators.operator import Operator
Philipp Arras's avatar
Philipp Arras committed
39
from ..operators.simple_linear_operators import ducktape
40
from ..probing import StatCalculator
Philipp Arras's avatar
Philipp Arras committed
41
from ..sugar import full, makeDomain, makeField, makeOp
42

43

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

53

Martin Reinecke's avatar
Martin Reinecke committed
54
def _lognormal_moments(mean, sig, N=0):
Philipp Haim's avatar
Philipp Haim committed
55 56 57 58
    if N == 0:
        mean, sig = np.asfarray(mean), np.asfarray(sig)
    else:
        mean, sig = (_reshaper(param, N) for param in (mean, sig))
59
    if not np.all(mean > 0):
60
        raise ValueError("mean must be greater 0; got {!r}".format(mean))
61
    if not np.all(sig > 0):
62
        raise ValueError("sig must be greater 0; got {!r}".format(sig))
63

64
    logsig = np.sqrt(np.log1p((sig/mean)**2))
Philipp Arras's avatar
Philipp Arras committed
65
    logmean = np.log(mean) - logsig**2/2
66
    return logmean, logsig
Philipp Arras's avatar
Philipp Arras committed
67 68


Martin Reinecke's avatar
Martin Reinecke committed
69
def _normal(mean, sig, key, N=0):
Philipp Haim's avatar
Philipp Haim committed
70
    if N == 0:
Philipp Haim's avatar
Philipp Haim committed
71
        domain = DomainTuple.scalar_domain()
Philipp Haim's avatar
Philipp Haim committed
72
        mean, sig = np.asfarray(mean), np.asfarray(sig)
73 74 75 76 77
        return Adder(makeField(domain, mean)) @ (
            sig * ducktape(domain, None, key))

    domain = UnstructuredDomain(N)
    mean, sig = (_reshaper(param, N) for param in (mean, sig))
Martin Reinecke's avatar
Martin Reinecke committed
78 79
    return Adder(makeField(domain, mean)) @ (DiagonalOperator(
        makeField(domain, sig)) @ ducktape(domain, None, key))
Philipp Arras's avatar
Philipp Arras committed
80 81


Philipp Arras's avatar
Philipp Arras committed
82
def _log_k_lengths(pspace):
Philipp Arras's avatar
Philipp Arras committed
83
    """Log(k_lengths) without zeromode"""
Philipp Arras's avatar
Philipp Arras committed
84 85 86
    return np.log(pspace.k_lengths[1:])


Philipp Arras's avatar
Philipp Arras committed
87
def _relative_log_k_lengths(power_space):
Philipp Arras's avatar
Philipp Arras committed
88 89
    """Log-distance to first bin
    logkl.shape==power_space.shape, logkl[0]=logkl[1]=0"""
Philipp Arras's avatar
Philipp Arras committed
90 91 92 93 94 95
    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
96
    return np.insert(logkl, 0, 0)
Philipp Arras's avatar
Philipp Arras committed
97 98


Philipp Arras's avatar
Philipp Arras committed
99
def _log_vol(power_space):
100
    power_space = makeDomain(power_space)
Philipp Arras's avatar
Philipp Arras committed
101 102 103 104 105
    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
106 107 108 109 110 111
def _structured_spaces(domain):
    if isinstance(domain[0], UnstructuredDomain):
        return np.arange(1, len(domain))
    return np.arange(len(domain))


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


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


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

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

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

Philipp Arras's avatar
Philipp Arras committed
166 167

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

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

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

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


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

    def apply(self, x):
        self._check_input(x)
Martin Reinecke's avatar
Martin Reinecke committed
227
        amp = x.ptw("exp")
228
        spec = amp**2
Philipp Arras's avatar
Philipp Arras committed
229 230
        # FIXME This normalizes also the zeromode which is supposed to be left
        # untouched by this operator
Philipp Arras's avatar
Philipp Arras committed
231
        return self._specsum(self._mode_multiplicity(spec))**(-0.5)*amp
Philipp Arras's avatar
Philipp Arras committed
232 233 234


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

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


Philipp Haim's avatar
Philipp Haim committed
245
class _Distributor(LinearOperator):
Lukas Platz's avatar
Lukas Platz committed
246
    def __init__(self, dofdex, domain, target):
Philipp Haim's avatar
Philipp Haim committed
247 248 249 250 251 252 253 254
        self._dofdex = dofdex

        self._target = makeDomain(target)
        self._domain = makeDomain(domain)
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
255
        x = x.val
Philipp Haim's avatar
Philipp Haim committed
256 257 258 259 260
        if mode == self.TIMES:
            res = x[self._dofdex]
        else:
            res = np.empty(self._tgt(mode).shape)
            res[self._dofdex] = x
Martin Reinecke's avatar
Martin Reinecke committed
261
        return makeField(self._tgt(mode), res)
Martin Reinecke's avatar
Martin Reinecke committed
262

263

264 265
class _Amplitude(Operator):
    def __init__(self, target, fluctuations, flexibility, asperity,
266
                 loglogavgslope, azm, totvol, key, dofdex):
Philipp Arras's avatar
Philipp Arras committed
267 268 269 270 271 272 273 274 275 276 277
        """
        fluctuations > 0
        flexibility > 0
        asperity > 0
        loglogavgslope probably negative
        """
        assert isinstance(fluctuations, Operator)
        assert isinstance(flexibility, Operator)
        assert isinstance(asperity, Operator)
        assert isinstance(loglogavgslope, Operator)

Philipp Haim's avatar
Philipp Haim committed
278 279
        if len(dofdex) > 0:
            N_copies = max(dofdex) + 1
Philipp Haim's avatar
Philipp Haim committed
280
            space = 1
Philipp Frank's avatar
cleanup  
Philipp Frank committed
281 282
            distributed_tgt = makeDomain((UnstructuredDomain(len(dofdex)),
                                          target))
Philipp Haim's avatar
Philipp Haim committed
283
            target = makeDomain((UnstructuredDomain(N_copies), target))
Lukas Platz's avatar
Lukas Platz committed
284
            Distributor = _Distributor(dofdex, target, distributed_tgt)
Philipp Haim's avatar
Philipp Haim committed
285
        else:
Philipp Haim's avatar
Philipp Haim committed
286
            N_copies = 0
Philipp Haim's avatar
Philipp Haim committed
287
            space = 0
Philipp Haim's avatar
Philipp Haim committed
288
            distributed_tgt = target = makeDomain(target)
Martin Reinecke's avatar
Martin Reinecke committed
289
        azm_expander = ContractionOperator(distributed_tgt, spaces=space).adjoint
Philipp Haim's avatar
Philipp Haim committed
290
        assert isinstance(target[space], PowerSpace)
Martin Reinecke's avatar
Martin Reinecke committed
291

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

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

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

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

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

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

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

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

327 328 329 330
        slope = vslope @ ps_expander @ loglogavgslope
        sig_flex = vflex @ expander @ flexibility
        sig_asp = vasp @ expander @ asperity
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Haim's avatar
Philipp Haim committed
331
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Arras's avatar
Philipp Arras committed
332 333

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

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

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

355 356

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

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

368
    @staticmethod
Lukas Platz's avatar
Lukas Platz committed
369
    def make(offset_mean, offset_std_mean, offset_std_std, prefix,
Martin Reinecke's avatar
Martin Reinecke committed
370 371
             total_N=0,
             dofdex=None):
Lukas Platz's avatar
Lukas Platz committed
372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388
        """Returns a CorrelatedFieldMaker object.

        Parameters
        ----------
        offset_mean : float
            Mean offset from zero of the correlated field to be made.
        offset_std_mean : float
            Mean standard deviation of the offset value.
        offset_std_std : float
            Standard deviation of the stddev of the offset value.
        prefix : string
            Prefix to the names of the domains of the cf operator to be made.
        total_N : integer
            ?
        dofdex : np.array
            ?
        """
Philipp Frank's avatar
Philipp Frank committed
389 390
        if dofdex is None:
            dofdex = np.full(total_N, 0)
391 392
        elif len(dofdex) != total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Frank's avatar
Philipp Frank committed
393
        N = max(dofdex) + 1 if total_N > 0 else 0
Lukas Platz's avatar
Lukas Platz committed
394 395
        zm = _LognormalMomentMatching(offset_std_mean,
                                      offset_std_std,
Philipp Haim's avatar
Philipp Haim committed
396
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
397
                                      N)
Philipp Frank's avatar
fixup  
Philipp Frank committed
398
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
399
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
400
        return CorrelatedFieldMaker(offset_mean, zm, prefix, total_N)
401 402

    def add_fluctuations(self,
403
                         position_space,
404 405 406 407 408 409 410 411
                         fluctuations_mean,
                         fluctuations_stddev,
                         flexibility_mean,
                         flexibility_stddev,
                         asperity_mean,
                         asperity_stddev,
                         loglogavgslope_mean,
                         loglogavgslope_stddev,
Martin Reinecke's avatar
Martin Reinecke committed
412 413 414 415
                         prefix='',
                         index=None,
                         dofdex=None,
                         harmonic_partner=None):
Philipp Frank's avatar
Philipp Frank committed
416 417
        if harmonic_partner is None:
            harmonic_partner = position_space.get_default_codomain()
Philipp Frank's avatar
Fixup  
Philipp Frank committed
418 419 420
        else:
            position_space.check_codomain(harmonic_partner)
            harmonic_partner.check_codomain(position_space)
421

Philipp Haim's avatar
Philipp Haim committed
422 423
        if dofdex is None:
            dofdex = np.full(self._total_N, 0)
424 425
        elif len(dofdex) != self._total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Haim's avatar
Philipp Haim committed
426

Philipp Haim's avatar
Philipp Haim committed
427 428 429
        if self._total_N > 0:
            N = max(dofdex) + 1
            position_space = makeDomain((UnstructuredDomain(N), position_space))
Philipp Haim's avatar
Philipp Haim committed
430
        else:
Philipp Haim's avatar
Philipp Haim committed
431
            N = 0
Philipp Haim's avatar
Philipp Haim committed
432
            position_space = makeDomain(position_space)
Philipp Arras's avatar
Philipp Arras committed
433
        prefix = str(prefix)
Martin Reinecke's avatar
Martin Reinecke committed
434
        # assert isinstance(position_space[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
435

Philipp Arras's avatar
Philipp Arras committed
436 437
        fluct = _LognormalMomentMatching(fluctuations_mean,
                                         fluctuations_stddev,
438
                                         self._prefix + prefix + 'fluctuations',
Philipp Haim's avatar
Philipp Haim committed
439
                                         N)
Philipp Arras's avatar
Philipp Arras committed
440
        flex = _LognormalMomentMatching(flexibility_mean, flexibility_stddev,
441
                                        self._prefix + prefix + 'flexibility',
Philipp Haim's avatar
Philipp Haim committed
442
                                        N)
Philipp Arras's avatar
Philipp Arras committed
443
        asp = _LognormalMomentMatching(asperity_mean, asperity_stddev,
Philipp Arras's avatar
Philipp Arras committed
444
                                       self._prefix + prefix + 'asperity', N)
445
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
446
                        self._prefix + prefix + 'loglogavgslope', N)
Philipp Arras's avatar
Philipp Arras committed
447 448
        amp = _Amplitude(PowerSpace(harmonic_partner), fluct, flex, asp, avgsl,
                         self._azm, position_space[-1].total_volume,
449
                         self._prefix + prefix + 'spectrum', dofdex)
Philipp Haim's avatar
Philipp Haim committed
450

451 452
        if index is not None:
            self._a.insert(index, amp)
453
            self._position_spaces.insert(index, position_space)
454 455
        else:
            self._a.append(amp)
456
            self._position_spaces.append(position_space)
457

Philipp Frank's avatar
fixup  
Philipp Frank committed
458
    def _finalize_from_op(self):
Philipp Haim's avatar
Philipp Haim committed
459
        n_amplitudes = len(self._a)
Philipp Haim's avatar
Philipp Haim committed
460
        if self._total_N > 0:
Philipp Arras's avatar
Philipp Arras committed
461 462 463
            hspace = makeDomain(
                [UnstructuredDomain(self._total_N)] +
                [dd.target[-1].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
464 465
            spaces = tuple(range(1, n_amplitudes + 1))
            amp_space = 1
Philipp Haim's avatar
Philipp Haim committed
466 467
        else:
            hspace = makeDomain(
Philipp Arras's avatar
Philipp Arras committed
468
                [dd.target[0].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
469
            spaces = tuple(range(n_amplitudes))
Philipp Haim's avatar
Philipp Haim committed
470
            amp_space = 0
471

Martin Reinecke's avatar
Martin Reinecke committed
472
        expander = ContractionOperator(hspace, spaces=spaces).adjoint
Philipp Frank's avatar
fixup  
Philipp Frank committed
473
        azm = expander @ self._azm
474

475
        ht = HarmonicTransformOperator(hspace,
Philipp Haim's avatar
Philipp Haim committed
476
                                       self._position_spaces[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
477
                                       space=spaces[0])
478
        for i in range(1, n_amplitudes):
479 480 481 482 483 484 485
            ht = HarmonicTransformOperator(ht.target,
                                           self._position_spaces[i][amp_space],
                                           space=spaces[i]) @ ht
        a = []
        for ii in range(n_amplitudes):
            co = ContractionOperator(hspace, spaces[:ii] + spaces[ii + 1:])
            pp = self._a[ii].target[amp_space]
Philipp Haim's avatar
Philipp Haim committed
486
            pd = PowerDistributor(co.target, pp, amp_space)
487 488 489
            a.append(co.adjoint @ pd @ self._a[ii])
        corr = reduce(mul, a)
        return ht(azm*corr*ducktape(hspace, None, self._prefix + 'xi'))
Philipp Arras's avatar
Philipp Arras committed
490

491
    def finalize(self, prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
492 493 494
        """
        offset vs zeromode: volume factor
        """
Philipp Frank's avatar
fixup  
Philipp Frank committed
495
        op = self._finalize_from_op()
496 497
        if self._offset_mean is not None:
            offset = self._offset_mean
498 499 500 501 502 503 504
            # 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
505
        self.statistics_summary(prior_info)
506 507
        return op

508 509 510 511 512 513
    def statistics_summary(self, prior_info):
        from ..sugar import from_random

        if prior_info == 0:
            return

514 515
        lst = [('Offset amplitude', self.amplitude_total_offset),
               ('Total fluctuation amplitude', self.total_fluctuation)]
516
        namps = len(self._a)
517 518 519 520 521 522 523 524
        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:
525 526 527
            sc = StatCalculator()
            for _ in range(prior_info):
                sc.add(op(from_random('normal', op.domain)))
Martin Reinecke's avatar
merge  
Martin Reinecke committed
528
            mean = sc.mean.val
Martin Reinecke's avatar
Martin Reinecke committed
529
            stddev = sc.var.ptw("sqrt").val
530
            for m, s in zip(mean.flatten(), stddev.flatten()):
531
                logger.info('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
532 533 534

    def moment_slice_to_average(self, fluctuations_slice_mean, nsamples=1000):
        fluctuations_slice_mean = float(fluctuations_slice_mean)
535 536 537
        if not fluctuations_slice_mean > 0:
            msg = "fluctuations_slice_mean must be greater zero; got {!r}"
            raise ValueError(msg.format(fluctuations_slice_mean))
538
        from ..sugar import from_random
539 540
        scm = 1.
        for a in self._a:
Martin Reinecke's avatar
Martin Reinecke committed
541
            op = a.fluctuation_amplitude*self._azm.ptw("reciprocal")
Martin Reinecke's avatar
stage 3  
Martin Reinecke committed
542
            res = np.array([op(from_random('normal', op.domain)).val
543 544
                            for _ in range(nsamples)])
            scm *= res**2 + 1.
545
        return fluctuations_slice_mean/np.mean(np.sqrt(scm))
546

Philipp Arras's avatar
Philipp Arras committed
547
    @property
Philipp Haim's avatar
Philipp Haim committed
548
    def normalized_amplitudes(self):
549
        return self._a
Philipp Arras's avatar
Philipp Arras committed
550

Philipp Haim's avatar
Philipp Haim committed
551 552 553 554 555 556 557
    @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
558 559
        dom = self._a[0].target
        expand = ContractionOperator(dom, len(dom)-1).adjoint
Philipp Haim's avatar
Philipp Haim committed
560 561
        return self._a[0]*(expand @ self.amplitude_total_offset)

562 563 564
    @property
    def amplitude_total_offset(self):
        return self._azm
Philipp Arras's avatar
Philipp Arras committed
565 566

    @property
567
    def total_fluctuation(self):
568
        """Returns operator which acts on prior or posterior samples"""
569
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
570
            raise NotImplementedError
571
        if len(self._a) == 1:
572
            return self.average_fluctuation(0)
573 574
        q = 1.
        for a in self._a:
Martin Reinecke's avatar
Martin Reinecke committed
575
            fl = a.fluctuation_amplitude*self._azm.ptw("reciprocal")
Philipp Arras's avatar
Philipp Arras committed
576
            q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Martin Reinecke's avatar
Martin Reinecke committed
577
        return (Adder(full(q.target, -1.)) @ q).ptw("sqrt")*self._azm
578

Philipp Arras's avatar
Philipp Arras committed
579
    def slice_fluctuation(self, space):
580
        """Returns operator which acts on prior or posterior samples"""
581
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
582
            raise NotImplementedError
583
        if space >= len(self._a):
584
            raise ValueError("invalid space specified; got {!r}".format(space))
585
        if len(self._a) == 1:
586
            return self.average_fluctuation(0)
587 588
        q = 1.
        for j in range(len(self._a)):
Martin Reinecke's avatar
Martin Reinecke committed
589
            fl = self._a[j].fluctuation_amplitude*self._azm.ptw("reciprocal")
590
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
591
                q = q*fl**2
592
            else:
Philipp Arras's avatar
Philipp Arras committed
593
                q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Martin Reinecke's avatar
Martin Reinecke committed
594
        return q.ptw("sqrt")*self._azm
Philipp Arras's avatar
Philipp Arras committed
595 596

    def average_fluctuation(self, space):
597
        """Returns operator which acts on prior or posterior samples"""
598
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
599
            raise NotImplementedError
600
        if space >= len(self._a):
601
            raise ValueError("invalid space specified; got {!r}".format(space))
602
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
603 604
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
605

606 607
    @staticmethod
    def offset_amplitude_realized(samples):
Philipp Haim's avatar
Philipp Haim committed
608
        spaces = _structured_spaces(samples[0].domain)
609 610
        res = 0.
        for s in samples:
Philipp Haim's avatar
Philipp Haim committed
611
            res = res + s.mean(spaces)**2
Philipp Haim's avatar
Philipp Haim committed
612 613
        res = res/len(samples)
        return np.sqrt(res if np.isscalar(res) else res.val)
Philipp Arras's avatar
Philipp Arras committed
614

615 616 617 618 619 620 621 622
    @staticmethod
    def total_fluctuation_realized(samples):
        return _total_fluctuation_realized(samples)

    @staticmethod
    def slice_fluctuation_realized(samples, space):
        """Computes slice fluctuations from collection of field (defined in signal
        space) realizations."""
Philipp Haim's avatar
Philipp Haim committed
623 624
        spaces = _structured_spaces(samples[0].domain)
        if space >= len(spaces):
625
            raise ValueError("invalid space specified; got {!r}".format(space))
Philipp Haim's avatar
Philipp Haim committed
626
        if len(spaces) == 1:
627
            return _total_fluctuation_realized(samples)
Philipp Haim's avatar
Philipp Haim committed
628
        space = space + spaces[0]
Philipp Arras's avatar
Philipp Arras committed
629
        res1, res2 = 0., 0.
630
        for s in samples:
Philipp Frank's avatar
fixes  
Philipp Frank committed
631 632 633 634
            res1 = res1 + s**2
            res2 = res2 + s.mean(space)**2
        res1 = res1/len(samples)
        res2 = res2/len(samples)
Philipp Haim's avatar
Philipp Haim committed
635
        res = res1.mean(spaces) - res2.mean(spaces[:-1])
Philipp Haim's avatar
Philipp Haim committed
636
        return np.sqrt(res if np.isscalar(res) else res.val)
Philipp Frank's avatar
fixes  
Philipp Frank committed
637

Philipp Arras's avatar
Philipp Arras committed
638
    @staticmethod
639 640 641
    def average_fluctuation_realized(samples, space):
        """Computes average fluctuations from collection of field (defined in signal
        space) realizations."""
Philipp Haim's avatar
Philipp Haim committed
642 643
        spaces = _structured_spaces(samples[0].domain)
        if space >= len(spaces):
644
            raise ValueError("invalid space specified; got {!r}".format(space))
Philipp Haim's avatar
Philipp Haim committed
645
        if len(spaces) == 1:
646
            return _total_fluctuation_realized(samples)
Philipp Haim's avatar
Philipp Haim committed
647 648 649
        space = space + spaces[0]
        sub_spaces = set(spaces)
        sub_spaces.remove(space)
Philipp Arras's avatar
Philipp Arras committed
650
        # Domain containing domain[space] and domain[0] iff total_N>0
Philipp Haim's avatar
Philipp Haim committed
651
        sub_dom = makeDomain([samples[0].domain[ind]
Philipp Arras's avatar
Philipp Arras committed
652
                              for ind in (set([0])-set(spaces)) | set([space])])
Philipp Haim's avatar
Philipp Haim committed
653
        co = ContractionOperator(sub_dom, len(sub_dom)-1)
654
        size = co.domain.size/co.target.size
Philipp Arras's avatar
Philipp Arras committed
655 656
        res = 0.
        for s in samples:
Philipp Haim's avatar
Philipp Haim committed
657
            r = s.mean(sub_spaces)
658
            res = res + (r - co.adjoint(co(r)/size))**2
Philipp Haim's avatar
Philipp Haim committed
659
        res = res.mean(spaces[0])/len(samples)
Philipp Haim's avatar
Philipp Haim committed
660
        return np.sqrt(res if np.isscalar(res) else res.val)