{
 "worksheets": [
  {
   "cells": [
    {
     "metadata": {},
     "cell_type": "code",
     "input": "%pylab qt",
     "prompt_number": 1,
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": "Populating the interactive namespace from numpy and matplotlib\n"
      }
     ],
     "language": "python",
     "trusted": true,
     "collapsed": false
    },
    {
     "metadata": {
      "heading_collapsed": false
     },
     "cell_type": "heading",
     "source": "Define Wigner function of arbitrary coherent state superposition (cat)",
     "level": 2
    },
    {
     "metadata": {},
     "cell_type": "markdown",
     "source": "The Wigner function of a cat state $|\\psi\\rangle = \\frac{1}{N} \\left( |\\alpha\\rangle + e^{i\\phi} |\\beta\\rangle \\right)$ is\n\n$$\nW_{|\\psi\\rangle} = \\frac{1}{N^2} \\left(W_{|\\alpha\\rangle \\langle\\alpha|} + W_{|\\beta\\rangle \\langle\\beta|} \n+ 2 \\mathrm{Re} \\left[ e^{-i\\phi} W_{|\\alpha\\rangle \\langle\\beta|} \\right] \\right)\n$$\n\nwith\n\n$$\nN = \\sqrt{2 + 2 e^{-(|\\alpha|^2 + |\\beta|^2)/2} \\mathrm{Re} \\left[ e^{\\alpha^\\ast \\beta + i\\phi} \\right]}\n$$\n\nand\n\n$$\n\\begin{align*}\nW_{|\\alpha\\rangle \\langle\\beta|}(x,p) \n &= \\frac{1}{\\pi} e^{-(x - \\sigma_x)^2 - (p - \\sigma_p)^2}\n  e^{i ( 2(x-\\sigma_x) \\delta_p - 2(p-\\sigma_p) \\delta_x - \n   \\frac{1}{2} \\alpha_x \\beta_p + \\frac{1}{2} \\alpha_p \\beta_x)}, \\\\\n &= \\frac{1}{\\pi} e^{-(x - (\\sigma_x + i\\delta_p))^2 - (p - (\\sigma_p - i\\delta_x))^2}\n   e^{-\\delta_x^2 - \\delta_p^2} e^{\\frac{i}{2} (\\alpha_p \\beta_x - \\alpha_x \\beta_p)},\n\\end{align*}\n$$\n\n$$\n\\alpha_x = \\sqrt{2} \\mathrm{Re}[\\alpha], \\quad \\alpha_p = \\sqrt{2} \\mathrm{Im} [\\alpha], \\quad\n\\beta_x = \\sqrt{2} \\mathrm{Re}[\\beta], \\quad \\beta_p = \\sqrt{2} \\mathrm{Im} [\\beta],\n$$\n\n$$\n\\sigma_{x,p} = \\frac{\\alpha_{x,p} + \\beta_{x,p}}{2}, \\quad\n\\delta_{x,p} = \\frac{\\alpha_{x,p} - \\beta_{x,p}}{2}.\n$$\n\nThis can be derived from direct application of the definition of the Wigner function,\n\n$$\nW_\\hat{O}(x,p) \\equiv \\frac{1}{2\\pi} \\int_{-\\infty}^{\\infty} e^{iyp} \n  \\langle x - \\tfrac{y}{2}| \\hat{O} |x + \\tfrac{y}{2} \\rangle dy\n$$\n\nand the coherent state wavefunction,\n\n$$\n\\langle x | \\alpha \\rangle = \\frac{1}{\\pi^{1/4}} e^{-\\frac{1}{2} (x - \\sqrt{2} \\mathrm{Re}\\alpha)^2\n  + i\\sqrt{2} \\mathrm{Im}\\alpha (x - \\mathrm{Re}\\alpha / \\sqrt{2})} .\n$$"
    },