getting_started_2.py 2.37 KB
Newer Older
Natalia Porqueres's avatar
Natalia Porqueres committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
import nifty5 as ift
import sys
import numpy as np
import global_newton as gn
from nifty5.library.nonlinearities import Exponential

np.random.seed(42)

N = 2
Nsamples = 5

s_space = ift.RGSpace(10024)
h_space = s_space.get_default_codomain()

domain = ift.MultiDomain.make({'xi': h_space})
position = ift.from_random('normal', domain)
HT = ift.HarmonicTransformOperator(h_space, s_space)

def sqrtpspec(k):
    return 16. / (20.+k**2)


# Define amplitude model
p_space = ift.PowerSpace(h_space)
pd = ift.PowerDistributor(h_space, p_space)
a = ift.PS_field(p_space, sqrtpspec)
A = pd(a)

# Define sky model
xi = ift.Variable(position)['xi']
logsky_h = xi * A
logsky = HT(logsky_h)
nonlin = Exponential()
sky = ift.PointwiseExponential(logsky)
R = ift.ScalingOperator(1., s_space)
d_space = R.target[0]
lamb = R(sky)

# Generate mock data
MOCK_POSITION = ift.from_random('normal', lamb.position.domain)
data = np.random.poisson(lamb.at(MOCK_POSITION).value.val.astype(np.float64))
data = ift.Field.from_local_data(d_space, data)

# Define Hamiltonian
position = ift.from_random('normal', lamb.position.domain)
likelihood = ift.library.PoissonLogLikelihood(lamb, data)

ic_cg = ift.GradientNormController(iteration_limit=50)
ic_samps = ift.GradientNormController(iteration_limit=500, tol_abs_gradnorm=1e-4)

ic_newton = ift.GradientNormController(name='Newton', tol_abs_gradnorm=1e-3)
minimizer = ift.RelaxedNewton(ic_newton)

H = gn.Hamiltonian(likelihood, ic_cg, ic_samps)

for _ in range(N):
    samples = [H.curvature.draw_sample(from_inverse=True)
               for _ in range(Nsamples)]
    sc_samplesky = ift.StatCalculator()
    for s in samples:
        sc_samplesky.add(sky.at(s+position).value)
    ift.plot(sc_samplesky.mean, name='sample_mean.png')

    KL = gn.SampledKullbachLeiblerDivergence(H, samples, ic_cg)
    KL, convergence = minimizer(KL)
    position = KL.position

# Plot results
E = KL
l1 = lamb.at(E.position).value
l2 = lamb.at(MOCK_POSITION).value
ift.plot([data, l2, l1], name="poisson.png",
         label=['Data', 'Mock signal', 'Reconstruction'],
         alpha=[.5, 1, 1])
if power_spectrum_estimation:
    a_mock = a.at(MOCK_POSITION).value
    a_recon = a.at(E.position).value
else:
    a_mock = a
    a_recon = a
ift.plot([a_mock**2, a_recon**2, ift.power_analyze(logsky_h.at(E.position).value)],
         name='power_spectrum.png', label=['Mock', 'Reconstruction', 'power_analyze'])