Merge branch 'working_on_demos' into 'NIFTy_4'

Working on demos

See merge request ift/NIFTy!260

demos/wiener_filter_data_space_noiseless.py

+import numpy as np
+import nifty4 as ift
+
+
+# TODO: MAKE RESPONSE MPI COMPATIBLE OR USE LOS RESPONSE INSTEAD
+
+class CustomResponse(ift.LinearOperator):
+    """
+    A custom operator that measures a specific points`
+
+    An operator that is a delta measurement at certain points
+    """
+    def __init__(self, domain, data_points):
+        self._domain = ift.DomainTuple.make(domain)
+        self._points = data_points
+        data_shape = ift.Field.full(domain, 0.).to_global_data()[data_points]\
+            .shape
+        self._target = ift.DomainTuple.make(ift.UnstructuredDomain(data_shape))
+
+    def _times(self, x):
+        d = np.zeros(self._target.shape, dtype=np.float64)
+        d += x.to_global_data()[self._points]
+        return ift.from_global_data(self._target, d)
+
+    def _adjoint_times(self, d):
+        x = np.zeros(self._domain.shape, dtype=np.float64)
+        x[self._points] += d.to_global_data()
+        return ift.from_global_data(self._domain, x)
+
+    @property
+    def domain(self):
+        return self._domain
+
+    @property
+    def target(self):
+        return self._target
+
+    def apply(self, x, mode):
+        self._check_input(x, mode)
+        return self._times(x) if mode == self.TIMES else self._adjoint_times(x)
+
+    @property
+    def capability(self):
+        return self.TIMES | self.ADJOINT_TIMES
+
+if __name__ == "__main__":
+    np.random.seed(43)
+    # Set up physical constants
+    # Total length of interval or volume the field lives on, e.g. in meters
+    L = 2.
+    # Typical distance over which the field is correlated (in same unit as L)
+    correlation_length = 0.3
+    # Variance of field in position space sqrt(<|s_x|^2>) (in same unit as s)
+    field_variance = 2.
+    # Smoothing length of response (in same unit as L)
+    response_sigma = 0.01
+    # typical noise amplitude of the measurement
+    noise_level = 0.
+
+    # Define resolution (pixels per dimension)
+    N_pixels = 256
+
+    # Set up derived constants
+    k_0 = 1./correlation_length
+    # defining a power spectrum with the right correlation length
+    # we later set the field variance to the desired value
+    unscaled_pow_spec = (lambda k: 1. / (1 + k/k_0) ** 4)
+    pixel_width = L/N_pixels
+
+    # Set up the geometry
+    s_space = ift.RGSpace([N_pixels, N_pixels], distances=pixel_width)
+    h_space = s_space.get_default_codomain()
+    s_var = ift.get_signal_variance(unscaled_pow_spec, h_space)
+    pow_spec = (lambda k: unscaled_pow_spec(k)/s_var*field_variance**2)
+
+    HT = ift.HarmonicTransformOperator(h_space, s_space)
+
+    # Create mock data
+
+    Sh = ift.create_power_operator(h_space, power_spectrum=pow_spec)
+    sh = Sh.draw_sample()
+
+    Rx = CustomResponse(s_space, [np.arange(0, N_pixels, 5)[:, np.newaxis],
+                                  np.arange(0, N_pixels, 2)[np.newaxis, :]])
+    ift.extra.consistency_check(Rx)
+    a = ift.Field.from_random('normal', s_space)
+    b = ift.Field.from_random('normal', Rx.target)
+    R = Rx * HT
+
+    noiseless_data = R(sh)
+    N = ift.ScalingOperator(noise_level**2, R.target)
+    n = N.draw_sample()
+
+    d = noiseless_data + n
+
+    # Wiener filter
+
+    IC = ift.GradientNormController(name="inverter", iteration_limit=1000,
+                                    tol_abs_gradnorm=0.0001)
+    inverter = ift.ConjugateGradient(controller=IC)
+    # setting up measurement precision matrix M
+    M = (ift.SandwichOperator.make(R.adjoint, Sh) + N)
+    M = ift.InversionEnabler(M, inverter)
+    m = Sh(R.adjoint(M.inverse_times(d)))
+
+    # Plotting
+    backprojection = Rx.adjoint(d)
+    reweighted_backprojection = (backprojection / backprojection.max() *
+                                 HT(sh).max())
+    zmax = max(HT(sh).max(), reweighted_backprojection.max(), HT(m).max())
+    zmin = min(HT(sh).min(), reweighted_backprojection.min(), HT(m).min())
+    plotdict = {"colormap": "Planck-like", "zmax": zmax, "zmin": zmin}
+    ift.plot(HT(sh), name="mock_signal.png", **plotdict)
+    ift.plot(backprojection, name="backprojected_data.png", **plotdict)
+    ift.plot(HT(m), name="reconstruction.png", **plotdict)

