There is a maintenance of MPCDF Gitlab on Thursday, April 22st 2020, 9:00 am CEST - Expect some service interruptions during this time

correlated_fields.py 23.3 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
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
24
from ..field import Field
Philipp Arras's avatar
Philipp Arras committed
25
from ..operators.adder import Adder
26
from ..operators.contraction_operator import ContractionOperator
Philipp Arras's avatar
Philipp Arras committed
27
from ..operators.distributors import PowerDistributor
Philipp Arras's avatar
Philipp Arras committed
28
from ..operators.endomorphic_operator import EndomorphicOperator
Martin Reinecke's avatar
Martin Reinecke committed
29
from ..operators.harmonic_operators import HarmonicTransformOperator
Philipp Arras's avatar
Philipp Arras committed
30
from ..operators.linear_operator import LinearOperator
31
from ..operators.diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
32 33
from ..operators.operator import Operator
from ..operators.simple_linear_operators import VdotOperator, ducktape
Philipp Arras's avatar
Philipp Arras committed
34
from ..operators.value_inserter import ValueInserter
35
from ..probing import StatCalculator
Philipp Haim's avatar
Philipp Haim committed
36
from ..sugar import from_global_data, from_random, full, makeDomain, get_default_codomain
37

38

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

48

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


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


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


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


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


97
def _total_fluctuation_realized(samples, space = 0):
98 99
    res = 0.
    for s in samples:
100 101
        res = res + (s - s.mean(space, keepdims = True))**2
    return np.sqrt((res/len(samples)).mean(space))
102 103 104 105 106 107 108 109 110


def _stats(op, samples):
    sc = StatCalculator()
    for s in samples:
        sc.add(op(s.extract(op.domain)))
    return sc.mean.to_global_data(), sc.var.sqrt().to_global_data()


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


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

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

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

Philipp Arras's avatar
Philipp Arras committed
154 155

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

    def apply(self, x, mode):
        self._check_input(x, mode)
169 170 171 172

        #Maybe make class properties
        axis = self._target.axes[self._space][0]
        sl = (slice(None),)*axis
Philipp Haim's avatar
Fixes  
Philipp Haim committed
173
        extender_sl = (None,)*axis + (slice(None),) + (None,)*(self._target.axes[-1][-1] - axis)
174 175 176 177 178
        first = sl + (0,)
        second = sl + (1,)
        from_third = sl + (slice(2,None),)
        no_border = sl + (slice(1,-1),)
        reverse = sl + (slice(None,None,-1),)
179 180

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


class _Normalization(Operator):
198
    def __init__(self, domain, space = 0):
Philipp Arras's avatar
Philipp Arras committed
199
        self._domain = self._target = makeDomain(domain)
200
        assert isinstance(self._domain[space], PowerSpace)
201 202 203 204
        hspace = list(self._domain)
        hspace[space] = hspace[space].harmonic_partner
        hspace = makeDomain(hspace)
        pd = PowerDistributor(hspace, power_space=self._domain[space], space = space)
Philipp Arras's avatar
Philipp Arras committed
205
        # TODO Does not work on sphere yet
206
        mode_multiplicity = pd.adjoint(full(pd.target, 1.)).to_global_data_rw()
207
        zero_mode = (slice(None),)*self._domain.axes[space][0] + (0,)
Philipp Haim's avatar
Philipp Haim committed
208
        mode_multiplicity[zero_mode] = 0
209 210
        self._mode_multiplicity = from_global_data(self._domain, mode_multiplicity)
        self._specsum = _SpecialSum(self._domain, space)
Philipp Arras's avatar
Philipp Arras committed
211 212 213 214 215 216 217

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


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

    def apply(self, x, mode):
        self._check_input(x, mode)
229
        return self._contractor.adjoint(self._contractor(x))
Philipp Arras's avatar
Philipp Arras committed
230 231


232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
class _slice_extractor(LinearOperator):
    #FIXME it should be tested if the the domain and target are consistent with the slice
    def __init__(self, domain, target, sl):
        self._domain = makeDomain(domain)
        self._target = makeDomain(target)
        self._sl = sl
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
        x = x.to_global_data()
        if mode == self.TIMES:
            res = x[self._sl]
            res = res.reshape(self._target.shape)
        else:
            res = np.zeros(self._domain.shape)
            res[self._sl] = x
        return from_global_data(self._tgt(mode), res)
