correlated_fields.py 30.2 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-2020 Max-Planck-Society
Philipp Arras's avatar
Philipp Arras committed
15
# Authors: Philipp Frank, Philipp Arras, Philipp Haim
Martin Reinecke's avatar
Martin Reinecke committed
16
#
17
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
Philipp Arras's avatar
Philipp Arras committed
18

19 20 21
from functools import reduce
from operator import mul

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

Philipp Arras's avatar
Philipp Arras committed
24
from ..domain_tuple import DomainTuple
Philipp Arras's avatar
Philipp Arras committed
25 26
from ..domains.power_space import PowerSpace
from ..domains.unstructured_domain import UnstructuredDomain
Philipp Arras's avatar
Philipp Arras committed
27
from ..field import Field
28
from ..logger import logger
Philipp Arras's avatar
Philipp Arras committed
29
from ..multi_field import MultiField
Philipp Arras's avatar
Philipp Arras committed
30
from ..operators.adder import Adder
31
from ..operators.contraction_operator import ContractionOperator
Philipp Arras's avatar
Philipp Arras committed
32
from ..operators.diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
33
from ..operators.distributors import PowerDistributor
Philipp Arras's avatar
Philipp Arras committed
34
from ..operators.endomorphic_operator import EndomorphicOperator
Martin Reinecke's avatar
Martin Reinecke committed
35
from ..operators.harmonic_operators import HarmonicTransformOperator
Philipp Arras's avatar
Philipp Arras committed
36
from ..operators.linear_operator import LinearOperator
Philipp Arras's avatar
Philipp Arras committed
37
from ..operators.operator import Operator
Philipp Arras's avatar
Philipp Arras committed
38
from ..operators.simple_linear_operators import ducktape
39
from ..probing import StatCalculator
Philipp Arras's avatar
Philipp Arras committed
40
from ..sugar import full, makeDomain, makeField, makeOp
41
from .. import utilities
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):
247 248 249
        self._dofdex = np.array(dofdex)
        self._target = DomainTuple.make(target)
        self._domain = DomainTuple.make(domain)
Philipp Haim's avatar
Philipp Haim committed
250 251 252 253
        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
254
        x = x.val
Philipp Haim's avatar
Philipp Haim committed
255 256 257
        if mode == self.TIMES:
            res = x[self._dofdex]
        else:
258
            res = np.zeros(self._tgt(mode).shape, dtype=x.dtype)
259
            res = utilities.special_add_at(res, 0, self._dofdex, x)
Martin Reinecke's avatar
Martin Reinecke committed
260
        return makeField(self._tgt(mode), res)
Martin Reinecke's avatar
Martin Reinecke committed
261

262

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

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

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

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

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

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

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

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

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

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

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

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

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

354 355

class CorrelatedFieldMaker:
Lukas Platz's avatar
Lukas Platz committed
356 357 358
    """Constrution helper for hirarchical correlated field models.

    The correlated field models are parametrized by creating
359 360
    power spectrum operators ("amplitudes") via calls to
    :func:`add_fluctuations` that act on the targeted field subdomains.
Lukas Platz's avatar
Lukas Platz committed
361
    During creation of the :class:`CorrelatedFieldMaker` via
362 363 364
    :func:`make`, a global offset from zero of the field model
    can be defined and an operator applying fluctuations
    around this offset is parametrized.
Lukas Platz's avatar
Lukas Platz committed
365 366

    The resulting correlated field model operator has a
Martin Reinecke's avatar
Martin Reinecke committed
367
    :class:`~nifty7.multi_domain.MultiDomain` as its domain and
Lukas Platz's avatar
Lukas Platz committed
368 369 370
    expects its input values to be univariately gaussian.

    The target of the constructed operator will be a
Martin Reinecke's avatar
merge  
Martin Reinecke committed
371
    :class:`~nifty7.domain_tuple.DomainTuple` containing the
372 373
    `target_subdomains` of the added fluctuations in the order of
    the `add_fluctuations` calls.
Lukas Platz's avatar
Lukas Platz committed
374

375
    Creation of the model operator is completed by calling the method
Lukas Platz's avatar
Lukas Platz committed
376 377
    :func:`finalize`, which returns the configured operator.

378 379 380 381 382 383 384 385 386 387
    An operator representing an array of correlated field models
    can be constructed by setting the `total_N` parameter of
    :func:`make`. It will have an
    :class:`~nifty.domains.unstructucture_domain.UnstructureDomain`
    of shape `(total_N,)` prepended to its target domain and represent
    `total_N` correlated fields simulataneously.
    The degree of information sharing between the correlated field
    models can be configured via the `dofdex` parameters
    of :func:`make` and :func:`add_fluctuations`.

Lukas Platz's avatar
Lukas Platz committed
388
    See the methods :func:`make`, :func:`add_fluctuations`
389
    and :func:`finalize` for further usage information."""
