correlated_fields.py 15.5 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
from functools import reduce
Philipp Arras's avatar
Philipp Arras committed
21
from numpy.testing import assert_allclose
22

Philipp Arras's avatar
Philipp Arras committed
23
from ..domain_tuple import DomainTuple
Philipp Arras's avatar
Philipp Arras committed
24
25
from ..domains.power_space import PowerSpace
from ..domains.unstructured_domain import UnstructuredDomain
Philipp Arras's avatar
Philipp Arras committed
26
from ..extra import check_jacobian_consistency, consistency_check
27
from ..field import Field
Philipp Arras's avatar
Philipp Arras committed
28
from ..multi_domain import MultiDomain
Philipp Arras's avatar
Philipp Arras committed
29
from ..operators.adder import Adder
30
from ..operators.contraction_operator import ContractionOperator
Philipp Arras's avatar
Philipp Arras committed
31
from ..operators.distributors import PowerDistributor
Philipp Arras's avatar
Philipp Arras committed
32
from ..operators.endomorphic_operator import EndomorphicOperator
Martin Reinecke's avatar
Martin Reinecke committed
33
from ..operators.harmonic_operators import HarmonicTransformOperator
Philipp Arras's avatar
Philipp Arras committed
34
from ..operators.linear_operator import LinearOperator
35
from ..operators.diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
36
37
from ..operators.operator import Operator
from ..operators.simple_linear_operators import VdotOperator, ducktape
Philipp Arras's avatar
Philipp Arras committed
38
from ..operators.value_inserter import ValueInserter
39
40
from ..sugar import from_global_data, from_random, full, makeDomain, get_default_codomain

41
def _reshaper(x, shape):
42
    x = np.array(x)
43

44
45
46
47
48
49
50
51
52
53
    if x.shape == shape:
        return np.asfarray(x)
    elif x.shape in [(), (1,)]:
        return np.full(shape, x, dtype=np.float)
    else:
        raise TypeError("Shape of parameters cannot be interpreted")

def _lognormal_moment_matching(mean, sig, key,
        domain = DomainTuple.scalar_domain(), space = 0):
    domain = makeDomain(domain)
54
    mean, sig = (_reshaper(param, domain.shape) for param in (mean, sig))
55
56
57
    key = str(key)
    assert np.all(mean > 0)
    assert np.all(sig > 0)
Philipp Arras's avatar
Philipp Arras committed
58
59
    logsig = np.sqrt(np.log((sig/mean)**2 + 1))
    logmean = np.log(mean) - logsig**2/2
60
    return _normal(logmean, logsig, key, domain).exp()
Philipp Arras's avatar
Philipp Arras committed
61
62


63
64
65
def _normal(mean, sig, key,
        domain = DomainTuple.scalar_domain(), space = 0):
    domain = makeDomain(domain)
66
    mean, sig = (_reshaper(param, domain.shape) for param in (mean, sig))
67
68
    assert np.all(sig > 0)
    return Adder(from_global_data(domain, mean)) @ (
Philipp Haim's avatar
Fixes    
Philipp Haim committed
69
        DiagonalOperator(from_global_data(domain,sig)) @ ducktape(domain, None, key))
Philipp Arras's avatar
Philipp Arras committed
70
71


Philipp Frank's avatar
Philipp Frank committed
72
class _SlopeRemover(EndomorphicOperator):
73
    def __init__(self, domain, cooridinates, space = 0):
Philipp Frank's avatar
Philipp Frank committed
74
        self._domain = makeDomain(domain)
75
        self._sc = cooridinates / float(cooridinates[-1])
Philipp Arras's avatar
Philipp Arras committed
76

77
        self._space = space
Philipp Haim's avatar
Philipp Haim committed
78
79
80
        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
81
        self._capability = self.TIMES | self.ADJOINT_TIMES
Philipp Arras's avatar
Philipp Arras committed
82

Philipp Frank's avatar
Philipp Frank committed
83
84
85
86
    def apply(self,x,mode):
        self._check_input(x,mode)
        x = x.to_global_data()
        if mode == self.TIMES:
