{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Speed of single-shot vs reuse of a HankelTransform object\n\nFor a simple case (as in `sphx_glr_auto_examples_one_shot_example.py`)\nthere are two simple forward :func:`~.one_shot.qdht` and [inverse :func:`~.one_shot.iqdht`]\nfunctions which can be used to calculate the [inverse] Hankel transform of a function sampled\nat an arbitrary set of points in radius [wave-number] space.\n\nHere we will use the same example application as `sphx_glr_auto_examples_usage_example.py`:\na beam-propagation method propagation of a radially-symmetric Gaussian beam.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import time\n\nimport numpy as np\nfrom scipy import interpolate\nimport matplotlib.pyplot as plt\n\nfrom pyhank import HankelTransform, qdht, iqdht\nfrom helper import gauss1d, imagesc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Initialise radius  and $z$ grids and beam parameters as in\n`sphx_glr_auto_examples_usage_example.py`.\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)\nNz = 100  # Number of z positions\nz_max = 0.1  # Maximum propagation distance\n\nr = np.linspace(0, r_max, nr)\nz = np.linspace(0, z_max, Nz)\n\nDr = 100e-6  # Beam radius (100um)\nlambda_ = 488e-9  # wavelength 488nm\nk0 = 2 * np.pi / lambda_  # Vacuum k vector\n\nfield = gauss1d(r, 0, Dr)   # Initial field"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we need two functions that propagate the beam in two ways (giving the same answer).\nThe first will use single shot, the second will use a :class:`.HankelTransform` object.\nBelow we will run each of them in turn and compare the speed.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def propagate_using_object(r: np.ndarray, field: np.ndarray) -> np.ndarray:\n    transformer = HankelTransform(order=0, radial_grid=r)\n    field_for_transform = transformer.to_transform_r(field)  # Resampled field\n    hankel_transform = transformer.qdht(field_for_transform)\n\n    propagated_field = np.zeros((nr, Nz), dtype=complex)\n    kz = np.sqrt(k0 ** 2 - transformer.kr ** 2)\n    for n, z_loop in enumerate(z):\n        phi_z = kz * z_loop  # Propagation phase\n        hankel_transform_at_z = hankel_transform * np.exp(1j * phi_z)  # Apply propagation\n        field_at_z_transform_grid = transformer.iqdht(hankel_transform_at_z)  # iQDHT\n        propagated_field[:, n] = transformer.to_original_r(field_at_z_transform_grid)  # Interpolate output\n    intensity = np.abs(propagated_field) ** 2\n    return intensity\n\n\ndef propagate_using_single_shot(r: np.ndarray, field: np.ndarray) -> np.ndarray:\n    kr, hankel_transform = qdht(r, field)\n\n    propagated_field = np.zeros((nr, Nz), dtype=complex)\n    kz = np.sqrt(k0 ** 2 - kr ** 2)\n    for n, z_loop in enumerate(z):\n        phi_z = kz * z_loop  # Propagation phase\n        hankel_transform_at_z = hankel_transform * np.exp(1j * phi_z)  # Apply propagation\n        r_transform, field_at_z_transform_grid = iqdht(kr, hankel_transform_at_z)  # iQDHT\n        f = interpolate.interp1d(r_transform, field_at_z_transform_grid, axis=0,\n                                 fill_value='extrapolate', kind='cubic')\n        propagated_field[:, n] = f(r)\n    intensity = np.abs(propagated_field) ** 2\n    return intensity"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now run and time the two functions:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tic = time.time()\nsingle_shot_intensity = propagate_using_single_shot(r, field)\ntoc = time.time()\nprint(f'Single shot propagation took {toc-tic:.2f} s')\n\ntic = time.time()\nobject_intensity = propagate_using_object(r, field)\ntoc = time.time()\nprint(f'Object propagation took {toc-tic:.2f} s')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The single shot approach takes a *lot* longer!\n\nPlot the two results to check they are the same:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure()\nplt.subplot(2, 1, 1)\nimagesc(z * 1e3, r * 1e3, single_shot_intensity)\nplt.xlabel('Propagation distance ($z$) /mm')\nplt.ylabel('Radial position ($r$) /mm')\nplt.colorbar()\nplt.ylim([0, 1])\n\nplt.subplot(2, 1, 2)\nimagesc(z * 1e3, r * 1e3, object_intensity)\nplt.xlabel('Propagation distance ($z$) /mm')\nplt.ylabel('Radial position ($r$) /mm')\nplt.ylim([0, 1])\nplt.colorbar()\nplt.tight_layout()"
      ]
    }
  ],
  "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
}