390 391 392
    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")
393
        self._a = []
394
        self._target_subdomains = []
Philipp Arras's avatar
Formats  
Philipp Arras committed
395

396 397
        self._offset_mean = offset_mean
        self._azm = offset_fluctuations_op
398
        self._prefix = prefix
Philipp Haim's avatar
Philipp Haim committed
399
        self._total_N = total_N
Philipp Arras's avatar
Formats  
Philipp Arras committed
400

401
    @staticmethod
Philipp Arras's avatar
Philipp Arras committed
402
    def make(offset_mean, offset_std_mean, offset_std_std, prefix, total_N=0,
Martin Reinecke's avatar
Martin Reinecke committed
403
             dofdex=None):
Lukas Platz's avatar
Lukas Platz committed
404 405 406 407 408 409 410
        """Returns a CorrelatedFieldMaker object.

        Parameters
        ----------
        offset_mean : float
            Mean offset from zero of the correlated field to be made.
        offset_std_mean : float
Lukas Platz's avatar
Lukas Platz committed
411
            Mean standard deviation of the offset.
Lukas Platz's avatar
Lukas Platz committed
412
        offset_std_std : float
Lukas Platz's avatar
Lukas Platz committed
413
            Standard deviation of the stddev of the offset.
Lukas Platz's avatar
Lukas Platz committed
414 415
        prefix : string
            Prefix to the names of the domains of the cf operator to be made.
Lukas Platz's avatar
Lukas Platz committed
416
            This determines the names of the operator domain.
417 418
        total_N : integer, optional
            Number of field models to create.
Lukas Platz's avatar
Lukas Platz committed
419 420 421
            If not 0, the first entry of the operators target will be an
            :class:`~nifty.domains.unstructured_domain.UnstructuredDomain`
            with length `total_N`.
422
        dofdex : np.array of integers, optional
Philipp Arras's avatar
Philipp Arras committed
423 424 425
            An integer array specifying the zero mode models used if
            total_N > 1. It needs to have length of total_N. If total_N=3 and
            dofdex=[0,0,1], that means that two models for the zero mode are
426 427 428 429
            instanciated; the first one is used for the first and second
            field model and the second is used for the third field model.
            *If not specified*, use the same zero mode model for all
            constructed field models.
Lukas Platz's avatar
Lukas Platz committed
430
        """
Philipp Frank's avatar
Philipp Frank committed
431 432
        if dofdex is None:
            dofdex = np.full(total_N, 0)
433 434
        elif len(dofdex) != total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Frank's avatar
Philipp Frank committed
435
        N = max(dofdex) + 1 if total_N > 0 else 0
Lukas Platz's avatar
Lukas Platz committed
436 437
        zm = _LognormalMomentMatching(offset_std_mean,
                                      offset_std_std,
Philipp Haim's avatar
Philipp Haim committed
438
                                      prefix + 'zeromode',
Philipp Frank's avatar
Philipp Frank committed
439
                                      N)
Philipp Frank's avatar
fixup  
Philipp Frank committed
440
        if total_N > 0:
Martin Reinecke's avatar
Martin Reinecke committed
441
            zm = _Distributor(dofdex, zm.target, UnstructuredDomain(total_N)) @ zm
442
        return CorrelatedFieldMaker(offset_mean, zm, prefix, total_N)
