Merge branch 'convenience_functions' into 'NIFTy_5'

Convenience functions

See merge request ift/nifty-dev!109

demos/getting_started_3.py

@@ -105,6 +105,7 @@ if __name__ == '__main__':
             title="power")
         plot.output(ny=1, ysize=6, xsize=16, name="loop.png")
 
+    KL = ift.KL_Energy(position, H, N_samples)
     plot = ift.Plot()
     sc = ift.StatCalculator()
     for sample in KL.samples:

nifty5/__init__.py

@@ -73,13 +73,14 @@ from .minimization.kl_energy import KL_Energy
 from .sugar import *
 from .plot import Plot
 
-from .library.amplitude_model import AmplitudeModel
+from .library.amplitude_model import AmplitudeModel, SmoothAmplitudeModel
 from .library.inverse_gamma_model import InverseGammaModel
 from .library.los_response import LOSResponse
 
 from .library.wiener_filter_curvature import WienerFilterCurvature
 from .library.correlated_fields import CorrelatedField, MfCorrelatedField
-from .library.adjust_variances import make_adjust_variances
+from .library.adjust_variances import (make_adjust_variances,
+                                       do_adjust_variances)
 
 from . import extra
 

nifty5/domains/log_rg_space.py

@@ -28,7 +28,21 @@ from .structured_domain import StructuredDomain
 
 
 class LogRGSpace(StructuredDomain):
-
+    """NIFTy subclass for logarithmic Cartesian grids.
+
+    Parameters
+    ----------
+    shape : int or tuple of int
+        Number of grid points or numbers of gridpoints along each axis.
+    bindistances : float or tuple of float
+        Distance between two grid points along each axis. These are
+        measured on logarithmic scale and are constant therfore.
+    t_0 : float or tuple of float
+        FIXME
+    harmonic : bool, optional
+        Whether the space represents a grid in position or harmonic space.
+        (default: False).
+    """
     _needed_for_hash = ['_shape', '_bindistances', '_t_0', '_harmonic']
 
     def __init__(self, shape, bindistances, t_0, harmonic=False):
@@ -41,8 +55,8 @@ class LogRGSpace(StructuredDomain):
         self._bindistances = tuple(bindistances)
         self._t_0 = tuple(t_0)
 
-        self._dim = int(reduce(lambda x, y: x * y, self._shape))
-        self._dvol = float(reduce(lambda x, y: x * y, self._bindistances))
+        self._dim = int(reduce(lambda x, y: x*y, self._shape))
+        self._dvol = float(reduce(lambda x, y: x*y, self._bindistances))
 
     @property
     def harmonic(self):
@@ -69,24 +83,23 @@ class LogRGSpace(StructuredDomain):
         return np.array(self._t_0)
 
     def __repr__(self):
-        return ("LogRGSpace(shape={}, harmonic={})"
-                .format(self.shape, self.harmonic))
+        return ("LogRGSpace(shape={}, harmonic={})".format(
+            self.shape, self.harmonic))
 
     def get_default_codomain(self):
-        codomain_bindistances = 1. / (self.bindistances * self.shape)
-        return LogRGSpace(self.shape, codomain_bindistances,
-                          self._t_0, True)
+        codomain_bindistances = 1./(self.bindistances*self.shape)
+        return LogRGSpace(self.shape, codomain_bindistances, self._t_0, True)
 
     def get_k_length_array(self):
         ib = dobj.ibegin_from_shape(self._shape)
         res = np.arange(self.local_shape[0], dtype=np.float64) + ib[0]
-        res = np.minimum(res, self.shape[0]-res)*self.bindistances[0]
+        res = np.minimum(res, self.shape[0] - res)*self.bindistances[0]
         if len(self.shape) == 1:
             return Field.from_local_data(self, res)
         res *= res
         for i in range(1, len(self.shape)):
             tmp = np.arange(self.local_shape[i], dtype=np.float64) + ib[i]
-            tmp = np.minimum(tmp, self.shape[i]-tmp)*self.bindistances[i]
+            tmp = np.minimum(tmp, self.shape[i] - tmp)*self.bindistances[i]
             tmp *= tmp
             res = np.add.outer(res, tmp)
         return Field.from_local_data(self, np.sqrt(res))

nifty5/fft.py

@@ -20,7 +20,7 @@ if _use_fftw:
     _fft_extra_args = dict(planner_effort='FFTW_ESTIMATE', threads=nthreads)
 else:
     from numpy.fft import fftn, rfftn, ifftn
-    _fft_extra_args={}
+    _fft_extra_args = {}
 
 
 def hartley(a, axes=None):

nifty5/library/adjust_variances.py

