{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The plasma dispersion function\n",
    "==============================\n",
    "\n",
    "Let's import some basics (and `plasmapy`!)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from plasmapy.dispersion import plasma_dispersion_func, plasma_dispersion_func_deriv"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Take a look at the docs to :func:`~plasmapy.dispersion.dispersion_functions.plasma_dispersion_func` for more information on this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll now make some sample data to visualize the dispersion function:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-1, 1, 200)\n",
    "X, Y = np.meshgrid(x, x)\n",
    "Z = X + 1j * Y\n",
    "print(Z.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we start plotting, let's make a visualization function first:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def plot_complex(X, Y, Z, N=50):\n",
    "    fig, (real_axis, imag_axis) = plt.subplots(1, 2)\n",
    "    real_axis.contourf(X, Y, Z.real, N)\n",
    "    imag_axis.contourf(X, Y, Z.imag, N)\n",
    "    real_axis.set_title(\"Real values\")\n",
    "    imag_axis.set_title(\"Imaginary values\")\n",
    "    for ax in [real_axis, imag_axis]:\n",
    "        ax.set_xlabel(\"Real values\")\n",
    "        ax.set_ylabel(\"Imaginary values\")\n",
    "    fig.tight_layout()\n",
    "\n",
    "\n",
    "plot_complex(X, Y, Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now apply our visualization function to our simple dispersion relation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "# sphinx_gallery_thumbnail_number = 2\n",
    "F = plasma_dispersion_func(Z)\n",
    "plot_complex(X, Y, F)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's find the area where the dispersion function has a lesser than zero real part:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "plot_complex(X, Y, F.real < 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the derivative:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "F = plasma_dispersion_func_deriv(Z)\n",
    "plot_complex(X, Y, F)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting the same function on a larger area:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-2, 2, 400)\n",
    "X, Y = np.meshgrid(x, x)\n",
    "Z = X + 1j * Y\n",
    "print(Z.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "F = plasma_dispersion_func(Z)\n",
    "plot_complex(X, Y, F, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we examine the derivative of the dispersion function as a function of the phase velocity of an electromagnetic wave propagating through the plasma. This is recreating figure 5.1 in J. Sheffield, D. Froula, S. H. Glenzer, and N. C. Luhmann Jr, Plasma scattering of electromagnetic radiation: theory and measurement techniques, ch. 5, p. 106 (Academic press, 2010)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "xs = np.linspace(0, 4, 100)\n",
    "ws = (-1 / 2) * plasma_dispersion_func_deriv(xs)\n",
    "wRe = np.real(ws)\n",
    "wIm = np.imag(ws)\n",
    "\n",
    "plt.plot(xs, wRe, label=\"Re\")\n",
    "plt.plot(xs, wIm, label=\"Im\")\n",
    "plt.axis([0, 4, -0.3, 1])\n",
    "plt.legend(\n",
    "    loc=\"upper right\",\n",
    "    frameon=False,\n",
    "    labelspacing=0.001,\n",
    "    fontsize=14,\n",
    "    borderaxespad=0.1,\n",
    ")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}