poisson_demo.py 4.83 KB
Newer Older
1 2 3 4
# Program to generate figures of article "Information theory for fields"
# by Torsten Ensslin, Annalen der Physik, submitted to special edition
# "Physics of Information" in April 2018

Martin Reinecke's avatar
Martin Reinecke committed
5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
import numpy as np
import nifty4 as ift
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
22 23 24 25
    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
26
    sigma_n2 = sigma_n**2
27
    L = 1.          # Total length of interval or volume the field lives on
28
    nprobes = 1000  # Number of probes for uncertainty quantification
Martin Reinecke's avatar
Martin Reinecke committed
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

    # 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
57 58 59
    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
60 61 62 63 64 65 66 67
    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
68 69
    fmask = ift.Field.from_global_data(d_domain, mask)
    Mask = ift.DiagonalOperator(fmask)
Martin Reinecke's avatar
Martin Reinecke committed
70 71
    R0 = Mask*GeoRem
    R = R0*HT
72 73 74 75 76 77 78 79

    # 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
80 81 82
    IC = ift.GradientNormController(name="inverter", iteration_limit=500,
                                    tol_abs_gradnorm=1e-3)
    inverter = ift.ConjugateGradient(controller=IC)
Martin Reinecke's avatar
Martin Reinecke committed
83
    D = (ift.SandwichOperator.make(R, N.inverse) + Phi_h.inverse).inverse
84 85 86 87
    D = ift.InversionEnabler(D, inverter, approximation=Phi_h)
    m = HT(D(j))

    # Uncertainty
Martin Reinecke's avatar
Martin Reinecke committed
88
    D = ift.SandwichOperator.make(aHT, D)  # real space propagator
89 90 91 92 93 94 95
    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
96 97 98
    c1 = (m-sig).to_global_data()
    c2 = (m+sig).to_global_data()
    plt.fill_between(x, c1, c2, color='pink', alpha=None)
99 100 101 102 103 104 105 106 107 108
    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
109 110 111

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

    # initial guess
Martin Reinecke's avatar
Martin Reinecke committed
116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
    psi0 = ift.Field.full(h_domain, 1e-7)
    energy = ift.library.PoissonEnergy(psi0, data, R0, nonlin, HT, Phi_h,
                                       inverter)
    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')