Commit 81cab5d6 by Philipp Arras

### Wiener filter demo

Works only with python2 and on commit 1d10be46.
parent 63eb16f4
Pipeline #24279 passed with stage
in 9 minutes and 15 seconds
 { "cells": [ { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# A NIFTy demonstration" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## IFT: Big Picture\n", "IFT starting point:\n", "\n", "$$d = Rs+n$$\n", "\n", "Typically, $s$ continuous field, $d$ discrete data vector. Particularily, $R$ is not invertible.\n", "\n", "IFT aims at **inverting** the above uninvertible problem in the **best possible way** using Bayesian statistics.\n", "\n", "\n", "## NIFTy\n", "\n", "NIFTy (Numerical Information Field Theory, en. raffiniert) is a Python framework in which IFT problems can be tackeled easily.\n", "\n", "Main Interfaces:\n", "\n", "- **Spaces**: Cartesian, 2-Spheres (Healpix, Gauss-Legendre) and their respective harmonic spaces.\n", "- **Fields**: Defined on spaces.\n", "- **Operators**: Acting on spaces." ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## Wiener Filter: Formulae\n", "\n", "### Assumptions\n", "\n", "- $d=Rs+n$, $R$ linear operator.\n", "- $\\mathcal P (s) = \\mathcal G (s,S)$, $\\mathcal P (n) = \\mathcal G (n,N)$ where $S, N$ are positive definite matrices.\n", "\n", "### Posterior\n", "The Posterior is given by:\n", "\n", "$$\\mathcal P (s|d) \\propto P(s,d) = \\mathcal G(d-Rs,N) \\,\\mathcal G(s,S) \\propto \\mathcal G (m,D)$$\n", "\n", "where\n", "\\begin{align}\n", "m &= Dj \\\\\n", "D^{-1}&= (S^{-1} +R^\\dagger N^{-1} R )\\\\\n", "j &= R^\\dagger N^{-1} d\n", "\\end{align}\n", "\n", "Let us implement this in NIFTy!" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## Wiener Filter: Example\n", "\n", "- One-dimensional signal with powerspectrum: $$P(k) = P_0\\,\\left(1+\\left(\\frac{k}{k_0}\\right)^2\\right)^{-\\gamma /2},$$\n", "with $P_0 = 0.2, k_0 = 5, \\gamma = 4$. Recall: $P(k)$ defines an isotropic and homogeneous $S$.\n", "- $N = 0.5 \\cdot \\text{id}$.\n", "- Number data points $N_{pix} = 512$.\n", "- Response operator:\n", "$$R_x=\\begin{pmatrix} \\delta(x-0)\\\\\\delta(x-1)\\\\\\ldots\\\\ \\delta(x-511) \\end{pmatrix}.$$\n", "However, the signal space is also discrete on the computer and $R = \\text{id}$." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "N_pixels = 512 # Number of pixels\n", "sigma2 = .5 # Noise variance\n", "\n", "def pow_spec(k):\n", " P0, k0, gamma = [.2, 5, 6]\n", " return P0 * (1. + (k/k0)**2)**(- gamma / 2)" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "## Wiener Filter: Implementation" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "-" } }, "source": [ "### Import Modules" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "from nifty import (DiagonalOperator, EndomorphicOperator, FFTOperator, Field,\n", " InvertibleOperatorMixin, PowerSpace, RGSpace,\n", " create_power_operator, SmoothingOperator, DiagonalProberMixin, Prober)" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Implement Propagator" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "class PropagatorOperator(InvertibleOperatorMixin, EndomorphicOperator):\n", " def __init__(self, R, N, Sh, default_spaces=None):\n", " super(PropagatorOperator, self).__init__(default_spaces=default_spaces,\n", " preconditioner=lambda x : fft.adjoint_times(Sh.times(fft.times(x))))\n", "\n", " self.R = R\n", " self.N = N\n", " self.Sh = Sh\n", " self._domain = R.domain\n", " self.fft = FFTOperator(domain=R.domain, target=Sh.domain)\n", "\n", " def _inverse_times(self, x, spaces, x0=None):\n", " return self.R.adjoint_times(self.N.inverse_times(self.R(x))) \\\n", " + self.fft.adjoint_times(self.Sh.inverse_times(self.fft(x)))\n", "\n", " @property\n", " def domain(self):\n", " return self._domain\n", "\n", " @property\n", " def unitary(self):\n", " return False\n", "\n", " @property\n", " def symmetric(self):\n", " return False\n", "\n", " @property\n", " def self_adjoint(self):\n", " return True" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "skip" } }, "source": [ "### Conjugate Gradient Preconditioning\n", "\n", "- $D$ is defined via:\n", "$$D^{-1} = \\mathcal F^\\dagger S_h^{-1}\\mathcal F + R^\\dagger N^{-1} R.$$\n", "In the end, we want to apply $D$ to $j$, i.e. we need the inverse action of $D^{-1}$. This is done numerically (algorithm: *Conjugate Gradient*). \n", "\n", "- One can define the *condition number* of a non-singular and normal matrix $A$:\n", "$$\\kappa (A) := \\frac{|\\lambda_{\\text{max}}|}{|\\lambda_{\\text{min}}|},$$\n", "where $\\lambda_{\\text{max}}$ and $\\lambda_{\\text{min}}$ are the largest and smallest eigenvalue of $A$, respectively.\n", "\n", "- The larger $\\kappa$ the slower Conjugate Gradient.\n", "\n", "- By default, conjugate gradient solves: $D^{-1} m = j$ for $m$, where $D^{-1}$ can be bad conditioned. If one knows a non-singular matrix $T$ for which $TD^{-1}$ is better conditioned, one can solve the equivalent problem:\n", "$$\\tilde A m = \\tilde j,$$\n", "where $\\tilde A = T D^{-1}$ and $\\tilde j = Tj$.\n", "\n", "- In our case $S^{-1}$ is responsible for the bad conditioning of $D$ depending on the chosen power spectrum. Thus, we choose\n", "\n", "$$T = \\mathcal F^\\dagger S_h^{-1} \\mathcal F.$$" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Generate Mock data\n", "\n", "- Generate a field $s$ and $n$ with given covariances.\n", "- Calculate $d$." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "s_space = RGSpace(N_pixels)\n", "fft = FFTOperator(s_space)\n", "h_space = fft.target[0]\n", "p_space = PowerSpace(h_space)\n", "\n", "\n", "# Operators\n", "Sh = create_power_operator(h_space, power_spectrum=pow_spec)\n", "N = DiagonalOperator(s_space, diagonal=sigma2, bare=True)\n", "R = DiagonalOperator(s_space, diagonal=1.)\n", "D = PropagatorOperator(R=R, N=N, Sh=Sh)\n", "\n", "# Fields and data\n", "sh = Field(p_space, val=pow_spec).power_synthesize(real_signal=True)\n", "s = fft.adjoint_times(sh)\n", "n = Field.from_random(domain=s_space, random_type='normal',\n", " std=np.sqrt(sigma2), mean=0)\n", "d = R(s) + n\n", "j = R.adjoint_times(N.inverse_times(d))" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Run Wiener Filter" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "m = D(j)" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Create Power Spectra of Signal and Reconstruction" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "s_power = sh.power_analyze()\n", "m_power = fft(m).power_analyze()\n", "s_power_data = s_power.val.get_full_data().real\n", "m_power_data = m_power.val.get_full_data().real\n", "\n", "# Get signal data and reconstruction data\n", "s_data = s.val.get_full_data().real\n", "m_data = m.val.get_full_data().real\n", "\n", "d_data = d.val.get_full_data().real" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Signal Reconstruction" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "plt.plot(s_data, 'k', label=\"Signal\", alpha=.5, linewidth=.5)\n", "plt.plot(d_data, 'k+', label=\"Data\")\n", "plt.plot(m_data, 'r', label=\"Reconstruction\")\n", "plt.title(\"Reconstruction\")\n", "plt.legend()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "subslide" } }, "outputs": [], "source": [ "plt.figure()\n", "plt.plot(s_data - s_data, 'k', label=\"Signal\", alpha=.5, linewidth=.5)\n", "plt.plot(d_data - s_data, 'k+', label=\"Data\")\n", "plt.plot(m_data - s_data, 'r', label=\"Reconstruction\")\n", "plt.axhspan(-np.sqrt(sigma2),np.sqrt(sigma2), facecolor='0.9', alpha=.5)\n", "plt.title(\"Residuals\")\n", "plt.legend()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Power Spectrum" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "plt.loglog()\n", "plt.xlim(1, int(N_pixels/2))\n", "ymin = min(m_power_data)\n", "plt.ylim(ymin, 1)\n", "xs = np.arange(1,int(N_pixels/2),.1)\n", "plt.plot(xs, pow_spec(xs), label=\"True Power Spectrum\", linewidth=.7, color='k')\n", "plt.plot(s_power_data, 'k', label=\"Signal\", alpha=.5, linewidth=.5)\n", "plt.plot(m_power_data, 'r', label=\"Reconstruction\")\n", "plt.axhline(sigma2 / N_pixels, color=\"k\", linestyle='--', label=\"Noise level\", alpha=.5)\n", "plt.axhspan(sigma2 / N_pixels, ymin, facecolor='0.9', alpha=.5)\n", "plt.title(\"Power Spectrum\")\n", "plt.legend()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "## Wiener Filter on Incomplete Data" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "skip" } }, "outputs": [], "source": [ "# Operators\n", "Sh = create_power_operator(h_space, power_spectrum=pow_spec)\n", "N = DiagonalOperator(s_space, diagonal=sigma2, bare=True)\n", "# R is defined below\n", "\n", "# Fields\n", "sh = Field(p_space, val=pow_spec).power_synthesize(real_signal=True)\n", "s = fft.adjoint_times(sh)\n", "n = Field.from_random(domain=s_space, random_type='normal',\n", " std=np.sqrt(sigma2), mean=0)" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "skip" } }, "source": [ "### Partially Lose Data" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "-" } }, "outputs": [], "source": [ "l = int(N_pixels * 0.2)\n", "h = int(N_pixels * 0.2 * 4)\n", "\n", "mask = Field(s_space, val=1)\n", "mask.val[ l : h] = 0\n", "\n", "R = DiagonalOperator(s_space, diagonal = mask)\n", "n.val[l:h] = 0\n", "\n", "d = R(s) + n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "skip" } }, "outputs": [], "source": [ "D = PropagatorOperator(R=R, N=N, Sh=Sh)\n", "j = R.adjoint_times(N.inverse_times(d))\n", "m = D(j)" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Compute Uncertainty\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "class DiagonalProber(DiagonalProberMixin, Prober):\n", " def __init__(self, *args, **kwargs):\n", " super(DiagonalProber,self).__init__(*args, **kwargs)\n", "\n", "diagProber = DiagonalProber(domain=s_space, probe_dtype=np.complex, probe_count=200)\n", "diagProber(D)\n", "m_var = Field(s_space,val=diagProber.diagonal.val).weight(-1)" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "skip" } }, "source": [ "### Get data" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true, "slideshow": { "slide_type": "skip" } }, "outputs": [], "source": [ "s_power = sh.power_analyze()\n", "m_power = fft(m).power_analyze()\n", "s_power_data = s_power.val.get_full_data().real\n", "m_power_data = m_power.val.get_full_data().real\n", "\n", "# Get signal data and reconstruction data\n", "s_data = s.val.get_full_data().real\n", "m_data = m.val.get_full_data().real\n", "m_var_data = m_var.val.get_full_data().real\n", "uncertainty = np.sqrt(np.abs(m_var_data))\n", "\n", "d_data = d.val.get_full_data().real\n", "\n", "# Set lost data to NaN for proper plotting\n", "d_data[d_data == 0] = np.nan" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "skip" } }, "outputs": [], "source": [ "fig = plt.figure(figsize=(15,10))\n", "plt.plot(s_data, 'k', label=\"Signal\", alpha=.5, linewidth=1)\n", "plt.plot(d_data, 'k+', label=\"Data\", alpha=1)\n", "plt.axvspan(l, h, facecolor='0.8', alpha=.5)\n", "plt.title(\"Incomplete Data\")\n", "plt.legend()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "subslide" } }, "outputs": [], "source": [ "fig" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "skip" } }, "outputs": [], "source": [ "fig = plt.figure(figsize=(15,10))\n", "plt.plot(s_data, 'k', label=\"Signal\", alpha=1, linewidth=1)\n", "plt.plot(d_data, 'k+', label=\"Data\", alpha=.5)\n", "plt.plot(m_data, 'r', label=\"Reconstruction\")\n", "plt.axvspan(l, h, facecolor='0.8', alpha=.5)\n", "plt.fill_between(range(N_pixels), m_data - uncertainty, m_data + uncertainty, facecolor='0')\n", "plt.title(\"Reconstruction of incomplete data\")\n", "plt.legend()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "slideshow": { "slide_type": "subslide" } }, "outputs": [], "source": [ "fig" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# 2d Example" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": true }, "outputs": [], "source": [ "N_pixels = 256 # Number of pixels\n", "sigma2 = 1000 # Noise variance\n", "\n", "\n", "def pow_spec(k):\n", " P0, k0, gamma = [.2, 20, 4]\n", " return P0 * (1. + (k/k0)**2)**(- gamma / 2)\n",