443 444

    def add_fluctuations(self,
445
                         target_subdomain,
446 447 448 449 450 451 452 453
                         fluctuations_mean,
                         fluctuations_stddev,
                         flexibility_mean,
                         flexibility_stddev,
                         asperity_mean,
                         asperity_stddev,
                         loglogavgslope_mean,
                         loglogavgslope_stddev,
Martin Reinecke's avatar
Martin Reinecke committed
454 455 456 457
                         prefix='',
                         index=None,
                         dofdex=None,
                         harmonic_partner=None):
Lukas Platz's avatar
Lukas Platz committed
458 459 460 461 462 463
        """Function to add correlation structures to the field to be made.

        Correlations are described by their power spectrum and the subdomain
        on which they apply.

        The parameters `fluctuations`, `flexibility`, `asperity` and
464 465
        `loglogavgslope` configure the power spectrum model ("amplitude")
        used on the target field subdomain `target_subdomain`.
Lukas Platz's avatar
Lukas Platz committed
466 467 468 469 470 471 472 473 474 475 476
        It is assembled as the sum of a power law component
        (linear slope in log-log power-frequency-space),
        a smooth varying component (integrated wiener process) and
        a ragged componenent (unintegrated wiener process).

        Multiple calls to `add_fluctuations` are possible, in which case
        the constructed field will have the outer product of the individual
        power spectra as its global power spectrum.

        Parameters
        ----------
Martin Reinecke's avatar
Martin Reinecke committed
477 478
        target_subdomain : :class:`~nifty7.domain.Domain`, \
                           :class:`~nifty7.domain_tuple.DomainTuple`
Lukas Platz's avatar
Lukas Platz committed
479 480 481 482 483
            Target subdomain on which the correlation structure defined
            in this call should hold.
        fluctuations_{mean,stddev} : float
            Total spectral energy -> Amplitude of the fluctuations
        flexibility_{mean,stddev} : float
484
            Amplitude of the non-power-law power spectrum component
Lukas Platz's avatar
Lukas Platz committed
485
        asperity_{mean,stddev} : float
486
            Roughness of the non-power-law power spectrum component
Lukas Platz's avatar
Lukas Platz committed
487 488 489 490 491
            Used to accomodate single frequency peaks
        loglogavgslope_{mean,stddev} : float
            Power law component exponent
        prefix : string
            prefix of the power spectrum parameter domain names
Philipp Arras's avatar
Philipp Arras committed
492 493 494
        index : int
            Position target_subdomain in the final total domain of the
            correlated field operator.
495 496
        dofdex : np.array, optional
            An integer array specifying the power spectrum models used if
Philipp Arras's avatar
Philipp Arras committed
497
            total_N > 1. It needs to have length of total_N. If total_N=3 and
498 499 500 501 502
            dofdex=[0,0,1], that means that two power spectrum models are
            instanciated; the first one is used for the first and second
            field model and the second one is used for the third field model.
            *If not given*, use the same power spectrum model for all
            constructed field models.
Martin Reinecke's avatar
Martin Reinecke committed
503 504
        harmonic_partner : :class:`~nifty7.domain.Domain`, \
                           :class:`~nifty7.domain_tuple.DomainTuple`
Lukas Platz's avatar
Lukas Platz committed
505 506
            In which harmonic space to define the power spectrum
        """
Philipp Frank's avatar
Philipp Frank committed
507
        if harmonic_partner is None:
508
            harmonic_partner = target_subdomain.get_default_codomain()
Philipp Frank's avatar
Fixup  
Philipp Frank committed
509
        else:
510 511
            target_subdomain.check_codomain(harmonic_partner)
            harmonic_partner.check_codomain(target_subdomain)
512

Philipp Haim's avatar
Philipp Haim committed
513 514
        if dofdex is None:
            dofdex = np.full(self._total_N, 0)
515 516
        elif len(dofdex) != self._total_N:
            raise ValueError("length of dofdex needs to match total_N")
Philipp Haim's avatar
Philipp Haim committed
517

