poisson_demo.py 4.89 KB
Newer Older
1 2
# Program to generate figures of article "Information theory for fields"
# by Torsten Ensslin, Annalen der Physik, submitted to special edition
3
# "Physics of Information" in April 2018, arXiv:1804.03350
Martin Reinecke's avatar
Martin Reinecke committed
4
# to get exact figure of paper comment out lines marked by "COMMENT OUT"
5

Martin Reinecke's avatar
Martin Reinecke committed
6
import numpy as np
Philipp Arras's avatar
Philipp Arras committed
7
import nifty5 as ift
Martin Reinecke's avatar
Martin Reinecke committed
8 9 10 11 12 13 14 15 16 17 18 19 20 21 22
import matplotlib.pyplot as plt


class Exp3(object):
    def __call__(self, x):
        return ift.exp(3*x)

    def derivative(self, x):
        return 3*ift.exp(3*x)


if __name__ == "__main__":
    np.random.seed(20)

    # Set up physical constants
23 24 25 26
    nu = 1.         # excitation field level
    kappa = 10.     # diffusion constant
    eps = 1e-8      # small number to tame zero mode
    sigma_n = 0.2   # noise level
Martin Reinecke's avatar
Martin Reinecke committed
27
    sigma_n2 = sigma_n**2
28
    L = 1.          # Total length of interval or volume the field lives on
29 30
    nprobes = 1000  # Number of probes for uncertainty quantification used in paper
    nprobes = 10    # COMMENT OUT TO REPRODUCE PAPER FIGURE EXACTLY
Martin Reinecke's avatar
Martin Reinecke committed
31

Martin Reinecke's avatar
Martin Reinecke committed
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
    # Define resolution (pixels per dimension)
    N_pixels = 1024

    # Define data gaps
    N1a = int(0.6*N_pixels)
    N1b = int(0.64*N_pixels)
    N2a = int(0.67*N_pixels)
    N2b = int(0.8*N_pixels)

    # Set up derived constants
    amp = nu/(2*kappa)  # spectral normalization
    pow_spec = lambda k: amp / (eps + k**2)
    lambda2 = 2*kappa*sigma_n2/nu  # resulting correlation length squared
    lambda1 = np.sqrt(lambda2)
    pixel_width = L/N_pixels
    x = np.arange(0, 1, pixel_width)

    # Set up the geometry
    s_domain = ift.RGSpace([N_pixels], distances=pixel_width)
    h_domain = s_domain.get_default_codomain()
    HT = ift.HarmonicTransformOperator(h_domain, s_domain)
    aHT = HT.adjoint

    # Create mock signal
    Phi_h = ift.create_power_operator(h_domain, power_spectrum=pow_spec)
    phi_h = Phi_h.draw_sample()
    # remove zero mode
Martin Reinecke's avatar
Martin Reinecke committed
59 60 61
    glob = phi_h.to_global_data()
    glob[0] = 0.
    phi_h = ift.Field.from_global_data(phi_h.domain, glob)
Martin Reinecke's avatar
Martin Reinecke committed
62 63 64 65 66 67 68 69
    phi = HT(phi_h)

    # Setting up an exemplary response
    GeoRem = ift.GeometryRemover(s_domain)
    d_domain = GeoRem.target[0]
    mask = np.ones(d_domain.shape)
    mask[N1a:N1b] = 0.
    mask[N2a:N2b] = 0.
Martin Reinecke's avatar
Martin Reinecke committed
70 71
    fmask = ift.Field.from_global_data(d_domain, mask)
    Mask = ift.DiagonalOperator(fmask)
Martin Reinecke's avatar
Martin Reinecke committed
72 73
    R0 = Mask*GeoRem
    R = R0*HT
74 75 76 77 78 79 80 81

    # Linear measurement scenario
    N = ift.ScalingOperator(sigma_n2, d_domain)  # Noise covariance
    n = Mask(N.draw_sample())  # seting the noise to zero in masked region
    d = R(phi_h) + n

    # Wiener filter
    j = R.adjoint_times(N.inverse_times(d))