@@ -19,11 +19,20 @@
 from __future__ import absolute_import, division, print_function
 
 from ..compat import *
+from ..minimization.energy_adapter import EnergyAdapter
+from ..multi_domain import MultiDomain
+from ..multi_field import MultiField
+from ..operators.distributors import PowerDistributor
 from ..operators.energy_operators import Hamiltonian, InverseGammaLikelihood
 from ..operators.scaling_operator import ScalingOperator
+from ..operators.simple_linear_operators import FieldAdapter
 
 
-def make_adjust_variances(a, xi, position, samples=[], scaling=None,
+def make_adjust_variances(a,
+                          xi,
+                          position,
+                          samples=[],
+                          scaling=None,
                           ic_samp=None):
     """ Creates a Hamiltonian for constant likelihood optimizations.
 
@@ -67,3 +76,36 @@ def make_adjust_variances(a, xi, position, samples=[], scaling=None,
         x = ScalingOperator(scaling, x.target)(x)
 
     return Hamiltonian(InverseGammaLikelihood(d_eval)(x), ic_samp=ic_samp)
+
+
+def do_adjust_variances(position, xi, amplitude_model, minimizer, samples=[]):
+    h_space = xi.target[0]
+    pd = PowerDistributor(h_space, amplitude_model.target[0])
+    a = pd(amplitude_model)
+    xi = FieldAdapter(MultiDomain.make({"xi": h_space}), "xi")
+
+    axi_domain = MultiDomain.union([a.domain, xi.domain])
+    ham = make_adjust_variances(
+        a, xi, position.extract(axi_domain), samples=samples)
+    a_pos = position.extract(a.domain)
+
+    # Minimize
+    # FIXME Try also KL here
+    e = EnergyAdapter(a_pos, ham, constants=[], want_metric=True)
+    e, _ = minimizer(e)
+
+    # Update position
+    s_h_old = (a*xi)(position.extract(axi_domain))
+
+    position = position.to_dict()
+    position['xi'] = s_h_old/a(e.position)
+    position = MultiField.from_dict(position)
+    position = MultiField.union([position, e.position])
+
+    s_h_new = (a*xi)(position.extract(axi_domain))
+
+    import numpy as np
+    # TODO Move this into the tests
+    np.testing.assert_allclose(s_h_new.to_global_data(),
+                               s_h_old.to_global_data())
+    return position

nifty5/library/amplitude_model.py

@@ -150,3 +150,54 @@ class AmplitudeModel(Operator):
     @property
     def norm_phi_mean(self):
         return self._norm_phi_mean
+
+
+class SmoothAmplitudeModel(Operator):
+    '''
+    Computes a smooth power spectrum.
+    Output lives in PowerSpace.
+
+    Parameters
+    ----------
+
+    Npixdof : #pix in dof_space
+
+    ceps_a, ceps_k0 : Smoothness parameters in ceps_kernel
+                        eg. ceps_kernel(k) = (a/(1+(k/k0)**2))**2
+                        a = ceps_a,  k0 = ceps_k0
+    '''
+    def __init__(self, s_space, Npixdof, ceps_a, ceps_k, name='tau'):
+        from ..operators.exp_transform import ExpTransform
+        from ..operators.qht_operator import QHTOperator
+
+        h_space = s_space.get_default_codomain()
+        self._exp_transform = ExpTransform(PowerSpace(h_space), Npixdof)
+        logk_space = self._exp_transform.domain[0]
+        qht = QHTOperator(target=logk_space)
+        dof_space = qht.domain[0]
+
+        self._domain = MultiDomain.make({name: dof_space})
+        self._target = self._exp_transform.target
+
+        kern = lambda k: _ceps_kernel(dof_space, k, ceps_a, ceps_k)
+        cepstrum = create_cepstrum_amplitude_field(dof_space, kern)
+
+        self._smooth_op = qht(makeOp(sqrt(cepstrum)))
+        self._key = name
+
+        self._qht = qht
+        self._ceps = makeOp(sqrt(cepstrum))
+
+    def apply(self, x):
+        self._check_input(x)
+        smooth_spec = self._smooth_op(x[self._key])
+        loglog_spec = smooth_spec
+        return self._exp_transform((0.5*loglog_spec).exp())
+
+    @property
+    def qht(self):
+        return self._qht
+
+    @property
+    def ceps(self):
+        return self._ceps

nifty5/operators/operator.py

@@ -59,6 +59,9 @@ class Operator(NiftyMetaBase()):
     def apply(self, x):
         raise NotImplementedError
 
+    def force(self, x):
+        return self.apply(x.extract(self.domain))
+
     def _check_input(self, x):
         from ..linearization import Linearization
         d = x.target if isinstance(x, Linearization) else x.domain