Merge branch 'move_dynamic_prior' into 'NIFTy_5'

dynamic and light cone prior

See merge request ift/nifty-dev!152

nifty5/__init__.py

@@ -76,6 +76,8 @@ from .plot import Plot
 from .library.amplitude_operator import AmplitudeOperator
 from .library.inverse_gamma_operator import InverseGammaOperator
 from .library.los_response import LOSResponse
+from .library.dynamic_operator import dynamic_operator, dynamic_lightcone_operator
+from .library.light_cone_operator import LightConeOperator
 
 from .library.wiener_filter_curvature import WienerFilterCurvature
 from .library.correlated_fields import CorrelatedField, MfCorrelatedField

nifty5/library/dynamic_operator.py

+# 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/>.
+#
+# Copyright(C) 2013-2019 Max-Planck-Society
+#
+# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
+
+import numpy as np
+
+from ..domain_tuple import DomainTuple
+from ..domains.unstructured_domain import UnstructuredDomain
+from ..domains.rg_space import RGSpace
+from ..field import Field
+from ..operators.diagonal_operator import DiagonalOperator
+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
+from ..sugar import makeOp
+from .light_cone_operator import LightConeOperator, _field_from_function
+
+
+def _make_dynamic_operator(domain, harmonic_padding, sm_s0, sm_x0, keys=['f', 'c'],
+                          causal=True, cone=True, minimum_phase=False, sigc=3.,
+                          quant=5.):
+    dom = DomainTuple.make(domain)
+    if not isinstance(dom[0], RGSpace):
+        raise TypeError("RGSpace required")
+    ops = {}
+    FFT = FFTOperator(dom)
+    Real = Realizer(dom)
+    ops['FFT'] = FFT
+    ops['Real'] = Real
+    if harmonic_padding is None:
+        CentralPadd = ScalingOperator(1., FFT.target)
+    else:
+        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!")
+        shp = ()
+        for i in range(len(FFT.target.shape)):
+            shp += (FFT.target.shape[i] + harmonic_padding[i],)
+        CentralPadd = FieldZeroPadder(FFT.target, shp, central=True)
+    ops['central_padding'] = CentralPadd
+    sdom = CentralPadd.target[0].get_default_codomain()
+    FFTB = FFTOperator(sdom)(Realizer(sdom))
+
+    m = FieldAdapter(sdom, keys[0])
+
+    dists = m.target[0].distances
+    if isinstance(sm_x0, float):
+        sm_x0 = list((sm_x0,)*len(dists))
+    elif len(sm_x0) != len(dists):
+        raise (ValueError, "Shape mismatch!")
+
+    def smoothness_prior_func(x):
+        res = 1.
+        for i in range(len(dists)):
+            res = res + (x[i]/sm_x0[i]/dists[i])**2
+        return sm_s0/res
+
+    Sm = makeOp(_field_from_function(m.target, smoothness_prior_func))
+    m = CentralPadd.adjoint(FFTB(Sm(m)))
+    ops['smoothed_dynamics'] = m
+
+    m = -m.log()
+    if not minimum_phase:
+        m = m.exp()
+    if causal or minimum_phase:
+        m = Real.adjoint(FFT.inverse(Realizer(FFT.target).adjoint(m)))
+        kernel = makeOp(
+            _field_from_function(FFT.domain, (lambda x: 1. + np.sign(x[0]))))
+        m = kernel(m)
+
+    if cone and len(m.target.shape) > 1:
+        if isinstance(sigc, float):
+            sigc = list((sigc,)*(len(m.target.shape) - 1))
+        elif len(sigc) != len(m.target.shape) - 1:
+            raise (ValueError, "Shape mismatch!")
+        c = FieldAdapter(UnstructuredDomain(len(sigc)), keys[1])
+        c = makeOp(Field.from_global_data(c.target, np.array(sigc)))(c)
+
+        lightspeed = ScalingOperator(-0.5, c.target)(c).exp()
+        scaling = np.array(m.target[0].distances[1:])/m.target[0].distances[0]
+        scaling = DiagonalOperator(Field.from_global_data(c.target, scaling))
+        ops['lightspeed'] = scaling(lightspeed)
+
+        c = LightConeOperator(c.target, m.target, quant)(c.exp())
+        ops['light_cone'] = c
+        m = c*m
+
+    if causal or minimum_phase:
+        m = FFT(Real(m))
+    if minimum_phase:
+        m = m.exp()
+    return m, ops
+
+def dynamic_operator(domain, harmonic_padding, sm_s0, sm_x0, key,
+                     causal=True, minimum_phase=False):
+    '''
+    Constructs an operator encoding the Greens function of a linear homogeneous dynamic system.
+
+    Parameters
+    ----------
+    domain : RGSpace
+        The space under consideration
+    harmonic_padding : None, int, list of int
+        Amount of central padding in harmonic space in pixels. If None the field is not padded at all.
+    sm_s0 : float
+        Cutoff for dynamic smoothness prior
+    sm_x0 : float, List of float
+        Scaling of dynamic smoothness along each axis
+    key : String
+        key for dynamics encoding parameter.
+    causal : boolean
+        Whether or not the reconstructed dynamics should be causal in time
+    minimum_phase: boolean
+        Whether or not the reconstructed dynamics should be minimum phase
+
+    Returns
+    -------
+    Operator
+        The Operator encoding the dynamic Greens function in harmonic space.
+    Dictionary of Operator
+        A collection of sub-chains of Operators which can be used for plotting and evaluation.
+
+    Notes
+    -----
+    Currently only supports RGSpaces.
+    Note that the first axis of the space is interpreted as the time axis.
+    '''
+    return _make_dynamic_operator(domain, harmonic_padding, sm_s0, sm_x0,
+                                  keys=[key],
+                                  causal=causal, cone=False,
+                                  minimum_phase=minimum_phase)
+
+def dynamic_lightcone_operator(domain, harmonic_padding, sm_s0, sm_x0, key,
+                               lightcone_key, sigc, quant,
+                               causal=True, minimum_phase=False):
+    '''
+    Constructs an operator encoding the Greens function of a linear homogeneous dynamic system.
+    The Greens function is constrained to be within a light cone.
+
+    Parameters
+    ----------
+    domain : RGSpace
+        The space under consideration. Must have dim > 1.
+    harmonic_padding : None, int, list of int
+        Amount of central padding in harmonic space in pixels. If None the field is not padded at all.
+    sm_s0 : float
+        Cutoff for dynamic smoothness prior
+    sm_x0 : float, List of float
+        Scaling of dynamic smoothness along each axis
+    key : String
+        key for dynamics encoding parameter.
+    lightcone_key: String
+        key for lightspeed paramteter.
+    sigc : float, List of float
+        variance of lightspeed parameter.
+    quant : float
+        Quantization of the light cone in pixels.
+    causal : boolean
+        Whether or not the reconstructed dynamics should be causal in time
+    minimum_phase: boolean
+        Whether or not the reconstructed dynamics should be minimum phase
+
+    Returns
+    -------
+    Operator
+        The Operator encoding the dynamic Greens function in harmonic space when evaluated.
+    Dictionary of Operator
+        A collection of sub-chains of Operators which can be used for plotting and evaluation.
+
+    Notes
+    -----
+    Currently only supports RGSpaces.
+    Note that the first axis of the space is interpreted as the time axis.
+    '''
+    if len(domain.shape) < 2:
+        raise ValueError("Space must be at least 2 dimensional!")
+    return _make_dynamic_operator(domain,harmonic_padding,sm_s0,sm_x0,
+                                  keys=[key,lightcone_key],
+                                  causal=causal, cone=True,
+                                  minimum_phase = minimum_phase,
+                                  sigc = sigc, quant = quant)