demos/wiener_filter_easy.py

 import numpy as np
 import nifty4 as ift
 
+
 if __name__ == "__main__":
     np.random.seed(43)
     # Set up physical constants
@@ -12,21 +13,25 @@ if __name__ == "__main__":
     field_variance = 2.
     # Smoothing length of response (in same unit as L)
     response_sigma = 0.01
+    # typical noise amplitude of the measurement
+    noise_level = 1.
 
     # Define resolution (pixels per dimension)
     N_pixels = 256
 
     # Set up derived constants
     k_0 = 1./correlation_length
-    # Note that field_variance**2 = a*k_0/4. for this analytic form of power
-    # spectrum
-    a = field_variance**2/k_0*4.
-    pow_spec = (lambda k: a / (1 + k/k_0) ** 4)
+    #defining a power spectrum with the right correlation length
+    #we later set the field variance to the desired value
+    unscaled_pow_spec = (lambda k: 1. / (1 + k/k_0) ** 4)
     pixel_width = L/N_pixels
 
     # Set up the geometry
     s_space = ift.RGSpace([N_pixels, N_pixels], distances=pixel_width)
     h_space = s_space.get_default_codomain()
+    s_var = ift.get_signal_variance(unscaled_pow_spec, h_space)
+    pow_spec = (lambda k: unscaled_pow_spec(k)/s_var*field_variance**2)
+
     HT = ift.HarmonicTransformOperator(h_space, s_space)
 
     # Create mock data
@@ -36,11 +41,8 @@ if __name__ == "__main__":
 
     R = HT*ift.create_harmonic_smoothing_operator((h_space,), 0,
                                                   response_sigma)
-
     noiseless_data = R(sh)
-    signal_to_noise = 1.
-    noise_amplitude = noiseless_data.val.std()/signal_to_noise
-    N = ift.ScalingOperator(noise_amplitude**2, s_space)
+    N = ift.ScalingOperator(noise_level**2, s_space)
     n = N.draw_sample()
 
     d = noiseless_data + n

nifty4/sugar.py

@@ -31,7 +31,8 @@ from .logger import logger
 __all__ = ['PS_field', 'power_analyze', 'create_power_operator',
            'create_harmonic_smoothing_operator', 'from_random',
            'full', 'empty', 'from_global_data', 'from_local_data',
-           'makeDomain', 'sqrt', 'exp', 'log', 'tanh', 'conjugate']
+           'makeDomain', 'sqrt', 'exp', 'log', 'tanh', 'conjugate',
+           'get_signal_variance']
 
 
 def PS_field(pspace, func):
@@ -41,6 +42,34 @@ def PS_field(pspace, func):
     return Field(pspace, val=data)
 
 
+def get_signal_variance(spec, space):
+    """
+    Computes how much a field with a given power spectrum will vary in space
+
+    This is a small helper function that computes how the expected variance
+    of a harmonically transformed sample of this power spectrum.
+
+    Parameters
+    ---------
+    spec: method
+        a method that takes one k-value and returns the power spectrum at that
+        location
+    space: PowerSpace or any harmonic Domain
+        If this function is given a harmonic domain, it creates the naturally
+        binned PowerSpace to that domain.
+        The field, for which the signal variance is then computed, is assumed
+        to have this PowerSpace as naturally binned PowerSpace
+    """
+    if space.harmonic:
+        space = PowerSpace(space)
+    if not isinstance(space, PowerSpace):
+        raise ValueError(
+            "space must be either a harmonic space or Power space.")
+    field = PS_field(space, spec)
+    dist = PowerDistributor(space.harmonic_partner, space)
+    k_field = dist(field)
+    return k_field.weight(2).sum()
+
 def _single_power_analyze(field, idx, binbounds):
     power_domain = PowerSpace(field.domain[idx], binbounds)
     pd = PowerDistributor(field.domain, power_domain, idx)