
    {
     "metadata": {},
     "cell_type": "code",
     "input": "def Wab(alpha, beta, x, p):\n    ax, ap = sqrt(2) * alpha.real, sqrt(2) * alpha.imag\n    bx, bp = sqrt(2) * beta.real, sqrt(2) * beta.imag\n    sx, sp = .5*(ax + bx), .5*(ap + bp)\n    dx, dp = .5*(ax - bx), .5*(ap - bp)\n\n    return 1/pi * exp(-(x - (sx + 1j*dp))**2 - (p - (sp - 1j*dx))**2 - dx**2 - dp**2 - .5j * (ax*bp - bx*ap))\n\ndef Wcat(alpha, beta, phi, x, p):\n    N = sqrt(2 + 2*exp(-(abs(alpha)**2 + abs(beta)**2)/2) * real(exp(conj(alpha) * beta + 1j*phi)))\n    \n    return 1/N**2 * real(Wab(alpha, alpha, x, p) + Wab(beta, beta, x, p) +\n                         2 * real(exp(-1j*phi) * Wab(alpha, beta, x, p)))",
     "prompt_number": 2,
     "outputs": [],
     "language": "python",
     "trusted": true,
     "collapsed": false
    },
    {
     "metadata": {
      "heading_collapsed": false
     },
     "cell_type": "heading",
     "source": "Plot Wigner function",
     "level": 2
    },
    {
     "metadata": {},
     "cell_type": "code",
     "input": "class CatPlot:\n    def __init__(self, image, alpha, beta, phi):\n        self.image = image\n        self.alpha = alpha\n        self.beta = beta\n        self.phi = phi\n        image.figure.canvas.mpl_connect('motion_notify_event', self.onclick)\n        image.figure.canvas.mpl_connect('scroll_event', self.onscroll)\n        self.update()\n        \n    def onclick(self, event):\n        if event.button==1 and event.inaxes:\n            x, p = event.xdata, event.ydata\n            ax, ap = sqrt(2) * real(self.alpha), sqrt(2) * imag(self.alpha)\n            bx, bp = sqrt(2) * real(self.beta), sqrt(2) * imag(self.beta)\n            \n            da = sqrt((x - ax)**2 + (p - ap)**2)\n            db = sqrt((x - bx)**2 + (p - bp)**2)\n            dab = sqrt((ax - bx)**2 + (ap - bp)**2)\n            \n            if da > dab/3. and db > dab/3. and dab > .5:\n                self.alpha = (x + (ax - bx)/2. + 1j*(p + (ap - bp)/2.))/sqrt(2.)\n                self.beta = (x - (ax - bx)/2. + 1j*(p - (ap - bp)/2.))/sqrt(2.)\n            elif da < db:\n                self.alpha = (x + 1j*p)/sqrt(2.)\n                self.beta = (bx - (x-ax) + 1j*(bp - (p-ap)))/sqrt(2.)\n            else:\n                self.beta = (x + 1j*p)/sqrt(2.)\n                self.alpha = (ax - (x-bx) + 1j*(ap - (p-bp)))/sqrt(2.)\n            \n            self.update()\n            \n    def onscroll(self, event):\n        self.phi += pi/16. * event.step\n        self.update()\n        \n    def update(self):\n        wig = Wcat(self.alpha, self.beta, self.phi, X, Y)\n\n        self.image.set_data(wig.T)\n        self.image.axes.set_title(r'|{: 4.2f}> + exp(i {:d}pi/16) |{: 4.2f}>'.format(self.alpha, int(self.phi*16/pi), self.beta))\n        self.image.figure.canvas.draw()\n        \n            \nX, Y = mgrid[-5:5:101j, -5:5:101j]\nfig, ax = plt.subplots(figsize=(10,10))\nax.grid()\nimage = ax.imshow(X,\n                  vmin=-1/pi, vmax=1/pi, extent=(-5,5,-5,5), origin='lower',\n                  cmap='RdBu_r')\n\nCatPlot(image, 2, -2, 0)",
     "prompt_number": 3,
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 3,
       "metadata": {},
       "text": "<__main__.CatPlot instance at 0x7fa1b401bf80>"
      }
     ],
     "language": "python",
     "trusted": true,
     "collapsed": false
    }
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "name": "",
  "signature": "sha256:ae5099cde68e5b22e2ae597799d0af855dac3c025192a58d5e2ddc3607fe1ca7"
 },
 "nbformat": 3
}