Wiener filter demo

Works only with python2 and on commit 1d10be46.

demos/Wiener Filter.ipynb

+{
+ "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",