Philipp Haim's avatar
Philipp Haim committed
87
            res = x - x[self._last]*self._sc[self._extender]
Philipp Frank's avatar
Philipp Frank committed
88
        else:
89
            #NOTE Why not x.copy()?
Philipp Frank's avatar
Philipp Frank committed
90
            res = np.zeros(x.shape,dtype=x.dtype)
Philipp Frank's avatar
Philipp Frank committed
91
            res += x
Philipp Haim's avatar
Philipp Haim committed
92
            res[self._last] -= np.sum(x*self._sc[self._extender], axis = self._space, keepdims = True)
Philipp Frank's avatar
Philipp Frank committed
93
94
        return from_global_data(self._tgt(mode),res)

95
def _make_slope_Operator(smooth,loglogavgslope, space = 0):
Philipp Frank's avatar
Philipp Frank committed
96
    tg = smooth.target
97
    logkl = _log_k_lengths(tg[space])
Philipp Frank's avatar
Philipp Frank committed
98
99
    logkl -= logkl[0]
    logkl = np.insert(logkl, 0, 0)
Philipp Haim's avatar
Fixes    
Philipp Haim committed
100
    noslope = smooth
101
    noslope = _SlopeRemover(tg,logkl, space) @ smooth
Philipp Frank's avatar
Philipp Frank committed
102
    # FIXME Move to tests
Philipp Haim's avatar
Philipp Haim committed
103
    consistency_check(_SlopeRemover(tg,logkl, space))
Philipp Frank's avatar
Philipp Frank committed
104

105
    expander = ContractionOperator(tg, spaces = space).adjoint
Philipp Haim's avatar
Fixes    
Philipp Haim committed
106
    _t = DiagonalOperator(from_global_data(tg[space], logkl), tg, spaces = space)
107
    return _t @ expander @ loglogavgslope + noslope
Philipp Arras's avatar
Philipp Arras committed
108
109
110
111
112

def _log_k_lengths(pspace):
    return np.log(pspace.k_lengths[1:])

class _TwoLogIntegrations(LinearOperator):
113
    def __init__(self, target, space = None):
Philipp Arras's avatar
Philipp Arras committed
114
        self._target = makeDomain(target)
115
116
117
118
119
        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
Philipp Arras's avatar
Philipp Arras committed
120
        self._capability = self.TIMES | self.ADJOINT_TIMES
121
        logk_lengths = _log_k_lengths(self._target[space])
Philipp Arras's avatar
Philipp Arras committed
122
123
124
125
        self._logvol = logk_lengths[1:] - logk_lengths[:-1]

    def apply(self, x, mode):
        self._check_input(x, mode)
126
127
128
129

        #Maybe make class properties
        axis = self._target.axes[self._space][0]
        sl = (slice(None),)*axis
Philipp Haim's avatar
Fixes    
Philipp Haim committed
130
        extender_sl = (None,)*axis + (slice(None),) + (None,)*(self._target.axes[-1][-1] - axis)
131
132
133
134
135
        first = sl + (0,)
        second = sl + (1,)
        from_third = sl + (slice(2,None),)
        no_border = sl + (slice(1,-1),)
        reverse = sl + (slice(None,None,-1),)
Philipp Arras's avatar
Philipp Arras committed
136
137
138
        if mode == self.TIMES:
            x = x.to_global_data()
            res = np.empty(self._target.shape)
139
140
141
            res[first] = 0
            res[second] = 0
            res[from_third] = np.cumsum(x[second], axis = axis)
Philipp Haim's avatar
Fixes    
Philipp Haim committed
142
            res[from_third] = (res[from_third] + res[no_border])/2*self._logvol[extender_sl] + x[first]
143
            res[from_third] = np.cumsum(res[from_third], axis = axis)
Philipp Arras's avatar
Philipp Arras committed
144
145
146
        else:
            x = x.to_global_data_rw()
            res = np.zeros(self._domain.shape)
147
148
            x[from_third] = np.cumsum(x[from_third][reverse], axis = axis)[reverse]
            res[first] += x[from_third]
