dynamic_operator.py 9.93 KB
Newer Older
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/>.
#
Martin Reinecke's avatar
Martin Reinecke committed
14
# Copyright(C) 2013-2019 Max-Planck-Society
15
#
Martin Reinecke's avatar
Martin Reinecke committed
16
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
17 18 19

import numpy as np

Philipp Frank's avatar
Philipp Frank committed
20
from ..domain_tuple import DomainTuple
Philipp Frank's avatar
Philipp Frank committed
21
from ..domains.rg_space import RGSpace
Philipp Arras's avatar
Philipp Arras committed
22
from ..domains.unstructured_domain import UnstructuredDomain
Philipp Arras's avatar
Philipp Arras committed
23
from ..field import Field
Philipp Frank's avatar
cleanup  
Philipp Frank committed
24
from ..operators.diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
25 26 27 28
from ..operators.field_zero_padder import FieldZeroPadder
from ..operators.harmonic_operators import FFTOperator
from ..operators.scaling_operator import ScalingOperator
from ..operators.simple_linear_operators import FieldAdapter, Realizer
29
from ..sugar import makeOp
Philipp Arras's avatar
Fixup  
Philipp Arras committed
30
from .light_cone_operator import LightConeOperator, _field_from_function
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
31 32


Philipp Arras's avatar
Philipp Arras committed
33 34 35 36
def _float_or_listoffloat(inp):
    return [float(x) for x in inp] if isinstance(inp, list) else float(inp)


Philipp Frank's avatar
docs  
Philipp Frank committed
37
def _make_dynamic_operator(target,
Philipp Arras's avatar
Philipp Arras committed
38 39 40 41 42 43 44 45 46
                           harmonic_padding,
                           sm_s0,
                           sm_x0,
                           cone,
                           keys,
                           causal,
                           minimum_phase,
                           sigc=None,
                           quant=None):
Philipp Frank's avatar
docs  
Philipp Frank committed
47
    if not isinstance(target, RGSpace):
Philipp Arras's avatar
Philipp Arras committed
48
        raise TypeError("RGSpace required")
Philipp Frank's avatar
docs  
Philipp Frank committed
49 50
    if not target.harmonic:
        raise TypeError("Target space must be harmonic")
Philipp Arras's avatar
Philipp Arras committed
51 52 53 54 55 56 57 58 59 60 61
    if not (isinstance(harmonic_padding, int) or harmonic_padding is None
            or all(isinstance(ii, int) for ii in harmonic_padding)):
        raise TypeError
    sm_s0 = float(sm_s0)
    sm_x0 = _float_or_listoffloat(sm_x0)
    cone = bool(cone)
    if not all(isinstance(ss, str) for ss in keys):
        raise TypeError
    causal, minimum_phase = bool(causal), bool(minimum_phase)
    if sigc is not None:
        sigc = _float_or_listoffloat(sigc)
Philipp Arras's avatar
Philipp Arras committed
62 63
    if quant is not None:
        quant = float(quant)
Philipp Arras's avatar
Philipp Arras committed
64 65 66
    if cone and (sigc is None or quant is None):
        raise RuntimeError

Philipp Frank's avatar
docs  
Philipp Frank committed
67
    dom = DomainTuple.make(target.get_default_codomain())
68
    ops = {}
Philipp Frank's avatar
Philipp Frank committed
69 70
    FFT = FFTOperator(dom)
    Real = Realizer(dom)
Philipp Frank's avatar
Philipp Frank committed
71 72
    ops['FFT'] = FFT
    ops['Real'] = Real
73
    if harmonic_padding is None:
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
74
        CentralPadd = ScalingOperator(1., FFT.target)
75
    else:
Philipp Frank's avatar
Philipp Frank committed
76 77 78 79
        if isinstance(harmonic_padding, int):
            harmonic_padding = list((harmonic_padding,)*len(FFT.target.shape))
        elif len(harmonic_padding) != len(FFT.target.shape):
            raise (ValueError, "Shape mismatch!")
Philipp Arras's avatar
Philipp Arras committed
80 81 82
        shp = [
            sh + harmonic_padding[ii] for ii, sh in enumerate(FFT.target.shape)
        ]
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
83
        CentralPadd = FieldZeroPadder(FFT.target, shp, central=True)