Philipp Haim's avatar
Philipp Haim committed
250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270


class _Distributor(LinearOperator):
    def __init__(self, dofdex, domain, target, space = 0):
        self._dofdex = dofdex

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

    def apply(self, x, mode):
        self._check_input(x, mode)
        x = x.to_global_data()
        if mode == self.TIMES:
            res = x[self._dofdex]
        else:
            res = np.empty(self._tgt(mode).shape)
            res[self._dofdex] = x
        return from_global_data(self._tgt(mode), res)

271 272
    

273 274
class _Amplitude(Operator):
    def __init__(self, target, fluctuations, flexibility, asperity,
Philipp Haim's avatar
Philipp Haim committed
275
                 loglogavgslope, azm, key, dofdex):
Philipp Arras's avatar
Philipp Arras committed
276 277 278 279 280 281 282 283 284 285 286
        """
        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
287 288
        if len(dofdex) > 0:
            N_copies = max(dofdex) + 1
Philipp Haim's avatar
Philipp Haim committed
289 290 291 292 293
            space = 1
            distributed_tgt = makeDomain((UnstructuredDomain(len(dofdex)), target))
            target = makeDomain((UnstructuredDomain(N_copies), target))
            Distributor = _Distributor(dofdex, target, distributed_tgt, 0)
        else:
Philipp Haim's avatar
Philipp Haim committed
294
            N_copies = 0
Philipp Haim's avatar
Philipp Haim committed
295
            space = 0
Philipp Haim's avatar
Philipp Haim committed
296 297
            distributed_tgt = target = makeDomain(target)
        azm_expander = ContractionOperator(distributed_tgt, spaces = space).adjoint
Philipp Haim's avatar
Philipp Haim committed
298 299
        assert isinstance(target[space], PowerSpace)
        
300
        twolog = _TwoLogIntegrations(target, space)
Philipp Arras's avatar
Philipp Arras committed
301
        dom = twolog.domain
302 303 304 305
        shp = dom[space].shape
        totvol = target[space].harmonic_partner.get_default_codomain().total_volume
        expander = ContractionOperator(dom, spaces = space).adjoint
        ps_expander = ContractionOperator(twolog.target, spaces = space).adjoint
Philipp Arras's avatar
Philipp Arras committed
306 307 308

        # Prepare constant fields
        foo = np.zeros(shp)
309 310
        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
311 312 313

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

        foo = np.ones(shp)
317 318
        foo[0] = _log_vol(target[space])**2/12.
        shift = DiagonalOperator(from_global_data(dom[space], foo), dom, space)
Philipp Frank's avatar
Philipp Frank committed
319
        
320 321 322 323 324
        vslope = DiagonalOperator(
                    from_global_data(target[space], _relative_log_k_lengths(target[space])),
                    target, space)

        foo, bar = [np.zeros(target[space].shape) for _ in range(2)]
Philipp Arras's avatar
Philipp Arras committed
325
        bar[1:] = foo[0] = totvol
326 327 328 329 330 331 332
        vol0, vol1 = [DiagonalOperator(from_global_data(target[space], aa), 
                target, space) for aa in (foo, bar)]

        #Prepare fields for Adder
        #NOTE alternative would be adjoint contraction_operator acting
        #on every space except the specified on
        shift, vol0 = [op(full(op.domain, 1)) for op in (shift, vol0)]
Philipp Arras's avatar
Philipp Arras committed
333 334
        # End prepare constant fields

335 336 337 338
        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
339
        sig_fluc = vol1 @ ps_expander @ fluctuations
Philipp Arras's avatar
Philipp Arras committed
340 341

        xi = ducktape(dom, None, key)
Philipp Arras's avatar
Philipp Arras committed
342
        sigma = sig_flex*(Adder(shift) @ sig_asp).sqrt()
343 344
        smooth = _SlopeRemover(target, space) @ twolog @ (sigma*xi)
        op = _Normalization(target, space) @ (slope + smooth)
Philipp Haim's avatar
Philipp Haim committed
345
        if N_copies > 0:
Philipp Haim's avatar
Philipp Haim committed
346 347
            op = Distributor @ op
            sig_fluc = Distributor @ sig_fluc
Philipp Haim's avatar
Philipp Haim committed
348
            op = Adder(Distributor(vol0)) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Philipp Haim's avatar
Philipp Haim committed
349
        else:
Philipp Haim's avatar
Philipp Haim committed
350
            op = (Adder(vol0)) @ (sig_fluc*(azm_expander @ azm.one_over())*op)
Philipp Arras's avatar
Philipp Arras committed
351

Philipp Arras's avatar
Philipp Arras committed
352
        self.apply = op.apply
Philipp Arras's avatar
Philipp Arras committed
353
        self._fluc = fluctuations
Philipp Arras's avatar
Philipp Arras committed
354
        self._domain, self._target = op.domain, op.target
355
        self._space = space
Philipp Arras's avatar
Philipp Arras committed
356

Philipp Arras's avatar
Philipp Arras committed
357 358 359 360
    @property
    def fluctuation_amplitude(self):
        return self._fluc

361 362

class CorrelatedFieldMaker:
Philipp Haim's avatar
Philipp Haim committed
363
    def __init__(self, amplitude_offset, prefix, total_N):
364
        self._a = []
365
        self._spaces = []
366
        self._position_spaces = []
367 368 369
        
        self._azm = amplitude_offset
        self._prefix = prefix
Philipp Haim's avatar
Philipp Haim committed
370
        self._total_N = total_N
371 372
    
    @staticmethod
Philipp Haim's avatar
Philipp Haim committed
373
    def make(offset_amplitude_mean, offset_amplitude_stddev, prefix, total_N = 0):
374 375 376 377 378 379
        offset_amplitude_stddev = float(offset_amplitude_stddev)
        offset_amplitude_mean = float(offset_amplitude_mean)
        assert offset_amplitude_stddev > 0
        assert offset_amplitude_mean > 0
        zm = _LognormalMomentMatching(offset_amplitude_mean,
                                      offset_amplitude_stddev,
Philipp Haim's avatar
Philipp Haim committed
380 381 382
                                      prefix + 'zeromode',
                                      total_N)
        return CorrelatedFieldMaker(zm, prefix, total_N)
383 384

    def add_fluctuations(self,
385
                         position_space,
386 387 388 389 390 391 392 393
                         fluctuations_mean,
                         fluctuations_stddev,
                         flexibility_mean,
                         flexibility_stddev,
                         asperity_mean,
                         asperity_stddev,
                         loglogavgslope_mean,
                         loglogavgslope_stddev,
Philipp Haim's avatar
Philipp Haim committed
394 395 396 397 398 399 400 401
                         prefix = '',
                         index = None,
                         dofdex = None):
        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
402
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
403
            space = 1
Philipp Haim's avatar
Philipp Haim committed
404 405
            N = max(dofdex) + 1
            position_space = makeDomain((UnstructuredDomain(N), position_space))
Philipp Haim's avatar
Philipp Haim committed
406 407
        else:
            space = 0
Philipp Haim's avatar
Philipp Haim committed
408
            N = 0
Philipp Haim's avatar
Philipp Haim committed
409 410
            position_space = makeDomain(position_space)
        power_space = PowerSpace(position_space[space].get_default_codomain())
Philipp Arras's avatar
Philipp Arras committed
411
        prefix = str(prefix)
412
        #assert isinstance(position_space[space], (RGSpace, HPSpace, GLSpace)
Philipp Arras's avatar
Philipp Arras committed
413

Philipp Arras's avatar
Philipp Arras committed
414 415
        fluct = _LognormalMomentMatching(fluctuations_mean,
                                         fluctuations_stddev,
416
                                         prefix + 'fluctuations',
Philipp Haim's avatar
Philipp Haim committed
417
                                         N)
Philipp Arras's avatar
Philipp Arras committed
418
        flex = _LognormalMomentMatching(flexibility_mean, flexibility_stddev,
419
                                        prefix + 'flexibility',
Philipp Haim's avatar
Philipp Haim committed
420
                                        N)
Philipp Arras's avatar
Philipp Arras committed
421
        asp = _LognormalMomentMatching(asperity_mean, asperity_stddev,
422
                                       prefix + 'asperity', 
Philipp Haim's avatar
Philipp Haim committed
423
                                       N)
424
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
Philipp Haim's avatar
Philipp Haim committed
425
                        prefix + 'loglogavgslope', N)
426
        amp = _Amplitude(power_space,
Philipp Haim's avatar
Philipp Haim committed
427 428
                         fluct, flex, asp, avgsl, self._azm, prefix + 'spectrum', dofdex)

429 430
        if index is not None:
            self._a.insert(index, amp)
431
            self._position_spaces.insert(index, position_space)
432
            self._spaces.insert(index, space)
433 434
        else:
            self._a.append(amp)
435
            self._position_spaces.append(position_space)
Philipp Haim's avatar
Typo  
Philipp Haim committed
436
            self._spaces.append(space)
437

Philipp Haim's avatar
Philipp Haim committed
438
    def finalize_from_op(self, zeromode, prefix=''):
439
        assert isinstance(zeromode, Operator)
440
        self._azm = zeromode
Philipp Haim's avatar
Philipp Haim committed
441
        n_amplitudes = len(self._a)
Philipp Haim's avatar
Philipp Haim committed
442
        if self._total_N > 0:
Philipp Haim's avatar
Philipp Haim committed
443 444
            hspace = makeDomain([UnstructuredDomain(self._total_N)] +
                    [dd[-1].get_default_codomain() for dd in self._position_spaces])
Philipp Haim's avatar
Philipp Haim committed
445 446 447 448
            spaces = list(1 + np.arange(n_amplitudes))
            #spaces = tuple(len(dd) for dd in self._position_spaces)
            #spaces = 1 + np.cumsum(spaces)
            zeroind = (slice(None),) + (0,)*(len(hspace.shape)-1)
Philipp Haim's avatar
Philipp Haim committed
449 450 451 452
        else:
            hspace = makeDomain(
                    [dd[-1].get_default_codomain() for dd in self._position_spaces])
            spaces = tuple(range(n_amplitudes))
Philipp Haim's avatar
Philipp Haim committed
453 454
            spaces = list(np.arange(n_amplitudes))
            zeroind = (0,)*len(hspace.shape)
Philipp Haim's avatar
Philipp Haim committed
455

456 457
        foo = np.ones(hspace.shape)
        foo[zeroind] = 0
458

459 460 461
        ZeroModeInserter = _slice_extractor(hspace, 
                     self._azm.target, zeroind).adjoint
        azm = Adder(from_global_data(hspace, foo)) @ ZeroModeInserter @ zeromode
462

Philipp Haim's avatar
Philipp Haim committed
463
        #spaces = np.array(range(n_amplitudes)) + 1 - 1//self._total_N
464 465
        ht = HarmonicTransformOperator(hspace,
                                   self._position_spaces[0][self._spaces[0]],
Philipp Haim's avatar
Philipp Haim committed
466
                                   space=spaces[0])
467
        for i in range(1, n_amplitudes):
468
            ht = (HarmonicTransformOperator(ht.target,
469
                                    self._position_spaces[i][self._spaces[i]],
Philipp Haim's avatar
Philipp Haim committed
470
                                    space=spaces[i]) @ ht)
471

472
        pd = PowerDistributor(hspace, self._a[0].target[self._spaces[0]], self._spaces[0])
473
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
474
            pd = (pd @ PowerDistributor(pd.domain,
475
                                   self._a[i].target[self._spaces[i]],
Philipp Haim's avatar
Philipp Haim committed
476
                                   space=spaces[i]))
Philipp Arras's avatar
Philipp Arras committed
477

478 479
        a = ContractionOperator(pd.domain, spaces[1:]).adjoint @ self._a[0]
        for i in range(1, n_amplitudes):
Philipp Haim's avatar
Philipp Haim committed
480
            co = ContractionOperator(pd.domain,
Philipp Haim's avatar
Philipp Haim committed
481
                    spaces[:i] + spaces[i+1:])
482
            a = a*(co.adjoint @ self._a[i])
Philipp Arras's avatar
Philipp Arras committed
483

484
        return ht(azm*(pd @ a)*ducktape(hspace, None, prefix + 'xi'))
Philipp Arras's avatar
Philipp Arras committed
485 486

    def finalize(self,
487 488
                 offset=None,
                 prior_info=100):
Philipp Arras's avatar
Philipp Arras committed
489 490 491 492
        """
        offset vs zeromode: volume factor
        """
        if offset is not None:
493
            raise NotImplementedError
Philipp Arras's avatar
Philipp Arras committed
494
            offset = float(offset)
495

496
        op = self.finalize_from_op(self._azm, self._prefix)
497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518
        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)]

        namps = len(self.amplitudes)
        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)
519 520
            for m, s in zip(mean.flatten(), stddev.flatten()):
                print('{}: {:.02E} ± {:.02E}'.format(kk, m, s))
521 522 523 524

    def moment_slice_to_average(self, fluctuations_slice_mean, nsamples=1000):
        fluctuations_slice_mean = float(fluctuations_slice_mean)
        assert fluctuations_slice_mean > 0
525
        from ..sugar import from_random
526 527
        scm = 1.
        for a in self._a:
528 529 530 531
            op = a.fluctuation_amplitude
            res= np.array([op(from_random('normal',op.domain)).to_global_data()
                            for _ in range(nsamples)])
            scm *= res**2 + 1.
532
        return fluctuations_slice_mean/np.mean(np.sqrt(scm))
533

Philipp Arras's avatar
Philipp Arras committed
534 535
    @property
    def amplitudes(self):
536
        return self._a
Philipp Arras's avatar
Philipp Arras committed
537

538 539 540
    @property
    def amplitude_total_offset(self):
        return self._azm
Philipp Arras's avatar
Philipp Arras committed
541 542

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

Philipp Arras's avatar
Philipp Arras committed
555
    def slice_fluctuation(self, space):
556
        """Returns operator which acts on prior or posterior samples"""
557
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
558
            raise NotImplementedError
559 560
        assert space < len(self._a)
        if len(self._a) == 1:
561
            return self.average_fluctuation(0)
562 563 564 565
        q = 1.
        for j in range(len(self._a)):
            fl = self._a[j].fluctuation_amplitude
            if j == space:
Philipp Arras's avatar
Philipp Arras committed
566
                q = q*fl**2
567
            else:
Philipp Arras's avatar
Philipp Arras committed
568
                q = q*(Adder(full(fl.target, 1.)) @ fl**2)
569
        return q.sqrt() * self._azm
Philipp Arras's avatar
Philipp Arras committed
570 571

    def average_fluctuation(self, space):
572
        """Returns operator which acts on prior or posterior samples"""
573
        if len(self._a) == 0:
Philipp Arras's avatar
Philipp Arras committed
574
            raise NotImplementedError
575 576
        assert space < len(self._a)
        if len(self._a) == 1:
577 578
            return self._a[0].fluctuation_amplitude*self._azm
        return self._a[space].fluctuation_amplitude*self._azm
579

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

587 588 589 590 591 592 593 594
    @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."""
595 596 597
        ldom = len(samples[0].domain)
        assert space < ldom
        if ldom == 1:
598
            return _total_fluctuation_realized(samples)
Philipp Arras's avatar
Philipp Arras committed
599
        res1, res2 = 0., 0.
600
        for s in samples:
Philipp Frank's avatar
fixes  
Philipp Frank committed
601 602 603 604 605 606 607
            res1 = res1 + s**2
            res2 = res2 + s.mean(space)**2
        res1 = res1/len(samples)
        res2 = res2/len(samples)
        res = res1.mean() - res2.mean()
        return np.sqrt(res)

608

Philipp Arras's avatar
Philipp Arras committed
609
    @staticmethod
610 611 612 613 614 615 616 617 618 619 620
    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
621 622
        res = 0.
        for s in samples:
623 624 625 626
            r = s.mean(spaces)
            res = res + (r - r.mean())**2
        res = res/len(samples)
        return np.sqrt(res.mean())