critical_filtering.py 4.65 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
import nifty4 as ift
Martin Reinecke's avatar
Martin Reinecke committed
2
from nifty4.library.nonlinearities import Linear, Tanh, Exponential
3
import numpy as np
4
np.random.seed(42)
Martin Reinecke's avatar
adjust  
Martin Reinecke committed
5

Jakob Knollmueller's avatar
Jakob Knollmueller committed
6

7 8 9 10 11 12 13 14 15 16
def adjust_zero_mode(m0, t0):
    mtmp = m0.to_global_data().copy()
    zero_position = len(m0.shape)*(0,)
    zero_mode = mtmp[zero_position]
    mtmp[zero_position] = zero_mode / abs(zero_mode)
    ttmp = t0.to_global_data().copy()
    ttmp[0] += 2 * np.log(abs(zero_mode))
    return (ift.Field.from_global_data(m0.domain, mtmp),
            ift.Field.from_global_data(t0.domain, ttmp))

17
if __name__ == "__main__":
18 19 20 21 22

    noise_level = 1.
    p_spec = (lambda k: (.5 / (k + 1) ** 3))

    nonlinearity = Linear()
23
    # Set up position space
24 25
    s_space = ift.RGSpace((128, 128))
    h_space = s_space.get_default_codomain()
26 27

    # Define harmonic transformation and associated harmonic space
28
    HT = ift.HarmonicTransformOperator(h_space, target=s_space)
29

30
    # Setting up power space
Martin Reinecke's avatar
adjust  
Martin Reinecke committed
31 32 33
    p_space = ift.PowerSpace(h_space,
                             binbounds=ift.PowerSpace.useful_binbounds(
                                 h_space, logarithmic=True))
Martin Reinecke's avatar
Martin Reinecke committed
34
    s_spec = ift.Field.full(p_space, 1e-5)
35 36 37 38 39 40 41 42 43 44 45 46
    # Choosing the prior correlation structure and defining
    # correlation operator
    p = ift.PS_field(p_space, p_spec)
    log_p = ift.log(p)
    S = ift.create_power_operator(h_space, power_spectrum=s_spec)

    # Drawing a sample sh from the prior distribution in harmonic space
    sh = ift.power_synthesize(s_spec)

    # Choosing the measurement instrument
    # Instrument = SmoothingOperator(s_space, sigma=0.01)
    mask = np.ones(s_space.shape)
Martin Reinecke's avatar
Martin Reinecke committed
47
    #mask[6000:8000] = 0.
Martin Reinecke's avatar
Martin Reinecke committed
48
    mask[30:70,30:70] = 0.
49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
    mask = ift.Field.from_global_data(s_space, mask)

    MaskOperator = ift.DiagonalOperator(mask)
    R = ift.GeometryRemover(s_space)
    R = R*MaskOperator
    # R = R*HT
    # R = R * ift.create_harmonic_smoothing_operator((harmonic_space,), 0,
    #                                                response_sigma)
    MeasurementOperator = R

    d_space = MeasurementOperator.target

    Distributor = ift.PowerDistributor(target=h_space, power_space=p_space)
    power = Distributor(ift.exp(0.5*log_p))
    # Creating the mock data
    true_sky = nonlinearity(HT(power*sh))
    noiseless_data = MeasurementOperator(true_sky)
    noise_amplitude = noiseless_data.val.std()*noise_level
    N = ift.ScalingOperator(noise_amplitude**2, d_space)
    n = ift.Field.from_random(
        domain=d_space, random_type='normal',
        std=noise_amplitude, mean=0)
    # Creating the mock data
    d = noiseless_data + n
73

74 75 76
    m0 = ift.power_synthesize(ift.Field.full(p_space, 1e-7))
    t0 = ift.Field.full(p_space, -4.)
    power0 = Distributor.times(ift.exp(0.5 * t0))
77

Martin Reinecke's avatar
Martin Reinecke committed
78 79 80 81 82 83 84
    plotdict = {"colormap": "Planck-like"}
    zmin = true_sky.min()
    zmax = true_sky.max()
    ift.plot(true_sky, title="True sky", name="true_sky.png", **plotdict)
    ift.plot(MeasurementOperator.adjoint_times(d), title="Data",
             name="data.png", **plotdict)

85 86 87 88
    IC1 = ift.GradientNormController(name="IC1", iteration_limit=100,
                                     tol_abs_gradnorm=1e-3)
    LS = ift.LineSearchStrongWolfe(c2=0.02)
    minimizer = ift.RelaxedNewton(IC1, line_searcher=LS)
89

90 91 92
    ICI = ift.GradientNormController(iteration_limit=500,
                                     tol_abs_gradnorm=1e-3)
    inverter = ift.ConjugateGradient(controller=ICI)
93

94 95 96 97 98
    for i in range(20):
        power0 = Distributor(ift.exp(0.5*t0))
        map0_energy = ift.library.NonlinearWienerFilterEnergy(
            m0, d, MeasurementOperator, nonlinearity, HT, power0, N, S,
            inverter=inverter)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
99

100 101 102
        # Minimization with chosen minimizer
        map0_energy, convergence = minimizer(map0_energy)
        m0 = map0_energy.position
103

104 105
        # Updating parameters for correlation structure reconstruction
        D0 = map0_energy.curvature
Jakob Knollmueller's avatar
Jakob Knollmueller committed
106

107 108 109 110 111
        # Initializing the power energy with updated parameters
        power0_energy = ift.library.NonlinearPowerEnergy(
            position=t0, d=d, N=N, xi=m0, D=D0, ht=HT,
            Instrument=MeasurementOperator, nonlinearity=nonlinearity,
            Distribution=Distributor, sigma=1., samples=2, inverter=inverter)
112

113
        power0_energy = minimizer(power0_energy)[0]
Martin Reinecke's avatar
Martin Reinecke committed
114

115 116 117 118 119 120 121
        # Setting new power spectrum
        t0 = power0_energy.position

        # break degeneracy between amplitude and excitation by setting
        # excitation monopole to 1
        m0, t0 = adjust_zero_mode(m0, t0)

Martin Reinecke's avatar
Martin Reinecke committed
122
    ift.plot(nonlinearity(HT(power0*m0)), title="Reconstructed sky",
Martin Reinecke's avatar
Martin Reinecke committed
123 124 125 126
             name="reconstructed_sky.png", zmin=zmin, zmax=zmax, **plotdict)
    ymin = np.min(p.to_global_data())
    ift.plot([ift.exp(t0),p], title="Power spectra",
             name="ps.png", ymin=ymin, **plotdict)