Martin Reinecke's avatar
Martin Reinecke committed
82 83
    IC = ift.GradientNormController(name="inverter", iteration_limit=500,
                                    tol_abs_gradnorm=1e-3)
Martin Reinecke's avatar
Martin Reinecke committed
84
    D = (ift.SandwichOperator.make(R, N.inverse) + Phi_h.inverse).inverse
85
    D = ift.InversionEnabler(D, IC, approximation=Phi_h)
86 87 88
    m = HT(D(j))

    # Uncertainty
Martin Reinecke's avatar
Martin Reinecke committed
89
    D = ift.SandwichOperator.make(aHT, D)  # real space propagator
90 91 92 93 94 95 96
    Dhat = ift.probe_with_posterior_samples(D.inverse, None,
                                            nprobes=nprobes)[1]
    sig = ift.sqrt(Dhat)

    # Plotting
    x_mod = np.where(mask > 0, x, None)
    plt.rcParams["text.usetex"] = True
Martin Reinecke's avatar
Martin Reinecke committed
97 98 99
    c1 = (m-sig).to_global_data()
    c2 = (m+sig).to_global_data()
    plt.fill_between(x, c1, c2, color='pink', alpha=None)
100 101 102 103 104 105 106 107 108 109
    plt.plot(x, phi.to_global_data(), label=r"$\varphi$", color='black')
    plt.scatter(x_mod, d.to_global_data(), label=r'$d$', s=1, color='blue',
                alpha=0.5)
    plt.plot(x, m.to_global_data(), label=r'$m$', color='red')
    plt.xlim([0, L])
    plt.ylim([-1, 1])
    plt.title('Wiener filter reconstruction')
    plt.legend()
    plt.savefig('Wiener_filter.pdf')
    plt.close()
Martin Reinecke's avatar
Martin Reinecke committed
110 111 112

    nonlin = Exp3()
    lam = R0(nonlin(HT(phi_h)))
Martin Reinecke's avatar
Martin Reinecke committed
113 114
    data = ift.Field.from_local_data(
        d_domain, np.random.poisson(lam.local_data).astype(np.float64))
115 116

    # initial guess
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
117
    psi0 = ift.full(h_domain, 1e-7)
118
    energy = ift.library.PoissonEnergy(psi0, data, R0, nonlin, HT, Phi_h, IC)
Martin Reinecke's avatar
Martin Reinecke committed
119 120 121 122 123 124 125 126 127 128 129 130 131 132
    IC1 = ift.GradientNormController(name="IC1", iteration_limit=200,
                                     tol_abs_gradnorm=1e-4)
    minimizer = ift.RelaxedNewton(IC1)
    energy = minimizer(energy)[0]

    var = ift.probe_with_posterior_samples(energy.curvature, HT, nprobes)[1]
    sig = ift.sqrt(var)

    m = HT(energy.position)
    phi = HT(phi_h)
    plt.rcParams["text.usetex"] = True
    c1 = nonlin(m-sig).to_global_data()
    c2 = nonlin(m+sig).to_global_data()
    plt.fill_between(x, c1, c2, color='pink', alpha=None)
133
    plt.plot(x, nonlin(phi).to_global_data(), label=r"$e^{3\varphi}$",
Martin Reinecke's avatar
Martin Reinecke committed
134
             color='black')
Martin Reinecke's avatar
Martin Reinecke committed
135 136
    plt.scatter(x_mod, data.to_global_data(), label=r'$d$', s=1, color='blue',
                alpha=0.5)
137 138
    plt.plot(x, nonlin(m).to_global_data(),
             label=r'$e^{3\varphi_\mathrm{cl}}$', color='red')
Martin Reinecke's avatar
Martin Reinecke committed
139
    plt.xlim([0, L])
140
    plt.ylim([-0.1, 7.5])
Martin Reinecke's avatar
Martin Reinecke committed
141 142
    plt.title('Poisson log-normal reconstruction')
    plt.legend()
143
    plt.savefig('Poisson.pdf')