{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dispersion: A Full Two Fluid Solution\n",
    "\n",
    "[tfds]: ../../api/plasmapy.dispersion.analytical.two_fluid_.two_fluid.rst\n",
    "[bellan2012]: https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2012JA017856\n",
    "[stringer1963]: https://doi.org/10.1088/0368-3281/5/2/304\n",
    "\n",
    "This notebook walks through the functionality of the [two_fluid()][tfds] function.  This function computes the wave frequencies for given wavenumbers and plasma parameters based on the analytical solution presented by [Bellan 2012][bellan2012] to the [Stringer 1963][stringer1963] two fluid dispersion relation.  The two fluid dispersion equaiton assumes a uniform magnetic field, a zero D.C. electric field, and low-frequency waves $\\omega / k c \\ll 1$ which equates to\n",
    "\n",
    "$$\n",
    "    \\left( \\cos^2 \\theta - Q \\frac{\\omega^2}{k^2 {v_A}^2} \\right)\n",
    "    \\left[\n",
    "        \\left( \\cos^2 \\theta - \\frac{\\omega^2}{k^2 {c_s}^2} \\right)\n",
    "        - Q \\frac{\\omega^2}{k^2 {v_A}^2} \\left(\n",
    "            1 - \\frac{\\omega^2}{k^2 {c_s}^2}\n",
    "        \\right)\n",
    "    \\right] \\\\\n",
    "    = \\left(1 - \\frac{\\omega^2}{k^2 {c_s}^2} \\right)\n",
    "      \\frac{\\omega^2}{{\\omega_{ci}}^2} \\cos^2 \\theta\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$Q = 1 + k^2 c^2/{\\omega_{pe}}^2$$\n",
    "\n",
    "$$\\cos \\theta = \\frac{k_z}{k}$$\n",
    "\n",
    "$$\\mathbf{B_o} = B_{o} \\mathbf{\\hat{z}}$$\n",
    "\n",
    "$\\omega$ is the wave frequency, $k$ is the wavenumber, $v_A$ is the Alfvén velocity, $c_s$ is the sound speed, $\\omega_{ci}$ is the ion gyrofrequency, and $\\omega_{pe}$ is the electron plasma frequency.\n",
    "\n",
    "The approach outlined in Section 5 of [Bellan 2012][bellan2012] produces exact roots to the above dispersion equation for all three modes (fast, acoustic, and Alfvén) without having to make additional approximations.  The following dispersion relation is what the [two_fluid()][tfds] function computes.\n",
    "\n",
    "$$\n",
    "    \\frac{\\omega}{\\omega_{ci}} = \\sqrt{\n",
    "        2 \\Lambda \\sqrt{-\\frac{P}{3}} \\cos\\left(\n",
    "            \\frac{1}{3} \\cos^{-1}\\left(\n",
    "                \\frac{3q}{2p} \\sqrt{-\\frac{3}{p}}\n",
    "            \\right)\n",
    "            - \\frac{2 \\pi}{3}j\n",
    "        \\right)\n",
    "        + \\frac{\\Lambda A}{3}\n",
    "    }\n",
    "$$\n",
    "\n",
    "where $j = 0$ represents the fast mode, $j = 1$ represents the Alfvén mode, and $j = 2$ represents the acoustic mode.  Additionally,\n",
    "\n",
    "$$p = \\frac{3B-A^2}{3} \\; , \\; q = \\frac{9AB-2A^3-27C}{27}$$\n",
    "\n",
    "$$A = \\frac{Q + Q^2 \\beta + Q \\alpha + \\alpha \\Lambda}{Q^2} \\;\n",
    "    , \\; B = \\alpha \\frac{1 + 2 Q \\beta + \\Lambda \\beta}{Q^2} \\;\n",
    "    , \\; C = \\frac{\\alpha^2 \\beta}{Q^2}$$\n",
    "\n",
    "$$\\alpha = \\cos^2 \\theta \\;\n",
    "    , \\; \\beta = \\left( \\frac{c_s}{v_A}\\right)^2 \\;\n",
    "    , \\; \\Lambda = \\left( \\frac{k v_{A}}{\\omega_{ci}}\\right)^2$$\n",
    "\n",
    "## Contents:\n",
    "\n",
    "1. [Wave Propagating at 45 Degrees](#Wave-Propagating-at-45-Degrees)\n",
    "2. [Wave frequencies on the k-theta plane](#Wave-frequencies-on-the-k-theta-plane)\n",
    "3. [Reproduce Figure 1 from Bellan 2012](#Reproduce-Figure-1-from-Bellan-2012)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import astropy.units as u\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from astropy.constants.si import c\n",
    "from matplotlib.ticker import MultipleLocator\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "from plasmapy.dispersion.analytical.two_fluid_ import two_fluid\n",
    "from plasmapy.formulary import speeds\n",
    "from plasmapy.formulary.frequencies import gyrofrequency, plasma_frequency, wc_, wp_\n",
    "from plasmapy.formulary.lengths import inertial_length\n",
    "from plasmapy.particles import Particle\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = [10.5, 0.56 * 10.5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wave Propagating at 45 Degrees\n",
    "\n",
    "Below we define the required parameters to compute the wave frequencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define input parameters\n",
    "inputs = {\n",
    "    \"k\": np.linspace(10**-7, 10**-2, 10000) * u.rad / u.m,\n",
    "    \"theta\": 45 * u.deg,\n",
    "    \"n_i\": 5 * u.cm**-3,\n",
    "    \"B\": 8.3e-9 * u.T,\n",
    "    \"T_e\": 1.6e6 * u.K,\n",
    "    \"T_i\": 4.0e5 * u.K,\n",
    "    \"ion\": Particle(\"p+\"),\n",
    "}\n",
    "\n",
    "# a few useful plasma parameters\n",
    "params = {\n",
    "    \"n_e\": inputs[\"n_i\"] * abs(inputs[\"ion\"].charge_number),\n",
    "    \"cs\": speeds.ion_sound_speed(\n",
    "        inputs[\"T_e\"],\n",
    "        inputs[\"T_i\"],\n",
    "        inputs[\"ion\"],\n",
    "    ),\n",
    "    \"va\": speeds.Alfven_speed(\n",
    "        inputs[\"B\"],\n",
    "        inputs[\"n_i\"],\n",
    "        ion=inputs[\"ion\"],\n",
    "    ),\n",
    "    \"wci\": gyrofrequency(inputs[\"B\"], inputs[\"ion\"]),\n",
    "}\n",
    "params[\"lpe\"] = inertial_length(params[\"n_e\"], \"e-\")\n",
    "params[\"wpe\"] = plasma_frequency(params[\"n_e\"], \"e-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The computed wave frequencies ($rad/s$) are returned in a dictionary with keys representing the wave modes and the values being an Astropy [Quantity](https://docs.astropy.org/en/stable/api/astropy.units.Quantity.html#astropy.units.Quantity).  Since our inputs had a scalar $\\theta$ and a 1D array of $k$'s, the computed wave frequencies will be a 1D array of size equal to the size of the $k$ array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute\n",
    "omegas = two_fluid(**inputs)\n",
    "\n",
    "(list(omegas.keys()), omegas[\"fast_mode\"], omegas[\"fast_mode\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the results of each wave mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = 14  # default font size\n",
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figheight = 1.6 * figheight\n",
    "fig = plt.figure(figsize=[figwidth, figheight])\n",
    "\n",
    "# normalize data\n",
    "k_prime = inputs[\"k\"] * params[\"lpe\"]\n",
    "\n",
    "# plot\n",
    "plt.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas[\"fast_mode\"] / params[\"wpe\"]),\n",
    "    \"r.\",\n",
    "    ms=1,\n",
    "    label=\"Fast\",\n",
    ")\n",
    "ax = plt.gca()\n",
    "ax.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas[\"alfven_mode\"] / params[\"wpe\"]),\n",
    "    \"b.\",\n",
    "    ms=1,\n",
    "    label=\"Alfvén\",\n",
    ")\n",
    "ax.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas[\"acoustic_mode\"] / params[\"wpe\"]),\n",
    "    \"g.\",\n",
    "    ms=1,\n",
    "    label=\"Acoustic\",\n",
    ")\n",
    "\n",
    "# adjust axes\n",
    "ax.set_xlabel(r\"$kc / \\omega_{pe}$\", fontsize=fs)\n",
    "ax.set_ylabel(r\"$Re(\\omega / \\omega_{pe})$\", fontsize=fs)\n",
    "ax.set_yscale(\"log\")\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_ylim(1e-6, 2e-2)\n",
    "ax.tick_params(\n",
    "    which=\"both\",\n",
    "    direction=\"in\",\n",
    "    width=1,\n",
    "    labelsize=fs,\n",
    "    right=True,\n",
    "    length=5,\n",
    ")\n",
    "\n",
    "# annotate\n",
    "text = (\n",
    "    rf\"$v_A/c_s = {params['va'] / params['cs']:.1f} \\qquad \"\n",
    "    rf\"c/v_A = 10^{np.log10(c / params['va']):.0f} \\qquad \"\n",
    "    f\"\\\\theta = {inputs['theta'].value:.0f}\"\n",
    "    \"^{\\\\circ}$\"\n",
    ")\n",
    "ax.text(0.25, 0.95, text, transform=ax.transAxes, fontsize=18)\n",
    "ax.legend(loc=\"upper left\", markerscale=5, fontsize=fs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wave frequencies on the k-theta plane\n",
    "\n",
    "Let us now look at the distribution of $\\omega$ on a $k$-$\\theta$ plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define input parameters\n",
    "inputs = {\n",
    "    \"k\": np.linspace(10**-7, 10**-2, 10000) * u.rad / u.m,\n",
    "    \"theta\": np.linspace(5, 85, 100) * u.deg,\n",
    "    \"n_i\": 5 * u.cm**-3,\n",
    "    \"B\": 8.3e-9 * u.T,\n",
    "    \"T_e\": 1.6e6 * u.K,\n",
    "    \"T_i\": 4.0e5 * u.K,\n",
    "    \"ion\": Particle(\"p+\"),\n",
    "}\n",
    "\n",
    "# a few useful plasma parameters\n",
    "params = {\n",
    "    \"n_e\": inputs[\"n_i\"] * abs(inputs[\"ion\"].charge_number),\n",
    "    \"cs\": speeds.ion_sound_speed(\n",
    "        inputs[\"T_e\"],\n",
    "        inputs[\"T_i\"],\n",
    "        inputs[\"ion\"],\n",
    "    ),\n",
    "    \"va\": speeds.Alfven_speed(\n",
    "        inputs[\"B\"],\n",
    "        inputs[\"n_i\"],\n",
    "        ion=inputs[\"ion\"],\n",
    "    ),\n",
    "    \"wci\": gyrofrequency(inputs[\"B\"], inputs[\"ion\"]),\n",
    "}\n",
    "params[\"lpe\"] = inertial_length(params[\"n_e\"], \"e-\")\n",
    "params[\"wpe\"] = plasma_frequency(params[\"n_e\"], \"e-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the $\\theta$ and $k$ values are now 1-D arrays, the returned wave frequencies will be 2-D arrays with the first dimension matching the size of $k$ and the second dimension matching the size of $\\theta$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute\n",
    "omegas = two_fluid(**inputs)\n",
    "\n",
    "(\n",
    "    omegas[\"fast_mode\"].shape,\n",
    "    omegas[\"fast_mode\"].shape[0] == inputs[\"k\"].size,\n",
    "    omegas[\"fast_mode\"].shape[1] == inputs[\"theta\"].size,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot (the fast mode)!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = 14  # default font size\n",
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figheight = 1.6 * figheight\n",
    "fig = plt.figure(figsize=[figwidth, figheight])\n",
    "\n",
    "# normalize data\n",
    "k_prime = inputs[\"k\"] * params[\"lpe\"]\n",
    "zdata = np.transpose(np.real(omegas[\"fast_mode\"].value)) / params[\"wpe\"].value\n",
    "\n",
    "# plot\n",
    "im = plt.imshow(\n",
    "    zdata,\n",
    "    aspect=\"auto\",\n",
    "    origin=\"lower\",\n",
    "    extent=[\n",
    "        np.min(k_prime.value),\n",
    "        np.max(k_prime.value),\n",
    "        np.min(inputs[\"theta\"].value),\n",
    "        np.max(inputs[\"theta\"].value),\n",
    "    ],\n",
    "    interpolation=None,\n",
    "    cmap=plt.cm.Spectral,\n",
    ")\n",
    "ax = plt.gca()\n",
    "\n",
    "# # adjust axes\n",
    "ax.set_xscale(\"linear\")\n",
    "ax.set_xlabel(r\"$kc/\\omega_{pe}$\", fontsize=fs)\n",
    "ax.set_ylabel(r\"$\\theta$ [$deg.$]\", fontsize=fs)\n",
    "ax.tick_params(\n",
    "    which=\"both\",\n",
    "    direction=\"in\",\n",
    "    width=2,\n",
    "    labelsize=fs,\n",
    "    right=True,\n",
    "    top=True,\n",
    "    length=10,\n",
    ")\n",
    "\n",
    "# Add colorbar\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"top\", size=\"5%\", pad=0.07)\n",
    "cbar = plt.colorbar(\n",
    "    im,\n",
    "    cax=cax,\n",
    "    orientation=\"horizontal\",\n",
    "    ticks=None,\n",
    "    fraction=0.05,\n",
    "    pad=0.0,\n",
    ")\n",
    "cbar.ax.tick_params(\n",
    "    axis=\"x\",\n",
    "    direction=\"in\",\n",
    "    width=2,\n",
    "    length=10,\n",
    "    top=True,\n",
    "    bottom=False,\n",
    "    labelsize=fs,\n",
    "    pad=0.0,\n",
    "    labeltop=True,\n",
    "    labelbottom=False,\n",
    ")\n",
    "cbar.ax.xaxis.set_label_position(\"top\")\n",
    "cbar.set_label(r\"$\\omega/\\omega_{pe}$\", fontsize=fs, labelpad=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reproduce Figure 1 from Bellan 2012\n",
    "\n",
    "[bellan2012]: https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2012JA017856\n",
    "\n",
    "\n",
    "Figure 1 of [Bellan 2012][bellan2012] chooses parameters such that $\\beta = 0.4$ and $\\Lambda=0.4$.  Below we define parameters to approximate Bellan's assumptions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define input parameters\n",
    "inputs = {\n",
    "    \"B\": 400e-4 * u.T,\n",
    "    \"ion\": Particle(\"He+\"),\n",
    "    \"n_i\": 6.358e19 * u.m**-3,\n",
    "    \"T_e\": 20 * u.eV,\n",
    "    \"T_i\": 10 * u.eV,\n",
    "    \"theta\": np.linspace(0, 90) * u.deg,\n",
    "    \"k\": (2 * np.pi * u.rad) / (0.56547 * u.m),\n",
    "}\n",
    "\n",
    "# a few useful plasma parameters\n",
    "params = {\n",
    "    \"n_e\": inputs[\"n_i\"] * abs(inputs[\"ion\"].charge_number),\n",
    "    \"cs\": speeds.cs_(inputs[\"T_e\"], inputs[\"T_i\"], inputs[\"ion\"]),\n",
    "    \"wci\": wc_(inputs[\"B\"], inputs[\"ion\"]),\n",
    "    \"va\": speeds.va_(inputs[\"B\"], inputs[\"n_i\"], ion=inputs[\"ion\"]),\n",
    "}\n",
    "params[\"beta\"] = (params[\"cs\"] / params[\"va\"]).value ** 2\n",
    "params[\"wpe\"] = wp_(params[\"n_e\"], \"e-\")\n",
    "params[\"Lambda\"] = (inputs[\"k\"] * params[\"va\"] / params[\"wci\"]).value ** 2\n",
    "\n",
    "(params[\"beta\"], params[\"Lambda\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute\n",
    "omegas = two_fluid(**inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate data for plots\n",
    "plt_vals = {}\n",
    "for mode, arr in omegas.items():\n",
    "    norm = (np.absolute(arr) / (inputs[\"k\"] * params[\"va\"])).value ** 2\n",
    "    plt_vals[mode] = {\n",
    "        \"x\": norm * np.sin(inputs[\"theta\"].to(u.rad).value),\n",
    "        \"y\": norm * np.cos(inputs[\"theta\"].to(u.rad).value),\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx-thumbnail": {
     "tooltip": "Full Two Fluid Dispersion Solution"
    }
   },
   "outputs": [],
   "source": [
    "fs = 14  # default font size\n",
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figheight = 1.6 * figheight\n",
    "fig = plt.figure(figsize=[figwidth, figheight])\n",
    "\n",
    "# Fast mode\n",
    "plt.plot(\n",
    "    plt_vals[\"fast_mode\"][\"x\"],\n",
    "    plt_vals[\"fast_mode\"][\"y\"],\n",
    "    linewidth=2,\n",
    "    label=\"Fast\",\n",
    ")\n",
    "ax = plt.gca()\n",
    "\n",
    "# adjust axes\n",
    "ax.set_xlabel(r\"$(\\omega / k v_A)^2 \\, \\sin \\theta$\", fontsize=fs)\n",
    "ax.set_ylabel(r\"$(\\omega / k v_A)^2 \\, \\cos \\theta$\", fontsize=fs)\n",
    "ax.set_xlim(0.0, 1.5)\n",
    "ax.set_ylim(0.0, 2.0)\n",
    "for spine in ax.spines.values():\n",
    "    spine.set_linewidth(2)\n",
    "ax.minorticks_on()\n",
    "ax.tick_params(which=\"both\", labelsize=fs, width=2)\n",
    "ax.tick_params(which=\"major\", length=10)\n",
    "ax.tick_params(which=\"minor\", length=5)\n",
    "ax.xaxis.set_major_locator(MultipleLocator(0.5))\n",
    "ax.xaxis.set_minor_locator(MultipleLocator(0.1))\n",
    "ax.yaxis.set_major_locator(MultipleLocator(0.5))\n",
    "ax.yaxis.set_minor_locator(MultipleLocator(0.1))\n",
    "\n",
    "\n",
    "# Alfven mode\n",
    "plt.plot(\n",
    "    plt_vals[\"alfven_mode\"][\"x\"],\n",
    "    plt_vals[\"alfven_mode\"][\"y\"],\n",
    "    linewidth=2,\n",
    "    label=\"Alfvén\",\n",
    ")\n",
    "\n",
    "# Acoustic mode\n",
    "plt.plot(\n",
    "    plt_vals[\"acoustic_mode\"][\"x\"],\n",
    "    plt_vals[\"acoustic_mode\"][\"y\"],\n",
    "    linewidth=2,\n",
    "    label=\"Acoustic\",\n",
    ")\n",
    "\n",
    "# annotations\n",
    "plt.legend(fontsize=fs, loc=\"upper right\")"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}