{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Grids: Uniformly-Spaced Cartesian Grids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Grids are a datastructure that represent one or more physical quantities that share spatial coordinates. For example, the density or magnetic field in a plasma as specified on a Cartesian grid. In addition to storing data, grids have built-in interpolator functions for estimating the values of quantities in between grid vertices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a grid"
   ]
  },
  {
   "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",
    "\n",
    "from plasmapy.plasma import grids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A grid can be created either by providing three arrays of spatial coordinates for vertices (eg. x,yz positions) or using a `np.linspace`-like syntax. For example, the two following methods are equivalent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Method 1\n",
    "xaxis, yaxis, zaxis = [np.linspace(-1 * u.cm, 1 * u.cm, num=20)] * 3\n",
    "x, y, z = np.meshgrid(xaxis, yaxis, zaxis, indexing=\"ij\")\n",
    "grid = grids.CartesianGrid(x, y, z)\n",
    "\n",
    "# Method 2\n",
    "grid = grids.CartesianGrid(\n",
    "    np.array([-1, -1, -1]) * u.cm, np.array([1, 1, 1]) * u.cm, num=(150, 150, 150)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The grid object provides access to a number of properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Is the grid uniformly spaced? {grid.is_uniform}\")\n",
    "print(f\"Grid shape: {grid.shape}\")\n",
    "print(f\"Grid units: {grid.units}\")\n",
    "print(f\"Grid spacing on xaxis: {grid.dax0:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The grid points themselves can be explicitly accessed in one of two forms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y, z = grid.grids\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xyz = grid.grid\n",
    "xyz.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the axes can be accessed similarly. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xaxis = grid.ax0\n",
    "xaxis.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding Quantities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the grid has been initialized, we can add quantities to it that represent physical properties defined on the grid vertices. Each quantity must be a `u.Quantity` array of the same shape as the grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ex = np.random.rand(*grid.shape) * u.V / u.m\n",
    "Ey = np.random.rand(*grid.shape) * u.V / u.m\n",
    "Ez = np.random.rand(*grid.shape) * u.V / u.m\n",
    "Bz = np.random.rand(*grid.shape) * u.T\n",
    "Bz.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When quantities are added to the grid, they are associated with a key string (just like a dictionary). Any key string can be used, but PlasmaPy functions use a shared set of recognized quantities to automatically interperet quantities. Each entry is stored as a namedtuple with fields (\"key\", \"description\", \"unit\"). The full list of recognized quantities can be accessed in the module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in grid.recognized_quantities.keys():\n",
    "    rk = grid.recognized_quantities[key]\n",
    "    key, description, unit = rk.key, rk.description, rk.unit\n",
    "    print(f\"{key} -> {description} ({unit})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quantities can be added to the grid as keyword arguments. The keyword becomes the key string for the quantity in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid.add_quantities(B_z=Bz)\n",
    "grid.add_quantities(E_x=Ex, E_y=Ey, E_z=Ez)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding an unrecognized quantity will lead to a warning, but the quantity will still be added to the grid and can still be accessed by the user later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_quantity = np.random.rand(*grid.shape) * u.T * u.mm\n",
    "grid.add_quantities(int_B=custom_quantity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A summary of the grid, including the currently-defined quantities, can be produced by printing the grid object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A simple list of the defined quantity keys on the grid can also be easily accessed "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(grid.quantities)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A number of methods are built into grid objects, and are illustrated here.\n",
    "\n",
    "The `grid.on_grid` method determines which points in an array are within the bounds of the grid. Since our example grid is a cube spanning from -1 to -1 cm on each axis, the first of the following points is on the grid while the second is not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos = np.array([[0.1, -0.3, 0], [3, 0, 0]]) * u.cm\n",
    "print(grid.on_grid(pos))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, the `grid.vector_intersects` function determines whether the line between two points passes through the grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt0 = np.array([3, 0, 0]) * u.cm\n",
    "pt1 = np.array([-3, 0, 0]) * u.cm\n",
    "pt2 = np.array([3, 10, 0]) * u.cm\n",
    "\n",
    "print(f\"Line from pt0 to pt1 intersects: {grid.vector_intersects(pt0, pt1)}\")\n",
    "print(f\"Line from pt0 to pt2 intersects: {grid.vector_intersects(pt0, pt2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpolating Quantities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Grid objects contain several interpolator methods to evaluate quantites at positions between the grid vertices. These interpolators use the fact that all of the quantities are defined on the same grid to perform faster interpolations using a nearest-neighbor scheme. When an interpolation at a position is requested, the grid indices closest to that position are calculated. Then, the quantity arrays are evaluated at the interpolated indices. Using this method, many quantities can be interpolated at the same positions in almost the same amount of time as is required to interpolate one quantity. \n",
    "\n",
    "Positions are provided to the interpolator as a `u.Quantity` array of shape [N,3] where N is the number of positions and [i,:] represents the x,y,z coordinates of the ith position:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos = np.array([[0.1, -0.3, 0], [0.5, 0.25, 0.8]]) * u.cm\n",
    "print(f\"Pos shape: {pos.shape}\")\n",
    "print(f\"Position 1: {pos[0,:]}\")\n",
    "print(f\"Position 2: {pos[1,:]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simplest interpolator directly returns the nearest-neighbor values for each quantity. Positions that are out-of-bounds return an interpolated value of zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ex_vals = grid.nearest_neighbor_interpolator(pos, \"E_x\")\n",
    "print(f\"Ex at position 1: {Ex_vals[0]:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multiple values can be interpolated at the same time by including additional keys as arguments. In this case, the interpolator returns a tuple of arrays as a result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ex_vals, Ey_vals, Ez_vals = grid.nearest_neighbor_interpolator(pos, \"E_x\", \"E_y\", \"E_z\")\n",
    "print(f\"E at position 1: ({Ex_vals[0]:.2f}, {Ey_vals[0]:.2f}, {Ez_vals[0]:.2f})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a higher-order interpolation, some grids (such as the CartesianGrid subclass) also include a volume-weighted interpolator. This interpolator averages the values on the eight grid vertices surrounding the position (weighted by their distance). The syntax for using this interpolator is the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ex_vals, Ey_vals, Ez_vals = grid.volume_averaged_interpolator(pos, \"E_x\", \"E_y\", \"E_z\")\n",
    "print(f\"E at position 1: ({Ex_vals[0]:.2f}, {Ey_vals[0]:.2f}, {Ez_vals[0]:.2f})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If repeated identical calls are being made to the same interpolator at different positions (for example, in a simulation loop), setting the `persistent` keyword to True will increase performance by not repeatedly re-loading the same quantity arrays from the grid. Setting this keyword for a single interpolation will not improve performace, and is not recommended (and is only done here for illustration)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ex_vals, Ey_vals, Ez_vals = grid.volume_averaged_interpolator(\n",
    "    pos, \"E_x\", \"E_y\", \"E_z\", persistent=True\n",
    ")\n",
    "print(f\"E at position 1: ({Ex_vals[0]:.2f}, {Ey_vals[0]:.2f}, {Ez_vals[0]:.2f})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}