Philipp Haim's avatar
Fixes    
Philipp Haim committed
149
            x[from_third] *= (self._logvol/2.)[extender_sl]
150
151
152
            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
153
154
155


class _Normalization(Operator):
156
    def __init__(self, domain, space = 0):
Philipp Arras's avatar
Philipp Arras committed
157
        self._domain = self._target = makeDomain(domain)
158
159
160
161
        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
162
        # TODO Does not work on sphere yet
163
        mode_multiplicity = pd.adjoint(full(pd.target, 1.)).to_global_data_rw()
Philipp Haim's avatar
Philipp Haim committed
164
165
        zero_mode = (slice(None),)*domain.axes[space][0] + (0,)
        mode_multiplicity[zero_mode] = 0
166
167
        self._mode_multiplicity = from_global_data(self._domain, mode_multiplicity)
        self._specsum = _SpecialSum(self._domain, space)
Philipp Arras's avatar
Philipp Arras committed
168
169
        # FIXME Move to tests
        consistency_check(self._specsum)
Philipp Arras's avatar
Philipp Arras committed
170
171
172
173
174
175
176

    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
177
        return self._specsum(self._mode_multiplicity*spec)**(-0.5)*amp
Philipp Arras's avatar
Philipp Arras committed
178
179
180


class _SpecialSum(EndomorphicOperator):
181
    def __init__(self, domain, space = 0):
Philipp Arras's avatar
Philipp Arras committed
182
183
        self._domain = makeDomain(domain)
        self._capability = self.TIMES | self.ADJOINT_TIMES
184
        self._contractor = ContractionOperator(domain, space)
Philipp Arras's avatar
Philipp Arras committed
185
186
187

    def apply(self, x, mode):
        self._check_input(x, mode)
188
        return self._contractor.adjoint(self._contractor(x))
Philipp Arras's avatar
Philipp Arras committed
189
190


191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
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 Arras's avatar
Philipp Arras committed
211
212
213
class CorrelatedFieldMaker:
    def __init__(self):
        self._amplitudes = []
214
        self._spaces = []
Philipp Arras's avatar
Philipp Arras committed
215
216

    def add_fluctuations_from_ops(self, target, fluctuations, flexibility,
217
                                  asperity, loglogavgslope, key, space = 0):
