Implemented Uniformoperator and generalised InterpolationOperator to supply the correct jacobian

nifty6/__init__.py

@@ -72,7 +72,7 @@ from .sugar import *
 
 from .plot import Plot
 
-from .library.special_distributions import InverseGammaOperator
+from .library.special_distributions import InverseGammaOperator, UniformOperator
 from .library.los_response import LOSResponse
 from .library.dynamic_operator import (dynamic_operator,
                                        dynamic_lightcone_operator)

nifty6/library/special_distributions.py

@@ -23,20 +23,59 @@ from ..field import Field
 from ..linearization import Linearization
 from ..operators.operator import Operator
 from ..sugar import makeOp
-
+from .. import Adder
 
 class _InterpolationOperator(Operator):
-    def __init__(self, domain, func, xmin, xmax, delta, table_func=None, inverse_table_func=None):
+    """
+    Calculates a function pointwise on a field by interpolation.
+    Can be supplied with a table_func to increase the interpolation accuracy,
+    Best results are achieved when table_func(func) is roughly linear.
+
+    Parameters
+    ----------
+    domain : Domain, tuple of Domain or DomainTuple
+        The domain on which the field shall be defined. This is at the same
+        time the domain and the target of the operator.
+    func : function
+        The function which is applied on the field.
+    xmin : float
+        The smallest value for which func will be evaluated.
+    xmax : float
+        The largest value for which func will be evaluated.
+    delta : float
+        Distance between sampling points for linear interpolation.
+    table_func : {'None', 'exp', 'log', 'power'}
+
+    exponent : float
+        This is only used by the 'power' table_func.
+    """
+    def __init__(self, domain, func, xmin, xmax, delta, table_func=None, exponent=None):
         self._domain = self._target = DomainTuple.make(domain)
         self._xmin, self._xmax = float(xmin), float(xmax)
         self._d = float(delta)
         self._xs = np.arange(xmin, xmax+2*self._d, self._d)
-        if table_func is not None:
-            foo = func(np.random.randn(10))
-            np.testing.assert_allclose(foo, inverse_table_func(table_func(foo)))
-        self._table = table_func(func(self._xs))
-        self._deriv = (self._table[1:]-self._table[:-1]) / self._d
-        self._inv_table_func = inverse_table_func
+        self._table = func(self._xs)
+        self._transform = table_func is not None
+        self._args = []
+        if exponent is not None and table_func is not 'power':
+            raise Exception("exponent is only used when table_func is 'power'.")
+        if table_func is None:
+            pass
+        elif table_func == 'exp':
+            self._table = np.exp(self._table)
+            self._inv_table_func = 'log'
+        elif table_func == 'log':
+            self._table = np.log(self._table)
+            self._inv_table_func = 'exp'
+        elif table_func == 'power':
+            if not np.isscalar(exponent):
+                return NotImplemented
+            self._table = np.power(self._table, exponent)
+            self._inv_table_func = '__pow__'
+            self._args = [np.power(float(exponent), -1)]
+        else:
+            return NotImplemented
+        self._deriv = (self._table[1:] - self._table[:-1]) / self._d
 
     def apply(self, x):
         self._check_input(x)
@@ -45,16 +84,21 @@ class _InterpolationOperator(Operator):
         val = (np.clip(xval, self._xmin, self._xmax) - self._xmin) / self._d
         fi = np.floor(val).astype(int)
         w = val - fi
-        res = self._inv_table_func((1-w)*self._table[fi] + w*self._table[fi+1])
+        res = (1-w)*self._table[fi] + w*self._table[fi+1]
         resfld = Field(self._domain, res)
         if not lin:
+            if self._transform:
+                resfld = getattr(resfld, self._inv_table_func)(*self._args)
             return resfld
-        jac = makeOp(Field(self._domain, self._deriv[fi]*res))
-        return x.new(resfld, jac)
+        lin = Linearization.make_var(resfld)
+        if self._transform:
+            lin = getattr(lin, self._inv_table_func)(*self._args)
+        jac = makeOp(Field(self._domain, self._deriv[fi]))
+        return x.new(lin.val, lin.jac@jac)
 
 
 def InverseGammaOperator(domain, alpha, q, delta=0.001):
-    """Transforms a Gaussian into an inverse gamma distribution.
+    """Transforms a Gaussian with unit covariance and zero mean into an inverse gamma distribution.
 
     The pdf of the inverse gamma distribution is defined as follows:
 
@@ -67,7 +111,7 @@ def InverseGammaOperator(domain, alpha, q, delta=0.001):
     The mode is :math:`q / (\\alpha + 1)`.
 
     This transformation is implemented as a linear interpolation which maps a
-    Gaussian onto a inverse gamma distribution.
+    Gaussian onto an inverse gamma distribution.
 
     Parameters
     ----------
@@ -76,14 +120,39 @@ def InverseGammaOperator(domain, alpha, q, delta=0.001):
         time the domain and the target of the operator.
     alpha : float
         The alpha-parameter of the inverse-gamma distribution.
-    q : float
+    q : float or Field
         The q-parameter of the inverse-gamma distribution.
     delta : float
-        distance between sampling points for linear interpolation.
+        Distance between sampling points for linear interpolation.
     """
     op = _InterpolationOperator(domain,
                                 lambda x: invgamma.ppf(norm.cdf(x), float(alpha)),
-                                -8.2, 8.2, delta, np.log, np.exp)
+                                -8.2, 8.2, delta, 'log')
     if np.isscalar(q):
         return op.scale(q)
     return makeOp(q) @ op
+
+def UniformOperator(domain, loc=0, scale=1, delta=1e-3):
+    """
+    Transforms a Gaussian with unit covariance and zero mean into a uniform distribution.
+    The uniform distribution's support is ``[loc, loc + scale]``.
+
+    Parameters
+    ----------
+    domain : Domain, tuple of Domain or DomainTuple
+        The domain on which the field shall be defined. This is at the same
+        time the domain and the target of the operator.
+    loc: float or Field
+
+    scale: float or Field
+
+    delta : float
+        Distance between sampling points for linear interpolation.
+    """
+    op = _InterpolationOperator(domain,
+                                lambda x: norm.cdf(x),
+                                -8.2, 8.2, delta)
+    loc = Adder(loc, domain=domain)
+    if np.isscalar(scale):
+        return loc(op.scale(scale))
+    return loc(makeOp(scale) @ op)