Philipp Frank's avatar
Philipp Frank committed
84
    ops['central_padding'] = CentralPadd
85 86
    sdom = CentralPadd.target[0].get_default_codomain()
    FFTB = FFTOperator(sdom)(Realizer(sdom))
Philipp Frank's avatar
Philipp Frank committed
87

88
    m = FieldAdapter(sdom, keys[0])
Philipp Frank's avatar
Philipp Frank committed
89

90
    dists = m.target[0].distances
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
91
    if isinstance(sm_x0, float):
Philipp Frank's avatar
Philipp Frank committed
92
        sm_x0 = list((sm_x0,)*len(dists))
Philipp Frank's avatar
cleanup  
Philipp Frank committed
93
    elif len(sm_x0) != len(dists):
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
94 95
        raise (ValueError, "Shape mismatch!")

Philipp Frank's avatar
cleanup  
Philipp Frank committed
96
    def smoothness_prior_func(x):
97
        res = 1.
Philipp Frank's avatar
Philipp Frank committed
98
        for i in range(len(dists)):
99 100
            res = res + (x[i]/sm_x0[i]/dists[i])**2
        return sm_s0/res
Philipp Frank's avatar
Philipp Frank committed
101

Philipp Arras's avatar
Cleanup  
Philipp Arras committed
102
    Sm = makeOp(_field_from_function(m.target, smoothness_prior_func))
Philipp Frank's avatar
cleanup  
Philipp Frank committed
103
    m = CentralPadd.adjoint(FFTB(Sm(m)))
Philipp Frank's avatar
Philipp Frank committed
104
    ops['smoothed_dynamics'] = m
Philipp Frank's avatar
Philipp Frank committed
105

106 107 108
    m = -m.log()
    if not minimum_phase:
        m = m.exp()
Philipp Frank's avatar
Philipp Frank committed
109 110
    if causal or minimum_phase:
        m = Real.adjoint(FFT.inverse(Realizer(FFT.target).adjoint(m)))
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
111 112
        kernel = makeOp(
            _field_from_function(FFT.domain, (lambda x: 1. + np.sign(x[0]))))
113
        m = kernel(m)
Philipp Frank's avatar
Philipp Frank committed
114 115

    if cone and len(m.target.shape) > 1:
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
116 117 118 119
        if isinstance(sigc, float):
            sigc = list((sigc,)*(len(m.target.shape) - 1))
        elif len(sigc) != len(m.target.shape) - 1:
            raise (ValueError, "Shape mismatch!")
Philipp Frank's avatar
cleanup  
Philipp Frank committed
120 121 122
        c = FieldAdapter(UnstructuredDomain(len(sigc)), keys[1])
        c = makeOp(Field.from_global_data(c.target, np.array(sigc)))(c)

Philipp Arras's avatar
Cleanup  
Philipp Arras committed
123
        lightspeed = ScalingOperator(-0.5, c.target)(c).exp()
Philipp Frank's avatar
hotfix  
Philipp Frank committed
124
        scaling = np.array(m.target[0].distances[1:])/m.target[0].distances[0]
Philipp Arras's avatar
Cleanup  
Philipp Arras committed
125
        scaling = DiagonalOperator(Field.from_global_data(c.target, scaling))
Philipp Frank's avatar
cleanup  
Philipp Frank committed
126 127 128
        ops['lightspeed'] = scaling(lightspeed)

        c = LightConeOperator(c.target, m.target, quant)(c.exp())
Philipp Frank's avatar
Philipp Frank committed
129
        ops['light_cone'] = c
130 131
        m = c*m

Philipp Frank's avatar
hotfix  
Philipp Frank committed
132
    if causal or minimum_phase:
Philipp Frank's avatar
Philipp Frank committed
133
        m = FFT(Real(m))
134 135 136
    if minimum_phase:
        m = m.exp()
    return m, ops
Philipp Frank's avatar
Philipp Frank committed
137

Philipp Arras's avatar
Philipp Arras committed
138