nifty5/library/light_cone_operator.py

+# 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/>.
+#
+# Copyright(C) 2013-2019 Max-Planck-Society
+#
+# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
+
+import numpy as np
+
+from ..domain_tuple import DomainTuple
+from ..field import Field
+from ..linearization import Linearization
+from ..operators.linear_operator import LinearOperator
+from ..operators.operator import Operator
+
+
+def _field_from_function(domain, func, absolute=False):
+    domain = DomainTuple.make(domain)
+    k_array = _make_coords(domain, absolute=absolute)
+    return Field.from_global_data(domain, func(k_array))
+
+
+def _make_coords(domain, absolute=False):
+    domain = DomainTuple.make(domain)
+    dim = len(domain.shape)
+    dist = domain[0].distances
+    shape = domain.shape
+    k_array = np.zeros((dim,) + shape)
+    for i in range(dim):
+        ks = np.minimum(shape[i] - np.arange(shape[i]), np.arange(
+            shape[i]))*dist[i]
+        if not absolute:
+            ks[int(shape[i]/2) + 1:] *= -1
+        fst_dims = (1,)*i
+        lst_dims = (1,)*(dim - i - 1)
+        k_array[i] += ks.reshape(fst_dims + (shape[i],) + lst_dims)
+    return k_array
+
+
+class LightConeDerivative(LinearOperator):
+    def __init__(self, domain, target, derivatives):
+        super(LightConeDerivative, self).__init__()
+        self._domain = domain
+        self._target = target
+        self._derivatives = derivatives
+        self._capability = self.TIMES | self.ADJOINT_TIMES
+
+    def apply(self, x, mode):
+        self._check_input(x, mode)
+        x = x.to_global_data()
+        res = np.zeros(self._tgt(mode).shape, dtype=self._derivatives.dtype)
+        for i in range(self.domain.shape[0]):
+            if mode == self.TIMES:
+                res += self._derivatives[i]*x[i]
+            else:
+                res[i] = np.sum(self._derivatives[i]*x)
+        return Field.from_global_data(self._tgt(mode), res)
+
+
+def cone_arrays(c, domain, sigx, want_gradient):
+    x = _make_coords(domain)
+    a = np.zeros(domain.shape, dtype=np.complex)
+    if want_gradient:
+        derivs = np.zeros((c.size,) + domain.shape, dtype=np.complex)
+    else:
+        derivs = None
+    a -= (x[0]/(sigx*domain[0].distances[0]))**2
+    for i in range(c.size):
+        res = (x[i + 1]/(sigx*domain[0].distances[i + 1]))**2
+        a += c[i]*res
+        if want_gradient:
+            derivs[i] = res
+    a = np.sqrt(a)
+    if want_gradient:
+        derivs *= -0.5
+        for i in range(c.size):
+            derivs[i][a == 0] = 0.
+            derivs[i][a != 0] /= a[a != 0]
+    a = a.real
+    if want_gradient:
+        derivs *= a
+    a = np.exp(-0.5*a**2)
+    if want_gradient:
+        derivs = a*derivs.real
+    return a, derivs
+
+
+class LightConeOperator(Operator):
+    def __init__(self, domain, target, sigx):
+        self._domain = domain
+        self._target = target
+        self._sigx = sigx
+
+    def apply(self, x):
+        islin = isinstance(x, Linearization)
+        val = x.val.to_global_data() if islin else x.to_global_data()
+        a, derivs = cone_arrays(val, self.target, self._sigx, islin)
+        res = Field.from_global_data(self.target, a)
+        if not islin:
+            return res
+        jac = LightConeDerivative(x.jac.target, self.target, derivs)(x.jac)
+        return Linearization(res, jac, want_metric=x.want_metric)