Philipp Arras's avatar
Philipp Arras committed
218
219
220
221
222
223
224
225
226
227
228
        """
        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)
        target = makeDomain(target)
229
        assert isinstance(target[space], PowerSpace)
Philipp Arras's avatar
Philipp Arras committed
230

231
        twolog = _TwoLogIntegrations(target, space)
232
        dt = twolog._logvol
Philipp Haim's avatar
Fixes    
Philipp Haim committed
233
        axis = target.axes[space][0]
234
        first = (slice(None),)*axis + (0,)
Philipp Haim's avatar
Fixes    
Philipp Haim committed
235
        extender_sl = (None,)*axis + (slice(None),) + (None,)*(target.axes[-1][-1] - axis)
236
237
        expander = ContractionOperator(twolog.domain, spaces = space).adjoint
        
Philipp Haim's avatar
Fixes    
Philipp Haim committed
238
        sqrt_t = np.zeros(twolog.domain[space].shape)
Philipp Haim's avatar
Philipp Haim committed
239
        sqrt_t[0] = sqrt_t[1] = np.sqrt(dt)
Philipp Haim's avatar
Fixes    
Philipp Haim committed
240
        sqrt_t = from_global_data(twolog.domain[space], sqrt_t)
241
242
243
        sqrt_t = DiagonalOperator(sqrt_t, twolog.domain, spaces = space)
        sigmasq = sqrt_t @ expander @ flexibility

Philipp Haim's avatar
Fixes    
Philipp Haim committed
244
        dist = np.zeros(twolog.domain[space].shape)
245
        dist[0] += 1.
Philipp Haim's avatar
Fixes    
Philipp Haim committed
246
        dist = from_global_data(twolog.domain[space], dist)
247
        dist = DiagonalOperator(dist, twolog.domain, spaces = space)
Philipp Arras's avatar
Philipp Arras committed
248

249
        shift = np.ones(twolog.domain.shape)
Philipp Haim's avatar
Fixes    
Philipp Haim committed
250
        shift[first] = (dt**2/12.)[extender_sl]
251
252
        shift = from_global_data(twolog.domain, shift)
        scale = sigmasq*(Adder(shift) @ dist @ expander @ asperity).sqrt()
Philipp Arras's avatar
Philipp Arras committed
253
254

        smooth = twolog @ (scale*ducktape(scale.target, None, key))
Philipp Haim's avatar
Fixes    
Philipp Haim committed
255
        smoothslope = _make_slope_Operator(smooth,loglogavgslope, space)
Philipp Frank's avatar
Philipp Frank committed
256
        
Philipp Arras's avatar
Philipp Arras committed
257
        # move to tests
258
259
260
        assert_allclose( 
                smooth(from_random('normal', smooth.domain)).val[(
                slice(None),)*axis + (slice(0,2),)], 0)
Philipp Arras's avatar
Philipp Arras committed
261
        consistency_check(twolog)
Philipp Arras's avatar
Philipp Arras committed
262
263
        check_jacobian_consistency(smooth, from_random('normal',
                                                       smooth.domain))
Philipp Arras's avatar
Philipp Arras committed
264
265
        check_jacobian_consistency(smoothslope,
                                   from_random('normal', smoothslope.domain))
Philipp Arras's avatar
Philipp Arras committed
266
267
        # end move to tests

Philipp Haim's avatar
Fixes    
Philipp Haim committed
268
        normal_ampl = _Normalization(target, space) @ smoothslope
Philipp Haim's avatar
Philipp Haim committed
269
        vol = target[space].harmonic_partner.get_default_codomain().total_volume
Philipp Haim's avatar
Fixes    
Philipp Haim committed
270
        arr = np.zeros(target[space].shape)
Philipp Arras's avatar
Philipp Arras committed
271
        arr[1:] = vol
Philipp Haim's avatar
Fixes    
Philipp Haim committed
272
273
274
        expander = ContractionOperator(target, spaces = space).adjoint
        expander = DiagonalOperator(from_global_data(target[space], arr)
                , target, spaces = space) @ expander
Philipp Arras's avatar
Philipp Arras committed
275
        mask = np.zeros(target.shape)
276
        mask[first] = vol
Philipp Arras's avatar
Philipp Arras committed
277
278
279
280
281
282
283
        adder = Adder(from_global_data(target, mask))
        ampl = adder @ ((expander @ fluctuations)*normal_ampl)

        # Move to tests
        # FIXME This test fails but it is not relevant for the final result
        # assert_allclose(
        #     normal_ampl(from_random('normal', normal_ampl.domain)).val[0], 1)
Philipp Haim's avatar
Philipp Haim committed
284
        #assert_allclose(ampl(from_random('normal', ampl.domain)).val[space][0], vol)
Philipp Haim's avatar
Fixes    
Philipp Haim committed
285
        op = _Normalization(target, space)
Philipp Arras's avatar
Philipp Arras committed
286
        check_jacobian_consistency(op, from_random('normal', op.domain))
Philipp Arras's avatar
Philipp Arras committed
287
288
289
        # End move to tests

        self._amplitudes.append(ampl)
290
        self._spaces.append(space)
Philipp Arras's avatar
Philipp Arras committed
291
292
293
294

    def add_fluctuations(self, target, fluctuations_mean, fluctuations_stddev,
                         flexibility_mean, flexibility_stddev, asperity_mean,
                         asperity_stddev, loglogavgslope_mean,
295
                         loglogavgslope_stddev, prefix, space = 0):
Philipp Arras's avatar
Philipp Arras committed
296
297
        prefix = str(prefix)

298
299
300
301
302
303
304
        parameter_domain = list(makeDomain(target))
        del parameter_domain[space]
        if parameter_domain != []:
            parameter_domain = makeDomain(parameter_domain)
        else:
            parameter_domain = DomainTuple.scalar_domain()

305
        fluct = _lognormal_moment_matching(fluctuations_mean, fluctuations_stddev,
306
                        prefix + 'fluctuations', parameter_domain, space = space)
Philipp Arras's avatar
Philipp Arras committed
307
        flex = _lognormal_moment_matching(flexibility_mean, flexibility_stddev,
308
                        prefix + 'flexibility', parameter_domain, space = space)
Philipp Arras's avatar
Philipp Arras committed
309
        asp = _lognormal_moment_matching(asperity_mean, asperity_stddev,
310
                        prefix + 'asperity', parameter_domain, space = space)
311
        avgsl = _normal(loglogavgslope_mean, loglogavgslope_stddev,
312
313
                        prefix + 'loglogavgslope', parameter_domain, space = space)

Philipp Haim's avatar
Fixes    
Philipp Haim committed
314
        return self.add_fluctuations_from_ops(target, fluct, flex, asp, avgsl,
315
                                       prefix + 'spectrum', space)
Philipp Arras's avatar
Philipp Arras committed
316
317
318
319
320
321
322
323
324
325
326
327

    def finalize_from_op(self, zeromode):
        raise NotImplementedError

    def finalize(self,
                 offset_amplitude_mean,
                 offset_amplitude_stddev,
                 prefix,
                 offset=None):
        """
        offset vs zeromode: volume factor
        """
328
        prefix = str(prefix)
Philipp Arras's avatar
Philipp Arras committed
329
        if offset is not None:
330
            offset = float(offset) 
331
332
333
334
335
336
337
338
339
        hspace = []
        zeroind = ()
        for amp, space in zip(self._amplitudes, self._spaces):
            dd =  list(amp.target)
            dd[space] = dd[space].harmonic_partner
            hspace.extend(dd)
            zeroind += (slice(None),)*space + (0,)*len(dd[space].shape)
        hspace = makeDomain(hspace)
        spaces = np.cumsum(self._spaces) + np.arange(len(self._spaces))
Philipp Arras's avatar
Philipp Arras committed
340

341
342
343
344
345
346
347
348
        parameter_domain = list(makeDomain(hspace))
        for space in self._spaces:
            del parameter_domain[space]
        if parameter_domain != []:
            parameter_domain = makeDomain(parameter_domain)
        else:
            parameter_domain = DomainTuple.scalar_domain()

Philipp Arras's avatar
Philipp Arras committed
349
350
        azm = _lognormal_moment_matching(offset_amplitude_mean,
                                         offset_amplitude_stddev,
351
352
                                         prefix + 'zeromode', parameter_domain,
                                         space = tuple(self._spaces))
353

Philipp Arras's avatar
Philipp Arras committed
354
355
356
        foo = np.ones(hspace.shape)
        foo[zeroind] = 0

357
358
359
360
361
362
363
364
        ZeroModeInserter = _slice_extractor(hspace, azm.target, zeroind).adjoint

        azm = Adder(from_global_data(hspace, foo)) @ ZeroModeInserter @ azm

        #NOTE ht and pd operator able to act on several spaces might be nice
        ht = HarmonicTransformOperator(hspace, space = spaces[0])
        pd = PowerDistributor(hspace, 
                self._amplitudes[0].target[spaces[0]], spaces[0])
Philipp Arras's avatar
Philipp Arras committed
365
        for i in range(1, len(self._amplitudes)):
366
367
368
            ht = HarmonicTransformOperator(ht.target, space = spaces[i]) @ ht
            pd = pd @ PowerDistributor( pd.domain, 
                    self._amplitudes[i].target[spaces[i]], space = spaces[i])
Philipp Arras's avatar
Philipp Arras committed
369
370
371
372
373
374
375
376
377
378
379
380
381

        a = ContractionOperator(pd.domain,
                                spaces[1:]).adjoint(self._amplitudes[0])
        for i in range(1, len(self._amplitudes)):
            a = a*(ContractionOperator(pd.domain, spaces[:i] + spaces[
                (i + 1):]).adjoint(self._amplitudes[i]))

        A = pd @ a
        return ht(azm*A*ducktape(hspace, None, prefix + 'xi'))

    @property
    def amplitudes(self):
        return self._amplitudes