Philipp Frank's avatar
docs  
Philipp Frank committed
139 140
def dynamic_operator(*,
                     target,
Philipp Arras's avatar
Philipp Arras committed
141 142 143 144 145 146 147 148
                     harmonic_padding,
                     sm_s0,
                     sm_x0,
                     key,
                     causal=True,
                     minimum_phase=False):
    '''Constructs an operator encoding the Green's function of a linear
    homogeneous dynamic system.
Philipp Frank's avatar
docs  
Philipp Frank committed
149 150 151 152 153 154 155 156 157 158 159
    
    When evaluated, this operator returns the Green's function representation
    in harmonic space. This result can be used as a convolution kernel to
    construct solutions of the homogeneous stochastic differential equation
    encoded in this operator. Note that if causal is True, the Green's function
    is convolved with a step function in time, where the temporal axis is the
    first axis of the space. In this case the resulting function only extends
    up to half the length of the first axis of the space to avoid boundary
    effects during convolution. If minimum_phase is true then the spectrum of
    the Green's function is used to construct a corresponding minimum phase
    filter.
Philipp Frank's avatar
Philipp Frank committed
160 161 162

    Parameters
    ----------
Philipp Frank's avatar
docs  
Philipp Frank committed
163 164
    target : RGSpace
        The harmonic space in which the Green's function shall be constructed.
Philipp Frank's avatar
Philipp Frank committed
165
    harmonic_padding : None, int, list of int
Philipp Arras's avatar
Philipp Arras committed
166 167
        Amount of central padding in harmonic space in pixels. If None the
        field is not padded at all.
Philipp Frank's avatar
Philipp Frank committed
168
    sm_s0 : float
Philipp Arras's avatar
Philipp Arras committed
169
        Cutoff for dynamic smoothness prior.
Philipp Arras's avatar
Philipp Arras committed
170
    sm_x0 : float, list of float
Philipp Arras's avatar
Philipp Arras committed
171
        Scaling of dynamic smoothness along each axis.
Philipp Frank's avatar
Philipp Frank committed
172 173 174
    key : String
        key for dynamics encoding parameter.
    causal : boolean
Philipp Arras's avatar
Docs  
Philipp Arras committed
175
        Whether or not the Green's function shall be causal in time.
Philipp Frank's avatar
docs  
Philipp Frank committed
176
        Default is True.
Philipp Frank's avatar
Philipp Frank committed
177
    minimum_phase: boolean
Philipp Arras's avatar
Docs  
Philipp Arras committed
178
        Whether or not the Green's function shall be a minimum phase filter.
Philipp Frank's avatar
docs  
Philipp Frank committed
179
        Default is False.
Philipp Frank's avatar
Philipp Frank committed
180 181 182 183

    Returns
    -------
    Operator
Philipp Frank's avatar
docs  
Philipp Frank committed
184
        The Operator encoding the dynamic Green's function in target space.
Philipp Frank's avatar
Philipp Frank committed
185
    Dictionary of Operator
Philipp Arras's avatar
Philipp Arras committed
186
        A collection of sub-chains of Operator which can be used for plotting
Philipp Arras's avatar
Philipp Arras committed
187
        and evaluation.
Philipp Frank's avatar
Philipp Frank committed
188 189 190

    Notes
    -----
Philipp Arras's avatar
Philipp Arras committed
191
    The first axis of the domain is interpreted the time axis.
Philipp Frank's avatar
Philipp Frank committed
192
    '''
Philipp Arras's avatar
Philipp Arras committed
193
    dct = {
Philipp Frank's avatar
docs  
Philipp Frank committed
194
        'target': target,
Philipp Arras's avatar
Philipp Arras committed
195 196 197 198 199 200 201 202 203 204 205
        'harmonic_padding': harmonic_padding,
        'sm_s0': sm_s0,
        'sm_x0': sm_x0,
        'keys': [key],
        'causal': causal,
        'cone': False,
        'minimum_phase': minimum_phase,
    }
    return _make_dynamic_operator(**dct)


Philipp Frank's avatar
docs  
Philipp Frank committed
206 207
def dynamic_lightcone_operator(*,
                               target,
Philipp Arras's avatar
Philipp Arras committed
208 209 210 211 212 213 214 215 216
                               harmonic_padding,
                               sm_s0,
                               sm_x0,
                               key,
                               lightcone_key,
                               sigc,
                               quant,
                               causal=True,
                               minimum_phase=False):