test/test_model_gradients.py

@@ -111,3 +111,29 @@ def testPointModel(space, seed):
     model = ift.InverseGammaOperator(space, alpha, q)
     # FIXME All those cdfs and ppfs are not very accurate
     ift.extra.check_value_gradient_consistency(model, pos, tol=1e-2, ntries=20)
+
+@pmp('domain', [ift.RGSpace(64, distances=.789),
+                ift.RGSpace([32, 32], distances=.789),
+                ift.RGSpace([32, 32, 8], distances=.789)])
+@pmp('causal', [True, False])
+@pmp('minimum_phase', [True, False])
+@pmp('seed', [4, 78, 23])
+def testDynamicModel(domain, causal, minimum_phase, seed):
+    model, _ = ift.dynamic_operator(domain,None,1.,1.,'f',
+                                    causal = causal,
+                                    minimum_phase = minimum_phase)
+    S = ift.ScalingOperator(1., model.domain)
+    pos = S.draw_sample()
+    # FIXME I dont know why smaller tol fails for 3D example
+    ift.extra.check_value_gradient_consistency(model, pos, tol=1e-5,
+                                               ntries=20)
+    if len(domain.shape) > 1:
+        model, _ = ift.dynamic_lightcone_operator(domain,None,3.,1.,
+                                                  'f','c',1.,5,
+                                                  causal = causal,
+                                                  minimum_phase = minimum_phase)
+        S = ift.ScalingOperator(1., model.domain)
+        pos = S.draw_sample()
+        # FIXME I dont know why smaller tol fails for 3D example
+        ift.extra.check_value_gradient_consistency(model, pos, tol=1e-5,
+                                                   ntries=20)