Commit 81cab5d6 authored by Philipp Arras's avatar 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",
"\n",
"\n",
"s_space = RGSpace([N_pixels, N_pixels])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true,
"slideshow": {
"slide_type": "skip"
}
},
"outputs": [],
"source": [
"fft = FFTOperator(s_space)\n",
"h_space = fft.target[0]\n",
"p_space = PowerSpace(h_space)\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 = SmoothingOperator(s_space, sigma=.01)\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",