{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Typical usage\n\nTo demonstrate the use of the Hankel Transform class, we will give an example\nof propagating a radially-symmetric beam using the beam propagation method.\n\nIn this case, it will be a simple Gaussian beam propagating way from focus and\ndiverging.\n\nFirst we will use a loop over $z$ position, and then we will demonstrate\nthe vectorisation of the :func:`.HankelTransforms.iqdht` (and\n:func:`~.HankelTransforms.qdht`) functions.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "First import the standard libraries\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then the functions from this package\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from pyhank import HankelTransform\n# noinspection PyUnresolvedReferences\nfrom helper import gauss1d, imagesc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Initialise radius grid\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nr = 1024  # Number of sample points\nr_max = 5e-3  # Maximum radius (5mm)\nr = np.linspace(0, r_max, nr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Initialise $z$ grid\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Nz = 200  # Number of z positions\nz_max = 0.1  # Maximum propagation distance\nz = np.linspace(0, z_max, Nz)  # Propagation axis"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set up beam parameters\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Dr = 100e-6  # Beam radius (100um)\nlambda_ = 488e-9  # wavelength 488nm\nk0 = 2 * np.pi / lambda_  # Vacuum k vector"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set up a :class:`.HankelTransform` object, telling it the order (``0``) and\nthe radial grid.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "H = HankelTransform(order=0, radial_grid=r)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set up the electric field profile at $z = 0$, and resample onto the correct radial grid\n(``transformer.r``) as required for the QDHT.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Er = gauss1d(r, 0, Dr)   # Initial field\nErH = H.to_transform_r(Er)  # Resampled field"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Perform Hankel Transform\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Convert from physical field to physical wavevector\nEkrH = H.qdht(ErH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Propagate the beam - loop\nDo the propagation in a loop over $z$\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Pre-allocate an array for field as a function of r and z\nErz = np.zeros((nr, Nz), dtype=complex)\nkz = np.sqrt(k0 ** 2 - H.kr ** 2)\nfor i, z_loop in enumerate(z):\n    phi_z = kz * z_loop  # Propagation phase\n    EkrHz = EkrH * np.exp(1j * phi_z)  # Apply propagation\n    ErHz = H.iqdht(EkrHz)  # iQDHT\n    Erz[:, i] = H.to_original_r(ErHz)  # Interpolate output\nIrz = np.abs(Erz) ** 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plotting\nPlot the initial field and radial wavevector distribution (given by the\nHankel transform)\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure()\nplt.plot(r * 1e3, np.abs(Er) ** 2, r * 1e3, np.unwrap(np.angle(Er)),\n         H.r * 1e3, np.abs(ErH) ** 2, H.r * 1e3, np.unwrap(np.angle(ErH)), '+')\nplt.title('Initial electric field distribution')\nplt.xlabel('Radial co-ordinate (r) /mm')\nplt.ylabel('Field intensity /arb.')\nplt.legend(['$|E(r)|^2$', '$\\\\phi(r)$', '$|E(H.r)|^2$', '$\\\\phi(H.r)$'])\nplt.axis([0, 1, 0, 1])\n\nplt.figure()\nplt.plot(H.kr, np.abs(EkrH) ** 2)\nplt.title('Radial wave-vector distribution')\nplt.xlabel(r'Radial wave-vector ($k_r$) /rad $m^{-1}$')\nplt.ylabel('Field intensity /arb.')\nplt.axis([0, 3e4, 0, np.max(np.abs(EkrH) ** 2)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now plot an image showing the intensity as a function of\nradius and propagation distance.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure()\nimagesc(z * 1e3, r * 1e3, Irz)\nplt.title('Radial field intensity as a function of propagation for annular beam')\nplt.xlabel('Propagation distance ($z$) /mm')\nplt.ylabel('Radial position ($r$) /mm')\nplt.ylim([0, 1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The plot above shows a reduction of intensity with $z$, but it is\nbit difficult to see the beam growing in $r$. To show that, let's\nplot the intensity normalised such that the peak intensity at each $z$\ncoordinate is the same.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Irz_norm = Irz / Irz[0, :]\n\nplt.figure()\nimagesc(z * 1e3, r * 1e3, Irz_norm)\nplt.xlabel('Propagation distance ($z$) /mm')\nplt.ylabel('Radial position ($r$) /mm')\nplt.ylim([0, 1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Propagate the beam - vectorised\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "kz = np.sqrt(k0 ** 2 - H.kr ** 2)\nphi_z = kz[:, np.newaxis] * z[np.newaxis, :]  # Propagation phase\nEkrHz = EkrH[:, np.newaxis] * np.exp(1j * phi_z)  # Apply propagation\nErHz = H.iqdht(EkrHz)  # iQDHT\nErz = H.to_original_r(ErHz)  # Interpolate output\nIrz_vectorised = np.abs(Erz) ** 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now plot the result to check it is the same as the loop approach\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure()\nimagesc(z * 1e3, r * 1e3, Irz_vectorised)\nplt.title('Radial field intensity as a function of propagation for annular beam')\nplt.xlabel('Propagation distance ($z$) /mm')\nplt.ylabel('Radial position ($r$) /mm')\nplt.ylim([0, 1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Assert the two approaches produce the same intensity\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "assert np.allclose(Irz, Irz_vectorised)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}