Philipp Frank's avatar
docs  
Philipp Frank committed
217 218
    '''Extends the functionality of :function: dynamic_operator to a Green's
    function which is constrained to be within a light cone.
Philipp Frank's avatar
docs  
Philipp Frank committed
219 220 221 222 223
    
    The resulting Green's function is constrained to be within a light cone.
    This is achieved via convolution of the function with a light cone in
    space-time. Thereby the first axis of the space is set to be the teporal
    axis.
Philipp Frank's avatar
Philipp Frank committed
224 225 226

    Parameters
    ----------
Philipp Frank's avatar
docs  
Philipp Frank committed
227 228
    target : RGSpace
        The harmonic space in which the Green's function shall be constructed.
Philipp Arras's avatar
Docs  
Philipp Arras committed
229
        It needs to have at least two dimensions.
Philipp Frank's avatar
Philipp Frank committed
230
    harmonic_padding : None, int, list of int
Philipp Arras's avatar
Philipp Arras committed
231 232
        Amount of central padding in harmonic space in pixels. If None the
        field is not padded at all.
Philipp Frank's avatar
Philipp Frank committed
233
    sm_s0 : float
Philipp Arras's avatar
Philipp Arras committed
234
        Cutoff for dynamic smoothness prior.
Philipp Arras's avatar
Philipp Arras committed
235
    sm_x0 : float, list of float
Philipp Arras's avatar
Philipp Arras committed
236
        Scaling of dynamic smoothness along each axis.
Philipp Frank's avatar
Philipp Frank committed
237
    key : String
Philipp Arras's avatar
Philipp Arras committed
238
        Key for dynamics encoding parameter.
Philipp Frank's avatar
Philipp Frank committed
239
    lightcone_key: String
Philipp Arras's avatar
Philipp Arras committed
240
        Key for lightspeed paramteter.
Philipp Arras's avatar
Philipp Arras committed
241
    sigc : float, list of float
Philipp Arras's avatar
Philipp Arras committed
242
        Variance of lightspeed parameter.
Philipp Frank's avatar
Philipp Frank committed
243 244 245
    quant : float
        Quantization of the light cone in pixels.
    causal : boolean
Philipp Arras's avatar
Docs  
Philipp Arras committed
246
        Whether or not the Green's function shall be causal in time.
Philipp Frank's avatar
docs  
Philipp Frank committed
247
        Default is True.
Philipp Frank's avatar
Philipp Frank committed
248
    minimum_phase: boolean
Philipp Arras's avatar
Docs  
Philipp Arras committed
249
        Whether or not the Green's function shall be a minimum phase filter.
Philipp Frank's avatar
docs  
Philipp Frank committed
250
        Default is False.
Philipp Frank's avatar
Philipp Frank committed
251 252 253 254

    Returns
    -------
    Operator
Philipp Frank's avatar
docs  
Philipp Frank committed
255 256 257 258
        The Operator encoding the dynamic Green's function in harmonic space.
    Dictionary of Operator
        A collection of sub-chains of Operator which can be used for plotting
        and evaluation.
Philipp Frank's avatar
Philipp Frank committed
259 260 261

    Notes
    -----
Philipp Arras's avatar
Philipp Arras committed
262
    The first axis of the domain is interpreted the time axis.
Philipp Frank's avatar
Philipp Frank committed
263
    '''
Philipp Arras's avatar
Philipp Arras committed
264

Philipp Frank's avatar
docs  
Philipp Frank committed
265
    if len(target.shape) < 2:
Philipp Frank's avatar
Philipp Frank committed
266
        raise ValueError("Space must be at least 2 dimensional!")
Philipp Arras's avatar
Philipp Arras committed
267
    dct = {
Philipp Frank's avatar
docs  
Philipp Frank committed
268
        'target': target,
Philipp Arras's avatar
Philipp Arras committed
269 270 271 272 273 274 275 276 277 278 279
        'harmonic_padding': harmonic_padding,
        'sm_s0': sm_s0,
        'sm_x0': sm_x0,
        'keys': [key, lightcone_key],
        'causal': causal,
        'cone': True,
        'minimum_phase': minimum_phase,
        'sigc': sigc,
        'quant': quant
    }
    return _make_dynamic_operator(**dct)