{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FLEKS Python Visualization Toolkit\n",
    "\n",
    "flekspy is a Python package for processing FLEKS data.\n",
    "\n",
    "## FLEKS data format\n",
    "\n",
    "* Field: *.out format or AMREX built-in format, whose directory name is assumed to end with \"_amrex\"\n",
    "* PIC particle: AMREX built-in format \n",
    "* Test particle: binary data format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing the package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import flekspy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Downloading demo data\n",
    "\n",
    "If you don't have FLEKS data to start with, you can download demo field data with the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flekspy.util import download_testfile\n",
    "\n",
    "url = \"https://raw.githubusercontent.com/henry2004y/batsrus_data/master/batsrus_data.tar.gz\"\n",
    "download_testfile(url, \"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example test particle data can be downloaded as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://raw.githubusercontent.com/henry2004y/batsrus_data/master/test_particles.tar.gz\"\n",
    "download_testfile(url, \"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example PIC particle data can be downloaded as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://raw.githubusercontent.com/henry2004y/batsrus_data/master/3d_particle.tar.gz\"\n",
    "download_testfile(url, \"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "`flekspy.load` is the interface to read files of all formats. It returns a different object for different formats."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = \"data/1d__raw_2_t25.60000_n00000258.out\"\n",
    "ds = flekspy.load(file)\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The variables can be accessed via dictionary keys:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = ds.data[\"p\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extracting data\n",
    "\n",
    "Variables at a series of locations can be extracted via `extract_data`, or at a single location via `get_data`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "file = \"data/z=0_fluid_region0_0_t00001640_n00010142.out\"\n",
    "ds = flekspy.load(file)\n",
    "sat = np.array([[-28000.0, 0.0], [9000.0, 0.0]])\n",
    "d = ds.extract_data(sat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc = np.array([0.0, 0.0])\n",
    "d = ds.get_data(loc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing fields\n",
    "\n",
    "We provide convenient `pp` and `pcolormesh` method for IDL 1D and 2D outputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = \"data/1d__raw_2_t25.60000_n00000258.out\"\n",
    "ds = flekspy.load(file)\n",
    "ds.plot(\"p\", \"Bx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = \"data/z=0_fluid_region0_0_t00001640_n00010142.out\"\n",
    "ds = flekspy.load(file)\n",
    "ds.pcolormesh(\"x\", \"y\", \"Bx\", \"By\", \"Bz\", scale=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Native AMReX format can be read using YT:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = flekspy.load(\"data/3d*amrex\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can take a slice at a given location, which works both for IDL and AMReX data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dc = ds.get_slice(\"z\", 0.5)\n",
    "dc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and plot the selected variables on the slice with colored contours, streamlines, and contour lines:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, axes = dc.plot(\"By\", pcolor=True)\n",
    "dc.add_stream(axes[0], \"Bx\", \"By\", color=\"w\")\n",
    "dc.add_contour(axes[0], \"Bx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Users can also obtain the data of a 2D plane, and visualize it with matplotlib. A complete example is given below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "ds = flekspy.load(\"data/3*amrex\")\n",
    "dc = ds.get_slice(\"z\", 0.5)\n",
    "\n",
    "f, axes = plt.subplots(\n",
    "    1, 2, figsize=(12, 4), constrained_layout=True, sharex=True, sharey=True\n",
    ")\n",
    "\n",
    "fields = [\"By\", \"Bz\"]\n",
    "for ivar in range(2):\n",
    "    v = dc.evaluate_expression(fields[ivar])\n",
    "    vmin = v.min().v\n",
    "    vmax = v.max().v\n",
    "    ax = axes[ivar]\n",
    "    ax.set_title(fields[ivar], fontsize=16)\n",
    "    ax.set_ylabel(\"Y\", fontsize=16)\n",
    "    ax.set_xlabel(\"X\", fontsize=16)\n",
    "    c = ax.pcolormesh(dc.x.value, dc.y.value, np.array(v.T), cmap=\"turbo\")\n",
    "    cb = f.colorbar(c, ax=ax, pad=0.01)\n",
    "\n",
    "    ax.set_xlim(np.min(dc.x.value), np.max(dc.x.value))\n",
    "    ax.set_xlim(np.min(dc.y.value), np.max(dc.y.value))\n",
    "    dc.add_stream(ax, \"Bx\", \"By\", density=0.5, color=\"w\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing phase space distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from flekspy.util import unit_one\n",
    "\n",
    "data_file = \"data/3d_*amrex\"\n",
    "ds = flekspy.load(data_file)\n",
    "\n",
    "# Add a user defined field. See yt document for more information about derived field.\n",
    "ds.add_field(\n",
    "    ds.pvar(\"unit_one\"),\n",
    "    function=unit_one,\n",
    "    sampling_type=\"particle\",\n",
    "    units=\"dimensionless\",\n",
    ")\n",
    "\n",
    "x_field = \"p_uy\"\n",
    "y_field = \"p_uz\"\n",
    "# z_field = \"unit_one\"\n",
    "z_field = \"p_w\"\n",
    "xleft = [-0.016, -0.01, 0.0]\n",
    "xright = [0.016, 0.01, 1.0]\n",
    "\n",
    "zmin, zmax = 1e-5, 2e-3\n",
    "\n",
    "## Select and plot the particles inside a box defined by xleft and xright\n",
    "region = ds.box(xleft, xright)\n",
    "pp = ds.plot_phase(\n",
    "    x_field,\n",
    "    y_field,\n",
    "    z_field,\n",
    "    region=region,\n",
    "    unit_type=\"si\",\n",
    "    x_bins=100,\n",
    "    y_bins=32,\n",
    "    domain_size=(xleft[0], xright[0], xleft[1], xright[1]),\n",
    ")\n",
    "\n",
    "pp.set_cmap(pp.fields[0], \"turbo\")\n",
    "pp.set_font(\n",
    "    {\n",
    "        \"size\": 34,\n",
    "        \"family\": \"DejaVu Sans\",\n",
    "    }\n",
    ")\n",
    "\n",
    "pp.set_zlim(pp.fields[0], zmin=zmin, zmax=zmax)\n",
    "\n",
    "pp.set_xlabel(r\"$V_y$\")\n",
    "pp.set_ylabel(r\"$V_z$\")\n",
    "pp.set_colorbar_label(pp.fields[0], \"weight\")\n",
    "str_title = (\n",
    "    rf\"x = [{xleft[0]:.1e}, {xright[0]:.1e}], y = [{xleft[1]:.1e}, {xright[1]:.1e}]\"\n",
    ")\n",
    "pp.set_title(pp.fields[0], str_title)\n",
    "pp.set_log(pp.fields[0], True)\n",
    "pp.show()\n",
    "# pp.save(\"test\")\n",
    "# pp.plots[(\"particles\", z_field)].axes.xaxis.set_major_locator(\n",
    "# ticker.MaxNLocator(4))\n",
    "# pp.plots[(\"particles\", z_field)].axes.yaxis.set_major_locator(\n",
    "# ticker.MaxNLocator(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the location of particles that are inside a sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "center = [0, 0, 0]\n",
    "radius = 1\n",
    "# Object sphere is defined in yt/data_objects/selection_objects/spheroids.py\n",
    "sp = ds.sphere(center, radius)\n",
    "pp = ds.plot_particles(\n",
    "    \"p_x\", \"p_y\", \"p_w\", region=sp, unit_type=\"si\", x_bins=64, y_bins=64\n",
    ")\n",
    "pp.show()\n",
    "# pp.save(\"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the velocity space of particles that are inside a sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp = ds.plot_phase(\n",
    "    \"p_uy\", \"p_uz\", \"p_w\", region=sp, unit_type=\"si\", x_bins=64, y_bins=64\n",
    ")\n",
    "\n",
    "pp.set_cmap(pp.fields[0], \"turbo\")\n",
    "\n",
    "pp.set_font(\n",
    "\n",
    "    {\n",
    "\n",
    "        \"size\": 34,\n",
    "\n",
    "        \"family\": \"DejaVu Sans\",\n",
    "\n",
    "    }\n",
    ")\n",
    "\n",
    "pp.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the location of particles that are inside a disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "center = [0, 0, 0]\n",
    "normal = [0, 0, 1]  # normal direction of the disk\n",
    "radius = 0.5\n",
    "height = 1.0\n",
    "# Object sphere is defined in yt/data_objects/selection_objects/disk.py\n",
    "disk = ds.disk(center, normal, radius, height)\n",
    "pp = ds.plot_particles(\n",
    "    \"p_x\", \"p_y\", \"p_w\", region=disk, unit_type=\"si\", x_bins=64, y_bins=64\n",
    ")\n",
    "pp.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the velocity space of particles that are inside a disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp = ds.plot_phase(\n",
    "    \"p_uy\", \"p_uz\", \"p_w\", region=disk, unit_type=\"si\", x_bins=64, y_bins=64\n",
    ")\n",
    "\n",
    "pp.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the 2D array from the phase plot and reconstruct from scratch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for var_name in pp.profile.field_data:\n",
    "    val = pp.profile.field_data[var_name]\n",
    "\n",
    "x = pp.profile.x\n",
    "y = pp.profile.y\n",
    "\n",
    "plt.pcolormesh(x, y, val)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Test Particle Trajectories"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading particle data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flekspy import FLEKSTP\n",
    "\n",
    "tp = FLEKSTP(\"data/test_particles\", iSpecies=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Obtaining particle IDs\n",
    "\n",
    "Test particle IDs consists of a CPU index and a particle index attached to the CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pIDs = tp.getIDs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reading particle trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj = tp.read_particle_trajectory(pIDs[10])\n",
    "print(\"time:\\n\")\n",
    "print(traj[:, 0])\n",
    "print(\"x:\\n\")\n",
    "print(traj[:, 1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting initial location"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tp.read_initial_location(pIDs[10])\n",
    "print(\"time, X, Y, Z, Ux, Uy, Uz\")\n",
    "print(\n",
    "    f\"{x[0]:.2e}, {x[1]:.2e}, {x[2]:.2e}, {x[3]:.2e}, {x[4]:.2e}, {x[5]:.2e}, {x[6]:.2e}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tp.plot_trajectory(pIDs[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reading and visualizing all particle's location at a given snapshot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ids, pData = tp.read_particles_at_time(0.0, doSave=False)\n",
    "tp.plot_location(pData)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Selecting particles starting in a region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_select(tp, pid):\n",
    "    pData = tp.read_initial_location(pid)\n",
    "    inRegion = pData[FLEKSTP.ix_] > 0 and pData[FLEKSTP.iy_] > 0\n",
    "    return inRegion\n",
    "\n",
    "\n",
    "pSelected = tp.select_particles(f_select)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving trajectories to CSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tp.save_trajectory_to_csv(pIDs[0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.7 64-bit",
   "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.9.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "f22a20af907fde35ff19e1e892fdb271353fb19b11c7ebd774491472e685293c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}