Split AmplitudeModel into several models

nifty5/library/amplitude_model.py

@@ -29,7 +29,7 @@ from ..sugar import makeOp, sqrt
 
 
 def _ceps_kernel(dof_space, k, a, k0):
-    return a**2/(1+(k/k0)**2)**2
+    return a**2/(1 + (k/k0)**2)**2
 
 
 def create_cepstrum_amplitude_field(domain, cepstrum):
@@ -51,7 +51,7 @@ def create_cepstrum_amplitude_field(domain, cepstrum):
     q_array = domain.get_k_array()
 
     # Fill cepstrum field (all non-zero modes)
-    no_zero_modes = (slice(1, None),) * dim
+    no_zero_modes = (slice(1, None),)*dim
     ks = q_array[(slice(None),) + no_zero_modes]
     cepstrum_field = np.zeros(shape)
     cepstrum_field[no_zero_modes] = cepstrum(ks)
@@ -69,6 +69,76 @@ def create_cepstrum_amplitude_field(domain, cepstrum):
     return Field.from_global_data(domain, cepstrum_field)
 
 
+class CepstrumModel(Operator):
+    '''
+    Parameters
+    ----------
+    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, logk_space, ceps_a, ceps_k):
+        from ..operators.qht_operator import QHTOperator
+        from ..operators.symmetrizing_operator import SymmetrizingOperator
+        qht = QHTOperator(target=logk_space)
+        dof_space = qht.domain[0]
+        sym = SymmetrizingOperator(logk_space)
+        kern = lambda k: _ceps_kernel(dof_space, k, ceps_a, ceps_k)
+        cepstrum = create_cepstrum_amplitude_field(dof_space, kern)
+        self._qht = qht
+        self._ceps = makeOp(sqrt(cepstrum))
+        self._op = sym(qht(makeOp(sqrt(cepstrum))))
+        self._domain, self._target = self._op.domain, self._op.target
+
+    def apply(self, x):
+        self._check_input(x)
+        return self._op(x)
+
+    @property
+    def qht(self):
+        return self._qht
+
+    @property
+    def ceps(self):
+        return self._ceps
+
+
+class SlopeModel(Operator):
+    '''
+    Parameters
+    ----------
+
+    sm, sv : slope_mean = expected exponent of power law (e.g. -4),
+                slope_variance (default=1)
+
+    im, iv : y-intercept_mean, y-intercept_variance  of power_slope
+    '''
+
+    def __init__(self, logk_space, sm, sv, im, iv):
+        from ..operators.slope_operator import SlopeOperator
+
+        phi_mean = np.array([sm, im + sm*logk_space.t_0[0]])
+        phi_sig = np.array([sv, iv])
+
+        self._slope = SlopeOperator(logk_space)
+        self._slope = self._slope(
+            makeOp(Field.from_global_data(self._slope.domain, phi_sig)))
+        self._norm_phi_mean = Field.from_global_data(self._slope.domain,
+                                                     phi_mean/phi_sig)
+
+        self._domain = self._slope.domain
+        self._target = self._slope.target
+
+    def apply(self, x):
+        self._check_input(x)
+        return self._slope(x + self._norm_phi_mean)
+
+    @property
+    def norm_phi_mean(self):
+        return self._norm_phi_mean
+
+
 class AmplitudeModel(Operator):
     '''
     Computes a smooth power spectrum.
@@ -88,49 +158,31 @@ class AmplitudeModel(Operator):
 
     im, iv : y-intercept_mean, y-intercept_variance  of power_slope
     '''
-    def __init__(self, s_space, Npixdof, ceps_a, ceps_k, sm, sv, im, iv,
-                 keys=['tau', 'phi']):
+    def __init__(self, s_space, Npixdof, ceps_a, ceps_k, sm, sv, im, iv, keys=['tau', 'phi']):
         from ..operators.exp_transform import ExpTransform
-        from ..operators.qht_operator import QHTOperator
-        from ..operators.slope_operator import SlopeOperator
-        from ..operators.symmetrizing_operator import SymmetrizingOperator
+        from ..operators.simple_linear_operators import FieldAdapter
+        from ..operators.scaling_operator import ScalingOperator
 
         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]
-        sym = SymmetrizingOperator(logk_space)
-
-        phi_mean = np.array([sm, im+sm*logk_space.t_0[0]])
-        phi_sig = np.array([sv, iv])
+        et = ExpTransform(PowerSpace(h_space), Npixdof)
+        logk_space = et.domain[0]
 
-        self._slope = SlopeOperator(logk_space)
-        self._slope = self._slope(makeOp(Field.from_global_data(
-                                    self._slope.domain,phi_sig)))
-        self._norm_phi_mean = Field.from_global_data(self._slope.domain,
-                                                     phi_mean/phi_sig)
+        smooth = CepstrumModel(logk_space, ceps_a, ceps_k)
+        linear = SlopeModel(logk_space, sm, sv, im, iv)
 
-        self._domain = MultiDomain.make({keys[0]: dof_space,
-                                         keys[1]: self._slope.domain})
-        self._target = self._exp_transform.target
+        self._qht, self._ceps = smooth.qht, smooth.ceps
+        self._norm_phi_mean = linear.norm_phi_mean
 
-        kern = lambda k: _ceps_kernel(dof_space, k, ceps_a, ceps_k)
-        cepstrum = create_cepstrum_amplitude_field(dof_space, kern)
-
-        self._smooth_op = sym(qht(makeOp(sqrt(cepstrum))))
-        self._keys = tuple(keys)
+        fa_smooth = FieldAdapter(smooth.domain, keys[0])
+        fa_linear = FieldAdapter(linear.domain, keys[1])
 
-        self._qht = qht
-        self._ceps = makeOp(sqrt(cepstrum))
+        fac = ScalingOperator(0.5, smooth.target)
+        self._op = et((fac(smooth(fa_smooth) + linear(fa_linear))).exp())
+        self._domain, self._target = self._op.domain, self._op.target
 
     def apply(self, x):
         self._check_input(x)
-        smooth_spec = self._smooth_op(x[self._keys[0]])
-        phi = x[self._keys[1]] + self._norm_phi_mean
-        linear_spec = self._slope(phi)
-        loglog_spec = smooth_spec + linear_spec
-        return self._exp_transform((0.5*loglog_spec)).exp()
+        return self._op(x)
 
     @property
     def qht(self):