amplitude_model.py 4.17 KB
Newer Older
1 2
import numpy as np

Martin Reinecke's avatar
Martin Reinecke committed
3 4
from ..domains.power_space import PowerSpace
from ..domains.unstructured_domain import UnstructuredDomain
5
from ..field import Field
Martin Reinecke's avatar
Martin Reinecke committed
6
from ..multi.multi_field import MultiField
7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
from ..sugar import makeOp, sqrt


def _ceps_kernel(dof_space, k, a, k0):
    return a**2/(1+(k/(k0*dof_space.bindistances[0]))**2)**2


def make_amplitude_model(s_space, Npixdof, ceps_a, ceps_k, sm, sv, im, iv,
                         keys=['tau', 'phi']):
    '''
    Method for construction of amplitude model
    Computes a smooth power spectrum.
    Output lives in PowerSpace.

    Parameters
    ----------

    Npixdof : #pix in dof_space

Martin Reinecke's avatar
Martin Reinecke committed
26
    ceps_a, ceps_k0 : Smoothness parameters in ceps_kernel
27 28 29
                        eg. ceps_kernel(k) = (a/(1+(k/k0)**2))**2
                        a = ceps_a,  k0 = ceps_k0

Martin Reinecke's avatar
Martin Reinecke committed
30
    sm, sv : slope_mean = expected exponent of power law (e.g. -4),
31 32 33 34
                slope_variance (default=1)

    im, iv : y-intercept_mean, y-intercept_variance  of power_slope
    '''
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
35 36 37 38 39 40 41
    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 ..models.variable import Variable
    from ..models.constant import Constant
    from ..models.local_nonlinearity import PointwiseExponential
42

43 44 45 46 47 48 49 50 51 52 53 54 55
    h_space = s_space.get_default_codomain()
    p_space = PowerSpace(h_space)
    exp_transform = ExpTransform(p_space, Npixdof)
    logk_space = exp_transform.domain[0]
    dof_space = logk_space.get_default_codomain()
    param_space = UnstructuredDomain(2)
    qht = QHTOperator(dof_space, logk_space)
    sym = SymmetrizingOperator(logk_space)

    phi_mean = np.array([sm, im])
    phi_sig = np.array([sv, iv])

    slope = SlopeOperator(param_space, logk_space, phi_sig)
56
    norm_phi_mean = Field.from_global_data(param_space, phi_mean/phi_sig)
57 58 59 60 61 62 63 64 65 66 67 68 69

    fields = {keys[0]: Field.from_random('normal', dof_space),
              keys[1]: Field.from_random('normal', param_space)}

    position = MultiField(fields)

    dof_space = position[keys[0]].domain[0]
    kern = lambda k: _ceps_kernel(dof_space, k, ceps_a, ceps_k)
    cepstrum = create_cepstrum_amplitude_field(dof_space, kern)

    ceps = makeOp(sqrt(cepstrum))
    smooth_op = sym * qht * ceps
    smooth_spec = smooth_op(Variable(position)[keys[0]])
70
    phi = Variable(position)[keys[1]] + Constant(position, norm_phi_mean)
71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100
    linear_spec = slope(phi)
    loglog_spec = smooth_spec + linear_spec
    xlog_ampl = PointwiseExponential(0.5*loglog_spec)

    internals = {'loglog_spec': loglog_spec,
                 'qht': qht,
                 'ceps': ceps,
                 'norm_phi_mean': norm_phi_mean}
    return exp_transform(xlog_ampl), internals


def create_cepstrum_amplitude_field(domain, cepstrum):
    """Creates a ...
    Writes the sum of all modes into the zero-mode.

    Parameters
    ----------
    domain: ???
        ???
    cepstrum: Callable
        ???
    """

    dim = len(domain.shape)
    dist = domain.bindistances
    shape = domain.shape

    # Prepare q_array
    q_array = np.zeros((dim,) + shape)
    if dim == 1:
101
        ks = domain.get_k_length_array().to_global_data()
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127
        q_array = np.array([ks])
    else:
        for i in range(dim):
            ks = np.minimum(shape[i] - np.arange(shape[i]) +
                            1, np.arange(shape[i])) * dist[i]
            fst_dims = (1,) * i
            lst_dims = (1,) * (dim - i - 1)
            q_array[i] += ks.reshape(fst_dims + (shape[i],) + lst_dims)

    # Fill cepstrum field (all non-zero modes)
    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)

    # Fill cepstrum field (zero-mode subspaces)
    for i in range(dim):
        # Prepare indices
        fst_dims = (slice(None),) * i
        lst_dims = (slice(None),) * (dim - i - 1)
        sl = fst_dims + (slice(1, None),) + lst_dims
        sl2 = fst_dims + (0,) + lst_dims

        # Do summation
        cepstrum_field[sl2] = np.sum(cepstrum_field[sl], axis=i)

128
    return Field.from_global_data(domain, cepstrum_field)