Philipp Haim's avatar
Philipp Haim committed
518 519
        if self._total_N > 0:
            N = max(dofdex) + 1
520
            target_subdomain = makeDomain((UnstructuredDomain(N), target_subdomain))
Philipp Haim's avatar
Philipp Haim committed
521
        else:
Philipp Haim's avatar
Philipp Haim committed
522
            N = 0
523
            target_subdomain = makeDomain(target_subdomain)
Philipp Arras's avatar
Philipp Arras committed
524
        prefix = str(prefix)
525
        # assert isinstance(target_subdomain[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
526

Philipp Arras's avatar
Philipp Arras committed
527 528
        fluct = _LognormalMomentMatching(fluctuations_mean,
                                         fluctuations_stddev,
529
                                         self._prefix + prefix + 'fluctuations',
Philipp Haim's avatar
Philipp Haim committed
530
                                         N)
Philipp Arras's avatar
Philipp Arras committed
531
        flex = _LognormalMomentMatching(flexibility_mean, flexibility_stddev,
532
                                        self._prefix + prefix + 'flexibility',
Philipp Haim's avatar
Philipp Haim committed
533
                                        N)
Philipp Arras's avatar
Philipp Arras committed
534
        asp = _LognormalMomentMatching(asperity_mean, asperity_stddev,
Philipp Arras's avatar
Philipp Arras committed
535
                                       self._prefix + prefix + 'asperity', N)
536
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
537
                        self._prefix + prefix + 'loglogavgslope', N)
Philipp Arras's avatar
Philipp Arras committed
538
        amp = _Amplitude(PowerSpace(harmonic_partner), fluct, flex, asp, avgsl,
539
                         self._azm, target_subdomain[-1].total_volume,
540
                         self._prefix + prefix + 'spectrum', dofdex)
Philipp Haim's avatar
Philipp Haim committed
541

542 543
        if index is not None:
            self._a.insert(index, amp)
544
            self._target_subdomains.insert(index, target_subdomain)
545 546
        else:
            self._a.append(amp)
547
            self._target_subdomains.append(target_subdomain)
548

Philipp Arras's avatar
Philipp Arras committed
549 550 551 552 553 554 555 556 557 558
    def finalize(self, prior_info=100):
        """Finishes model construction process and returns the constructed
        operator.

        Parameters
        ----------
        prior_info : integer
            How many prior samples to draw for property verification statistics
            If zero, skips calculating and displaying statistics.
        """
Philipp Haim's avatar
Philipp Haim committed
559
        n_amplitudes = len(self._a)
Philipp Haim's avatar
Philipp Haim committed
560
        if self._total_N > 0:
