{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading Vlasiator Data Into VDFpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import os.path\n",
    "\n",
    "file = \"parallel_shock_1d_t100s.vlsv\"\n",
    "\n",
    "if not os.path.isfile(file):\n",
    "    url = f\"https://raw.githubusercontent.com/henry2004y/vlsv_data/master/{file}\"\n",
    "    testfiles = url.rsplit(\"/\", 1)[1]\n",
    "    r = requests.get(url, allow_redirects=True)\n",
    "    open(testfiles, \"wb\").write(r.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vlasiator works in the phase space directly, therefore, there is no concept of particles. Here we are collecting the moments computed from integrals of the phase space and performing the clustering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import vdfpy\n",
    "\n",
    "from pyvlasiator.vlsv import Vlsv\n",
    "import pyvlasiator.plot\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "ds = Vlsv(file)\n",
    "df = vdfpy.collect_moments(file)\n",
    "\n",
    "method = \"GMM\"\n",
    "labels = vdfpy.cluster(df, n_clusters=4, method=method)\n",
    "labels = labels[::20]\n",
    "\n",
    "xrange = np.linspace(ds.coordmin[0]+1, ds.coordmax[0]-1, labels.size)\n",
    "ylocs = np.zeros(labels.size)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 2), layout=\"constrained\")\n",
    "\n",
    "for g in np.unique(labels):\n",
    "    ix = np.where(labels == g)\n",
    "    ax.scatter(xrange[ix], ylocs[ix], label=g, s=30)\n",
    "\n",
    "ax.get_yaxis().set_visible(False)\n",
    "ax.legend(loc=\"upper center\", fancybox=True, shadow=True, ncol=4, fontsize=16)\n",
    "ax.set_title(\"1D Shock, 4-class GMM\", fontsize=14)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "from cycler import cycler\n",
    "custom_cycler = (cycler(color=[\"tab:purple\",\"tab:cyan\",\"tab:olive\",\"tomato\"]))\n",
    "\n",
    "# Create a figure with two subfigures\n",
    "fig = plt.figure(figsize=(12, 7), layout=\"constrained\")\n",
    "subfigs = fig.subfigures(2, 1, wspace=0.02, height_ratios=[1, 1.8])\n",
    "\n",
    "axsTop = subfigs[0].subplots(2, 1, sharex=True, height_ratios=[1, 3.5])\n",
    "axsTop[0].set_prop_cycle(custom_cycler)\n",
    "xrange = np.linspace(ds.coordmin[0]*1e-3+1, ds.coordmax[0]*1e-3-1, labels.size)\n",
    "\n",
    "for g in np.unique(labels):\n",
    "    ix = np.where(labels == g)\n",
    "    yloc = np.zeros_like(ix)\n",
    "    axsTop[0].scatter(xrange[ix], yloc, label=g, s=30)\n",
    "\n",
    "axsTop[0].legend(loc=(0.4, 0.0), fancybox=True, shadow=True, ncol=4, fontsize=8)\n",
    "axsTop[0].set_title(\"1D Shock, 4-class GMM\", fontsize=14)\n",
    "axsTop[0].axis(\"off\")\n",
    "\n",
    "x = np.linspace(ds.coordmin[0] * 1e-3, ds.coordmax[0] * 1e-3, ds.ncells[0])\n",
    "w = ds.read_variable(\"proton/vg_rho\") * 1e-6\n",
    "axsTop[1].plot(x, w)\n",
    "axsTop[1].set_xlabel(\"x [km]\", fontsize=14)\n",
    "axsTop[1].set_ylabel(r\"$n_p$ [/cc]\", fontsize=14)\n",
    "\n",
    "nVDF = 5\n",
    "vdfrange = [[-2000, 2000], [-2000, 2000]]\n",
    "xrange = np.linspace(ds.coordmin[0]+1, ds.coordmax[0]-1, nVDF)\n",
    "cmap = mpl.colormaps[\"turbo\"] # pyvlasiator 0.1.4\n",
    "\n",
    "axsBot = subfigs[1].subplots(2, nVDF, sharex=True, sharey=True)\n",
    "for i in range(nVDF):\n",
    "    ax1 = axsBot[0,i]\n",
    "    ax2 = axsBot[1,i]\n",
    "    loc = (xrange[i], 0, 0)\n",
    "    ds.vdfslice(loc, ax=ax1, range=vdfrange, addcolorbar=False, cmap=cmap)\n",
    "    ds.vdfslice(loc, ax=ax2, range=vdfrange, addcolorbar=False, slicetype=\"xy\", cmap=cmap)\n",
    "\n",
    "    cidReq = ds.getcell(loc)\n",
    "    cidNearest = ds.getnearestcellwithvdf(cidReq)\n",
    "    l = ds.getcellcoordinates(cidNearest) * 1e-3\n",
    "    ax1.set_xlim(vdfrange[0])\n",
    "    ax1.set_ylim(vdfrange[1])\n",
    "    ax2.set_xlim(vdfrange[0])\n",
    "    ax2.set_ylim(vdfrange[1])\n",
    "    ax1.set_title(f\"x={l[0]:.1e} km\")\n",
    "    ax2.set_title(\"\")\n",
    "    axsTop[1].axvline(l[0], color=\"k\", linestyle=\"--\")\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vdfpy-L8WReCYd-py3.11",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}