{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FLEKS Python Visualization Toolkit: AMReX Data\n",
    "\n",
    "flekspy is a Python package for processing FLEKS data. This notebook focuses on handling data in the AMReX format, which is used for both field and particle data. We will cover two main ways to load and analyze this data:\n",
    "\n",
    "1.  **Using the `yt`-based loader**: This is the primary method, leveraging the power of the `yt` library for slicing, plotting, and analyzing both field and particle data.\n",
    "2.  **Using the experimental native AMReX particle loader**: A faster, more direct way to load and plot particle data without the `yt` overhead."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup: Imports and Data Downloads"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import flekspy\n",
    "from flekspy.util import download_testfile, unit_one\n",
    "from flekspy.amrex import AMReXParticleData\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import yt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Downloading Demo Data\n",
    "\n",
    "We'll download two different AMReX datasets to demonstrate various features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset 1: General purpose field and particle data\n",
    "url_1 = \"https://raw.githubusercontent.com/henry2004y/batsrus_data/master/3d_particle.tar.gz\"\n",
    "download_testfile(url_1, \"data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset 2: Smaller particle dataset for more specific examples\n",
    "url_2 = \"https://raw.githubusercontent.com/henry2004y/batsrus_data/master/fleks_particle_small.tar.gz\"\n",
    "download_testfile(url_2, \"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Analyzing Field and PIC Particle Data with the `yt` Loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The primary way to load AMReX data is with `flekspy.load`, which uses `yt` on the backend. This returns a `yt` dataset object, giving you access to the full power of the `yt` analysis framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = flekspy.load(\"data/3d*amrex\", use_yt_loader=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Slicing and Plotting Field Data\n",
    "\n",
    "We can easily take a 2D slice of the 3D data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dc = ds.get_slice(\"z\", 0.5)\n",
    "dc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Flekspy provides convenient wrappers for creating plots from these slices."
   ]
  },
  {
   "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": [
    "For more control, you can also extract the data and plot it directly with Matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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\n",
    "\n",
    "We can also analyze the particle data. `yt` allows us to create derived fields, which are useful for weighting particles. Here, we create a `unit_one` field that simply assigns a value of 1 to each particle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.add_field(\n",
    "    ds.pvar(\"unit_one\"),\n",
    "    function=unit_one,\n",
    "    sampling_type=\"particle\",\n",
    "    units=\"dimensionless\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create a phase space plot. We'll select a spatial region and plot the distribution of y- and z-velocities (`p_uy`, `p_uz`), weighted by the particle weight (`p_w`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_field = \"p_uy\"\n",
    "y_field = \"p_uz\"\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",
    "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_zlim(pp.fields[0], zmin=zmin, zmax=zmax)\n",
    "pp.set_xlabel(r\"$V_y$\")\n",
    "pp.set_ylabel(r\"$V_z$\")\n",
    "pp.set_colorbar_label(pp.fields[0], \"weight\")\n",
    "pp.set_log(pp.fields[0], True)\n",
    "pp.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Selecting and Plotting Particles in Geometric Regions\n",
    "\n",
    "`yt` makes it easy to select particles within various geometric shapes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sphere\n",
    "Plot the spatial location and velocity space of particles within a sphere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp = ds.sphere(center=[0, 0, 0], radius=1)\n",
    "\n",
    "# Plot particle locations\n",
    "pp_loc = ds.plot_particles(\"p_x\", \"p_y\", \"p_w\", region=sp, unit_type=\"si\", x_bins=64, y_bins=64)\n",
    "pp_loc.show()\n",
    "\n",
    "# Plot phase space\n",
    "pp_phase = ds.plot_phase(\"p_uy\", \"p_uz\", \"p_w\", region=sp, unit_type=\"si\", x_bins=64, y_bins=64)\n",
    "pp_phase.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Disk\n",
    "Plot the spatial location and velocity space of particles within a disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "disk = ds.disk(center=[0, 0, 0], normal=[0, 0, 1], radius=0.5, height=1.0)\n",
    "\n",
    "# Plot particle locations\n",
    "pp_loc = ds.plot_particles(\"p_x\", \"p_y\", \"p_w\", region=disk, unit_type=\"si\", x_bins=64, y_bins=64)\n",
    "pp_loc.show()\n",
    "\n",
    "# Plot phase space\n",
    "pp_phase = ds.plot_phase(\"p_uy\", \"p_uz\", \"p_w\", region=disk, unit_type=\"si\", x_bins=64, y_bins=64)\n",
    "pp_phase.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced Example: Slicing with Range Limits\n",
    "\n",
    "Now we'll use the second dataset to show a more advanced slicing feature. Inheriting from IDL procedures, you can pass strings to `plot` to limit the plotting range (note: this syntax is experimental)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds2 = flekspy.load(\"data/fleks_particle_small/3d*amrex\")\n",
    "dc2 = ds2.get_slice(\"z\", 0.001)\n",
    "\n",
    "f, axes = dc2.plot(\"Bx>(2.2e5)<(3e5) Ex\", figsize=(12, 6))\n",
    "dc2.add_stream(axes[0], \"Bx\", \"By\", color=\"w\")\n",
    "dc2.add_contour(axes[1], \"Bx\", color=\"k\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced Example: Transforming Velocity Coordinates (WIP)\n",
    "\n",
    "This section demonstrates a work-in-progress for transforming particle velocities into a new coordinate system (e.g., field-aligned coordinates)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# l = [1, 0, 0]\n",
    "# m = [0, 1, 0]\n",
    "# n = [0, 0, 1]\n",
    "\n",
    "# def _vel_l(field, data):\n",
    "#     return (l[0] * data[(\"particles\", \"p_ux\")] + l[1] * data[(\"particles\", \"p_uy\")] + l[2] * data[(\"particles\", \"p_uz\")])\n",
    "\n",
    "# def _vel_m(field, data):\n",
    "#     return (m[0] * data[(\"particles\", \"p_ux\")] + m[1] * data[(\"particles\", \"p_uy\")] + m[2] * data[(\"particles\", \"p_uz\")])\n",
    "\n",
    "# ds2.add_field(ds2.pvar(\"vel_l\"), units=\"code_velocity\", function=_vel_l, sampling_type=\"particle\")\n",
    "# ds2.add_field(ds2.pvar(\"vel_m\"), units=\"code_velocity\", function=_vel_m, sampling_type=\"particle\")\n",
    "\n",
    "# sp2 = ds2.sphere(center=[0, 0, 0], radius=0.001)\n",
    "\n",
    "# x_field = ds2.pvar(\"vel_l\")\n",
    "# y_field = ds2.pvar(\"vel_m\")\n",
    "# z_field = ds2.pvar(\"p_w\")\n",
    "\n",
    "# pp = yt.create_profile(data_source=sp2, bin_fields=[x_field, y_field], fields=z_field, n_bins=[64, 64]).plot()\n",
    "# pp.set_unit(x_field, \"km/s\")\n",
    "# pp.set_unit(y_field, \"km/s\")\n",
    "# pp.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Using the Native AMReX Particle Loader (Experimental)\n",
    "\n",
    "For scenarios where you only need to load particle data and the overhead of `yt` is not desired, `flekspy` provides a direct, experimental loader called `AMReXParticleData`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_file = \"data/3d_particle_region0_1_t00000002_n00000007_amrex\"\n",
    "pd = AMReXParticleData(data_file)\n",
    "pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Particles within a given region can be extracted as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_range = (-0.008, 0.008)\n",
    "y_range = (-0.005, 0.005)\n",
    "\n",
    "rdata = pd.select_particles_in_region(x_range=x_range, y_range=y_range)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Phase space plotting is achieved via `plot_phase`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = pd.plot_phase(\"velocity_y\", \"velocity_z\", bins=(100,32))\n",
    "plt.show()"
   ]
  }
 ],
 "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
}