Philipp Arras's avatar
Philipp Arras committed
561 562 563
            hspace = makeDomain(
                [UnstructuredDomain(self._total_N)] +
                [dd.target[-1].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
564 565
            spaces = tuple(range(1, n_amplitudes + 1))
            amp_space = 1
Philipp Haim's avatar
Philipp Haim committed
566 567
        else:
            hspace = makeDomain(
Philipp Arras's avatar
Philipp Arras committed
568
                [dd.target[0].harmonic_partner for dd in self._a])
Philipp Haim's avatar
Philipp Haim committed
569
            spaces = tuple(range(n_amplitudes))
Philipp Haim's avatar
Philipp Haim committed
570
            amp_space = 0
571

Martin Reinecke's avatar
Martin Reinecke committed
572
        expander = ContractionOperator(hspace, spaces=spaces).adjoint
Philipp Frank's avatar
fixup  
Philipp Frank committed
573
        azm = expander @ self._azm
574

575
        ht = HarmonicTransformOperator(hspace,
576
                                       self._target_subdomains[0][amp_space],
Martin Reinecke's avatar
Martin Reinecke committed
577
                                       space=spaces[0])
578
        for i in range(1, n_amplitudes):
579
            ht = HarmonicTransformOperator(ht.target,
580
                                           self._target_subdomains[i][amp_space],
581 582 583 584 585
                                           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
586
            pd = PowerDistributor(co.target, pp, amp_space)
587 588
            a.append(co.adjoint @ pd @ self._a[ii])
        corr = reduce(mul, a)
Philipp Arras's avatar
Philipp Arras committed
589
        op = ht(azm*corr*ducktape(hspace, None, self._prefix + 'xi'))
Philipp Arras's avatar
Philipp Arras committed
590

591 592
        if self._offset_mean is not None:
            offset = self._offset_mean
593 594 595 596 597 598 599
            # 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
600
        self.statistics_summary(prior_info)
601 602
        return op

603 604 605 606 607 608
    def statistics_summary(self, prior_info):
        from ..sugar import from_random

        if prior_info == 0:
            return

609 610
        lst = [('Offset amplitude', self.amplitude_total_offset),
               ('Total fluctuation amplitude', self.total_fluctuation)]
611
        namps = len(self._a)
612 613 614 615 616 617 618 619
        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:
620 621
            sc = StatCalculator()
            for _ in range(prior_info):
622
                sc.add(op(from_random(op.domain, 'normal')))
Martin Reinecke's avatar
merge  
Martin Reinecke committed
623
            mean = sc.mean.val
Martin Reinecke's avatar
Martin Reinecke committed
624
            stddev = sc.var.ptw("sqrt").val
625
            for m, s in zip(mean.flatten(), stddev.flatten()):
626
                logger.info('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
627 628 629

    def moment_slice_to_average(self, fluctuations_slice_mean, nsamples=1000):
        fluctuations_slice_mean = float(fluctuations_slice_mean)
630 631 632
        if not fluctuations_slice_mean > 0:
            msg = "fluctuations_slice_mean must be greater zero; got {!r}"
            raise ValueError(msg.format(fluctuations_slice_mean))
633
        from ..sugar import from_random
634 635
        scm = 1.
        for a in self._a:
Martin Reinecke's avatar
Martin Reinecke committed
636
            op = a.fluctuation_amplitude*self._azm.ptw("reciprocal")
637
            res = np.array([op(from_random(op.domain, 'normal')).val
638 639
                            for _ in range(nsamples)])
            scm *= res**2 + 1.
640
        return fluctuations_slice_mean/np.mean(np.sqrt(scm))
641

Philipp Arras's avatar
Philipp Arras committed
642
    @property
Philipp Haim's avatar
Philipp Haim committed
643
    def normalized_amplitudes(self):
644
        """Returns the power spectrum operators used in the model"""
645
        return self._a
Philipp Arras's avatar
Philipp Arras committed
646

Philipp Haim's avatar
Philipp Haim committed
647 648 649 650 651 652 653
    @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
654 655
        dom = self._a[0].target
        expand = ContractionOperator(dom, len(dom)-1).adjoint
Philipp Haim's avatar
Philipp Haim committed
656 657
        return self._a[0]*(expand @ self.amplitude_total_offset)

658 659 660
    @property
    def amplitude_total_offset(self):
        return self._azm
Philipp Arras's avatar
Philipp Arras committed
661 662

    @property
663
    def total_fluctuation(self):
664
        """Returns operator which acts on prior or posterior samples"""
665
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
666
            raise NotImplementedError
667
        if len(self._a) == 1:
668
            return self.average_fluctuation(0)
669 670
        q = 1.
        for a in self._a:
Martin Reinecke's avatar
Martin Reinecke committed
671
            fl = a.fluctuation_amplitude*self._azm.ptw("reciprocal")
Philipp Arras's avatar
Philipp Arras committed
672
            q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Martin Reinecke's avatar
Martin Reinecke committed
673
        return (Adder(full(q.target, -1.)) @ q).ptw("sqrt")*self._azm
674

Philipp Arras's avatar
Philipp Arras committed
675
    def slice_fluctuation(self, space):
676
        """Returns operator which acts on prior or posterior samples"""
677
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
678
            raise NotImplementedError
679
        if space >= len(self._a):
680
            raise ValueError("invalid space specified; got {!r}".format(space))
681
        if len(self._a) == 1:
682
            return self.average_fluctuation(0)
683 684
        q = 1.
        for j in range(len(self._a)):
Martin Reinecke's avatar
Martin Reinecke committed
685
            fl = self._a[j].fluctuation_amplitude*self._azm.ptw("reciprocal")
686
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
687
                q = q*fl**2
688
            else:
Philipp Arras's avatar
Philipp Arras committed
689
                q = q*(Adder(full(fl.target, 1.)) @ fl**2)
Martin Reinecke's avatar
Martin Reinecke committed
690
        return q.ptw("sqrt")*self._azm
Philipp Arras's avatar
Philipp Arras committed
691 692

    def average_fluctuation(self, space):
693
        """Returns operator which acts on prior or posterior samples"""
694
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
695
            raise NotImplementedError
696
        if space >= len(self._a):
697
            raise ValueError("invalid space specified; got {!r}".format(space))
698
        if len(self._a) == 1:
Philipp Haim's avatar
Philipp Haim committed
699 700
            return self._a[0].fluctuation_amplitude
        return self._a[space].fluctuation_amplitude
701

702 703
    @staticmethod
    def offset_amplitude_realized(samples):
Philipp Haim's avatar
Philipp Haim committed
704
        spaces = _structured_spaces(samples[0].domain)
705 706
        res = 0.
        for s in samples:
Philipp Haim's avatar
Philipp Haim committed
707
            res = res + s.mean(spaces)**2
Philipp Haim's avatar
Philipp Haim committed
708 709
        res = res/len(samples)
        return np.sqrt(res if np.isscalar(res) else res.val)
Philipp Arras's avatar
Philipp Arras committed
710

711 712 713 714 715 716 717 718
    @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
719 720
        spaces = _structured_spaces(samples[0].domain)
        if space >= len(spaces):
721
            raise ValueError("invalid space specified; got {!r}".format(space))
Philipp Haim's avatar
Philipp Haim committed
722
        if len(spaces) == 1:
723
            return _total_fluctuation_realized(samples)
Philipp Haim's avatar
Philipp Haim committed
724
        space = space + spaces[0]
Philipp Arras's avatar
Philipp Arras committed
725
        res1, res2 = 0., 0.
726
        for s in samples:
Philipp Frank's avatar
fixes  
Philipp Frank committed
727 728 729 730
            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
731
        res = res1.mean(spaces) - res2.mean(spaces[:-1])
Philipp Haim's avatar
Philipp Haim committed
732
        return np.sqrt(res if np.isscalar(res) else res.val)
Philipp Frank's avatar
fixes  
Philipp Frank committed
733

Philipp Arras's avatar
Philipp Arras committed
734
    @staticmethod
735 736 737
    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
738 739
        spaces = _structured_spaces(samples[0].domain)
        if space >= len(spaces):
740
            raise ValueError("invalid space specified; got {!r}".format(space))
Philipp Haim's avatar
Philipp Haim committed
741
        if len(spaces) == 1:
742
            return _total_fluctuation_realized(samples)
Philipp Haim's avatar
Philipp Haim committed
743 744 745
        space = space + spaces[0]
        sub_spaces = set(spaces)
        sub_spaces.remove(space)
Philipp Arras's avatar
Philipp Arras committed
746
        # Domain containing domain[space] and domain[0] iff total_N>0
Philipp Haim's avatar
Philipp Haim committed
747
        sub_dom = makeDomain([samples[0].domain[ind]
Philipp Arras's avatar
Philipp Arras committed
748
                              for ind in (set([0])-set(spaces)) | set([space])])
Philipp Haim's avatar
Philipp Haim committed
749
        co = ContractionOperator(sub_dom, len(sub_dom)-1)
750
        size = co.domain.size/co.target.size
Philipp Arras's avatar
Philipp Arras committed
751 752
        res = 0.
        for s in samples:
Philipp Haim's avatar
Philipp Haim committed
753
            r = s.mean(sub_spaces)
754
            res = res + (r - co.adjoint(co(r)/size))**2
Philipp Haim's avatar
Philipp Haim committed
755
        res = res.mean(spaces[0])/len(samples)
Philipp Haim's avatar
Philipp Haim committed
756
        return np.sqrt(res if np.isscalar(res) else res.val)