{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 3 - Test particles\n",
    "\n",
    "Let us explore some options of the models [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov), and [GuidingCenter](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.GuidingCenter). In particular, we will\n",
    "\n",
    "1. Change the geometry\n",
    "2. Change the loading of the markers\n",
    "3. Set a static background magnetic field\n",
    "\n",
    "## Particles in a cylinder\n",
    "\n",
    "Similar to Tutorial 2, we shall re-create the launch files in the notebook for the purpose of discussion. The default launch files can be created from the console via\n",
    "\n",
    "```\n",
    "struphy params Vlasov\n",
    "struphy params GuidingCenter \n",
    "```\n",
    "\n",
    "This time, we set the simulation domain $\\Omega$ to be a cylinder. We explore two options for drawing markers in position space:\n",
    "\n",
    "- uniform in logical space $[0, 1]^3 = F^{-1}(\\Omega)$\n",
    "- uniform on the cylinder $\\Omega$\n",
    "\n",
    "We start with importing the API and the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from struphy import BaseUnits, EnvironmentOptions, Time\n",
    "from struphy import equils\n",
    "from struphy import domains\n",
    "from struphy import grids\n",
    "from struphy import DerhamOptions\n",
    "from struphy import maxwellians\n",
    "from struphy import BoundaryParameters, LoadingParameters, WeightsParameters\n",
    "from struphy import Simulation\n",
    "\n",
    "# import model\n",
    "from struphy.models import Vlasov"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "We shall create two simulations of the same model, which we store in different output folders, defined through the environment variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# light-weight model instances\n",
    "model = Vlasov()\n",
    "model_2 = Vlasov()\n",
    "\n",
    "# environment options\n",
    "env = EnvironmentOptions()\n",
    "env_2 = EnvironmentOptions(sim_folder=\"sim_2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "The other generic options will be the same for both simulations. Here, we just perform one time step and load a cylindrical geometry and we can leave the equilibrium empty:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# time stepping\n",
    "time_opts = Time(dt=0.2, Tend=0.2)\n",
    "\n",
    "# geometry\n",
    "a1 = 1e-6 # to avoid polar singularity, we cut a small hole around the axis.\n",
    "a2 = 5.0\n",
    "Lz = 20.0\n",
    "domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz)\n",
    "\n",
    "# fluid equilibrium (can be used as part of initial conditions)\n",
    "equil = None\n",
    "\n",
    "# grid\n",
    "grid = grids.TensorProductGrid()\n",
    "\n",
    "# derham options\n",
    "derham_opts = DerhamOptions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "Let us instantiate the first simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = Simulation(model,\n",
    "                 env=env,\n",
    "                 time_opts=time_opts,\n",
    "                 domain=domain,\n",
    "                 equil=equil,\n",
    "                 grid=grid,\n",
    "                 derham_opts=derham_opts,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "In Struphy there is a convenient method for generating similar simulations with different parameters. \n",
    "We can use it to generate the second simulation with different model and environment: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_2 = sim.spawn_sister(model= model_2, env=env_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Let us have a look at the simulation domain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "For the second simulation, in the parameters for particle loading we choose `spatial=\"disc\"` in order to draw uniformly on the cross section of the cylinder: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# species parameters\n",
    "loading_params = LoadingParameters(Np=1000)\n",
    "loading_params_2 = LoadingParameters(Np=1000, spatial=\"disc\")\n",
    "\n",
    "weights_params = WeightsParameters()\n",
    "boundary_params = BoundaryParameters()\n",
    "\n",
    "model.kinetic_ions.set_markers(\n",
    "    loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params\n",
    ")\n",
    "model_2.kinetic_ions.set_markers(\n",
    "    loading_params=loading_params_2, weights_params=weights_params, boundary_params=boundary_params\n",
    ")\n",
    "\n",
    "model.kinetic_ions.set_sorting_boxes()\n",
    "model_2.kinetic_ions.set_sorting_boxes()\n",
    "\n",
    "model.kinetic_ions.set_save_data(n_markers=1.0)\n",
    "model_2.kinetic_ions.set_save_data(n_markers=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "Propagator options and initial conditions shall be the same in both simulations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# propagator options\n",
    "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n",
    "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n",
    "\n",
    "model_2.propagators.push_vxb.options = model_2.propagators.push_vxb.Options()\n",
    "model_2.propagators.push_eta.options = model_2.propagators.push_eta.Options()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial conditions (background + perturbation)\n",
    "perturbation = None\n",
    "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n",
    "\n",
    "model.kinetic_ions.var.add_background(background)\n",
    "model_2.kinetic_ions.var.add_background(background)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "Let us now run the first simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "And now the second simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_2.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "We now post-process both simulations, load the generated data and plot the initial particle positions on a cross section of the cylinder:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_2.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_2.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "\n",
    "orbits = sim.orbits.kinetic_ions\n",
    "orbits_uni = sim_2.orbits.kinetic_ions\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.scatter(orbits[0, :, 0], orbits[0, :, 1], s=2.0)\n",
    "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n",
    "ax = plt.gca()\n",
    "ax.add_patch(circle1)\n",
    "ax.set_aspect(\"equal\")\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.title(\"sim_1: draw uniform in logical space\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.scatter(orbits_uni[0, :, 0], orbits_uni[0, :, 1], s=2.0)\n",
    "circle2 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n",
    "ax = plt.gca()\n",
    "ax.add_patch(circle2)\n",
    "ax.set_aspect(\"equal\")\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.title(\"sim_2: draw uniform on disc\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {},
   "source": [
    "## Reflecting boundary conditions \n",
    "\n",
    "Let us setup a new simulation and run for 50 time steps and with 15 particles in the cylinder.\n",
    "Moreover, we set reflecting boundary conditions in radial direction, which in Struphy is always the logical direction $\\eta_1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "time_opts = Time(dt=0.2, Tend=10.0)\n",
    "loading_params = LoadingParameters(Np=15, spatial=\"disc\")\n",
    "boundary_params = BoundaryParameters(bc=(\"reflect\", \"periodic\", \"periodic\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# light-weight model instance\n",
    "model = Vlasov()\n",
    "\n",
    "sim = Simulation(model,\n",
    "                 env=env,\n",
    "                 time_opts=time_opts,\n",
    "                 domain=domain,\n",
    "                 equil=equil,\n",
    "                 grid=grid,\n",
    "                 derham_opts=derham_opts,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.kinetic_ions.set_markers(\n",
    "    loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params\n",
    ")\n",
    "model.kinetic_ions.set_sorting_boxes()\n",
    "model.kinetic_ions.set_save_data(n_markers=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31",
   "metadata": {},
   "source": [
    "We still have to set the propagator options and the initial conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# propagator options\n",
    "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n",
    "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n",
    "\n",
    "# initial conditions (background + perturbation)\n",
    "perturbation = None\n",
    "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n",
    "\n",
    "model.kinetic_ions.var.add_background(background)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33",
   "metadata": {},
   "source": [
    "We can now run the simulation, then post-process the data and plot the resulting orbits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.load_plotting_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37",
   "metadata": {},
   "source": [
    "Under `sim.orbits[<species_name>]` one finds a three-dimensional numpy array; the first index refers to the time step, the second index to the particle and the third index to the particle attribute. The first three attributes are the particle positions, followed by the velocities and the (initial and time-dependent) weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "orbits = sim.orbits.kinetic_ions\n",
    "\n",
    "Nt = orbits.shape[0]\n",
    "Np = orbits.shape[1]\n",
    "Nattr = orbits.shape[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.gca()\n",
    "\n",
    "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n",
    "\n",
    "# create alpha for color scaling\n",
    "Tend = time_opts.Tend\n",
    "alpha = np.linspace(1.0, 0.0, Nt + 1)\n",
    "\n",
    "# loop through particles, plot all time steps\n",
    "for i in range(Np):\n",
    "    ax.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], alpha=alpha)\n",
    "\n",
    "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n",
    "\n",
    "ax.add_patch(circle1)\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.set_title(f\"{Nt - 1} time steps (full color at t=0)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40",
   "metadata": {},
   "source": [
    "## Particles in a cylinder with a magnetic field\n",
    "\n",
    "Let us perform as simulation with a magnetic background field. This can be done by setting an object from the `equils` module:\n",
    "\n",
    "<!-- Let $\\Omega \\subset \\mathbb R^3$ be a cylinder as before. Now, we search for trajectories $(\\mathbf x_p, \\mathbf v_p): [0,T] \\to \\Omega \\times \\mathbb R^3$, $p = 0, \\ldots, N-1$ that satisfy\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    " \\dot{\\mathbf x}_p &= \\mathbf v_p\\,,\\qquad && \\mathbf x_p(0) = \\mathbf x_{p0}\\,,\n",
    " \\\\[2mm]\n",
    " \\dot{\\mathbf v}_p &= \\mathbf v_p \\times \\mathbf B_0(\\mathbf x_p) \\qquad && \\mathbf v_p(0) = \\mathbf v_{p0}\\,,\n",
    " \\end{align}\n",
    "$$\n",
    "\n",
    "where $\\mathbf B_0$ is a given magnetic field from an [MHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils.html#mhd-equilibria). \n",
    "In addition to the Propagator [PushEta](https://struphy.pages.mpcdf.de/struphy/sections/subsections/propagators_markers.html#struphy.propagators.propagators_markers.PushEta) for the position update, we shall use [PushVxB](https://struphy.pages.mpcdf.de/struphy/sections/subsections/propagators_markers.html#struphy.propagators.propagators_markers.PushVxB) for the velocity update. -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "B0x = 0.0\n",
    "B0y = 0.0\n",
    "B0z = 1.0\n",
    "equil = equils.HomogenSlab(B0x=B0x, B0y=B0y, B0z=B0z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42",
   "metadata": {},
   "source": [
    "In order to project the equilibrium on the spline basis for fast evaluation in the particle kernels, we need a Derham complex:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43",
   "metadata": {},
   "outputs": [],
   "source": [
    "bcs = ((\"free\", \"free\"), None, None)\n",
    "derham_opts = DerhamOptions(bcs=bcs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44",
   "metadata": {},
   "source": [
    "Now we create the light-weight instance of the model and set the species options. We shall `remove` particles that hit the boundary in $\\eta_1$ (radial) direction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# light-weight model instance\n",
    "model = Vlasov()\n",
    "\n",
    "sim_withB = sim.spawn_sister(model=model,\n",
    "                             equil=equil,\n",
    "                             derham_opts=derham_opts,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "loading_params = LoadingParameters(Np=20)\n",
    "boundary_params = BoundaryParameters(bc=(\"remove\", \"periodic\", \"periodic\"))\n",
    "model.kinetic_ions.set_markers(\n",
    "    loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params\n",
    ")\n",
    "model.kinetic_ions.set_sorting_boxes()\n",
    "model.kinetic_ions.set_save_data(n_markers=1.0)\n",
    "\n",
    "# propagator options\n",
    "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n",
    "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n",
    "\n",
    "# initial conditions (background + perturbation)\n",
    "perturbation = None\n",
    "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n",
    "\n",
    "model.kinetic_ions.var.add_background(background)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47",
   "metadata": {},
   "source": [
    "Now the usual procedure: run, post-process, load data and finally plot the orbits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_withB.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_withB.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_withB.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "orbits = sim_withB.orbits.kinetic_ions\n",
    "\n",
    "Nt = orbits.shape[0]\n",
    "Np = orbits.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.gca()\n",
    "\n",
    "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n",
    "\n",
    "# create alpha for color scaling\n",
    "Tend = time_opts.Tend\n",
    "alpha = np.linspace(1.0, 0.0, Nt + 1)\n",
    "\n",
    "# loop through particles, plot all time steps\n",
    "for i in range(Np):\n",
    "    ax.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], alpha=alpha)\n",
    "\n",
    "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n",
    "\n",
    "ax.add_patch(circle1)\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.set_title(f\"{int(Nt - 1)} time steps (full color at t=0)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53",
   "metadata": {},
   "source": [
    "## Particles in a Tokamak equilibrium\n",
    "\n",
    "Let us try a more complicated `MHDequilibrium`, namely from an ASDEX-Upgrade equilibrium stored in an EQDSK file. We instantiate an `EQDSKequilibrium` with default parameters, except for the density:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {},
   "outputs": [],
   "source": [
    "n1 = 0.0\n",
    "n2 = 0.0\n",
    "na = 1.0\n",
    "equil = equils.EQDSKequilibrium(n1=n1, n2=n2, na=na)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55",
   "metadata": {},
   "source": [
    "Since `EQDSKequilibrium` is an `AxisymmMHDequilibrium`, which is a subclass of `CartesianMHDequilibrium`, we are free to choose any mapping for the simulation (e.g. a Cuboid for Cartesian coordinates). In order to be conforming to the boundary of the equilibrium, we shall choose the `Tokamak` mapping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_elements = (28, 72)\n",
    "degree = (3, 3)\n",
    "psi_power = 0.6\n",
    "psi_shifts = (1e-6, 1.0)\n",
    "domain = domains.Tokamak(equilibrium=equil, num_elements=num_elements, degree=degree, psi_power=psi_power, psi_shifts=psi_shifts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57",
   "metadata": {},
   "source": [
    "The `Tokamak` domain is a `PoloidalSplineTorus`, hence\n",
    "\n",
    "$$\n",
    " \\begin{aligned}\n",
    " x &= R \\cos(\\phi)\\,,\n",
    " \\\\\n",
    " y &= -R \\sin(\\phi)\\,,\n",
    " \\\\\n",
    " z &= Z\\,,\n",
    " \\end{aligned}\n",
    "$$\n",
    "\n",
    "between Cartesian $(x, y, z)$- and Tokamak $(R, Z, \\phi)$-coordinates holds, where $(R, Z)$ spans a poloidal plane. Moreover, the Tokamak coordinates are related to general torus coordinates $(r, \\theta, \\phi)$ via a polar mapping in the poloidal plane:\n",
    "\n",
    "$$\n",
    " \\begin{aligned}\n",
    " R &= R_0 + r \\cos(\\theta)\\,,\n",
    " \\\\\n",
    " Z &= r \\sin(\\theta)\\,,\n",
    " \\\\\n",
    " \\phi &= \\phi\\,.\n",
    " \\end{aligned}\n",
    "$$\n",
    "\n",
    "The torus coordinates are related to Struphy logical coordinates $\\boldsymbol \\eta = (\\eta_1, \\eta_2, \\eta_3) \\in [0, 1]^3$ as \n",
    "\n",
    "$$\n",
    " \\begin{aligned}\n",
    " r &= a_1 + (a_2 - a_1) \\eta_1\\,,\n",
    " \\\\\n",
    " \\theta &= 2\\pi \\eta_2\\,,\n",
    " \\\\\n",
    " \\phi &= 2\\pi \\eta_3\\,,\n",
    " \\end{aligned}\n",
    "$$\n",
    "\n",
    "where $a_2 > a_1 \\geq 0$ are boundaries in the radial $r$-direction.\n",
    "This can be seen for instance in the `HollowTorus` mapping (more complicated angle parametrizations $\\theta(\\eta_1, \\eta_2)$ are also available, but not discussed here).\n",
    "\n",
    "Let us plot the equilibrium magnetic field strength:\n",
    "\n",
    "1. in the poloidal plane at $\\phi = 0$\n",
    "2. in the top view at $z = 0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# logical grid on the unit cube\n",
    "e1 = np.linspace(0.0, 1.0, 101)\n",
    "e2 = np.linspace(0.0, 1.0, 101)\n",
    "e3 = np.linspace(0.0, 1.0, 101)\n",
    "\n",
    "# move away from the singular point r = 0\n",
    "e1[0] += 1e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# logical coordinates of the poloidal plane at phi = 0\n",
    "eta_poloidal = (e1, e2, 0.0)\n",
    "# logical coordinates of the top view at theta = 0\n",
    "eta_topview_1 = (e1, 0.0, e3)\n",
    "# logical coordinates of the top view at theta = pi\n",
    "eta_topview_2 = (e1, 0.5, e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cartesian coordinates (squeezed)\n",
    "x_pol, y_pol, z_pol = domain(*eta_poloidal, squeeze_out=True)\n",
    "x_top1, y_top1, z_top1 = domain(*eta_topview_1, squeeze_out=True)\n",
    "x_top2, y_top2, z_top2 = domain(*eta_topview_2, squeeze_out=True)\n",
    "\n",
    "print(f\"{x_pol.shape = }\")\n",
    "print(f\"{x_top1.shape = }\")\n",
    "print(f\"{x_top2.shape = }\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate two axes\n",
    "fig, axs = plt.subplots(2, 1, figsize=(8, 16))\n",
    "ax = axs[0]\n",
    "ax_top = axs[1]\n",
    "\n",
    "# min/max of field strength\n",
    "equil.domain = domain\n",
    "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n",
    "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n",
    "levels = np.linspace(Bmin, Bmax, 51)\n",
    "\n",
    "# absolute magnetic field at phi = 0\n",
    "im = ax.contourf(x_pol, z_pol, equil.absB0(*eta_poloidal, squeeze_out=True), levels=levels)\n",
    "\n",
    "# absolute magnetic field at Z = 0\n",
    "im_top = ax_top.contourf(x_top1, y_top1, equil.absB0(*eta_topview_1, squeeze_out=True), levels=levels)\n",
    "ax_top.contourf(x_top2, y_top2, equil.absB0(*eta_topview_2, squeeze_out=True), levels=levels)\n",
    "\n",
    "# last closed flux surface, poloidal\n",
    "ax.plot(x_pol[-1], z_pol[-1], color=\"k\")\n",
    "\n",
    "# last closed flux surface, toroidal\n",
    "ax_top.plot(x_top1[-1], y_top1[-1], color=\"k\")\n",
    "ax_top.plot(x_top2[-1], y_top2[-1], color=\"k\")\n",
    "\n",
    "# limiter, poloidal\n",
    "ax.plot(equil.limiter_pts_R, equil.limiter_pts_Z, \"tab:orange\")\n",
    "ax.axis(\"equal\")\n",
    "ax.set_xlabel(\"R\")\n",
    "ax.set_ylabel(\"Z\")\n",
    "ax.set_title(\"abs(B) at $\\phi=0$\")\n",
    "fig.colorbar(im)\n",
    "# limiter, toroidal\n",
    "limiter_Rmax = np.max(equil.limiter_pts_R)\n",
    "limiter_Rmin = np.min(equil.limiter_pts_R)\n",
    "\n",
    "thetas = 2 * np.pi * e2\n",
    "limiter_x_max = limiter_Rmax * np.cos(thetas)\n",
    "limiter_y_max = -limiter_Rmax * np.sin(thetas)\n",
    "limiter_x_min = limiter_Rmin * np.cos(thetas)\n",
    "limiter_y_min = -limiter_Rmin * np.sin(thetas)\n",
    "\n",
    "ax_top.plot(limiter_x_max, limiter_y_max, \"tab:orange\")\n",
    "ax_top.plot(limiter_x_min, limiter_y_min, \"tab:orange\")\n",
    "ax_top.axis(\"equal\")\n",
    "ax_top.set_xlabel(\"x\")\n",
    "ax_top.set_ylabel(\"y\")\n",
    "ax_top.set_title(\"abs(B) at $Z=0$\")\n",
    "fig.colorbar(im_top);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62",
   "metadata": {},
   "source": [
    "We need a Derham complex for the projection of the equilibrium onto the spline basis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_elements = (32, 72, 1)\n",
    "grid = grids.TensorProductGrid(num_elements=num_elements)\n",
    "\n",
    "degree = (3, 3, 1)\n",
    "bcs = ((\"free\", \"free\"), None, None)\n",
    "derham_opts = DerhamOptions(degree=degree, bcs=bcs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64",
   "metadata": {},
   "source": [
    "We aim to simulate 15000 time steps with a second-order splitting algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65",
   "metadata": {},
   "outputs": [],
   "source": [
    "time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66",
   "metadata": {},
   "source": [
    "We now set up a simulation of 4 specific particle orbits in this equilibrium:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# light-weight model instance\n",
    "model = Vlasov()\n",
    "\n",
    "sim_asdex = Simulation(model,\n",
    "                       env=env,\n",
    "                       time_opts=time_opts,\n",
    "                       domain=domain,\n",
    "                       equil=equil,\n",
    "                       grid=grid,\n",
    "                       derham_opts=derham_opts,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial particle positions in phase space\n",
    "initial = (\n",
    "    (0.501, 0.001, 0.001, 0.0, 0.0450, -0.04),  # co-passing particle\n",
    "    (0.511, 0.001, 0.001, 0.0, -0.0450, -0.04),  # counter passing particle\n",
    "    (0.521, 0.001, 0.001, 0.0, 0.0105, -0.04),  # co-trapped particle\n",
    "    (0.531, 0.001, 0.001, 0.0, -0.0155, -0.04),\n",
    ")\n",
    "\n",
    "loading_params = LoadingParameters(Np=4, seed=1608, specific_markers=initial)\n",
    "boundary_params = BoundaryParameters(bc=(\"remove\", \"periodic\", \"periodic\"))\n",
    "model.kinetic_ions.set_markers(\n",
    "    loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params, bufsize=2.0\n",
    ")\n",
    "model.kinetic_ions.set_sorting_boxes()\n",
    "model.kinetic_ions.set_save_data(n_markers=1.0)\n",
    "\n",
    "# propagator options\n",
    "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n",
    "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n",
    "\n",
    "# initial conditions (background + perturbation)\n",
    "perturbation = None\n",
    "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n",
    "\n",
    "model.kinetic_ions.var.add_background(background)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_asdex.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_asdex.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_asdex.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72",
   "metadata": {},
   "outputs": [],
   "source": [
    "orbits = sim_asdex.orbits.kinetic_ions\n",
    "\n",
    "Nt = orbits.shape[0]\n",
    "Np = orbits.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n",
    "\n",
    "dt = time_opts.dt\n",
    "Tend = time_opts.Tend\n",
    "\n",
    "for i in range(Np):\n",
    "    r = np.sqrt(orbits[:, i, 0] ** 2 + orbits[:, i, 1] ** 2)\n",
    "    # poloidal\n",
    "    ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n",
    "    # top view\n",
    "    ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n",
    "\n",
    "ax.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n",
    "ax_top.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74",
   "metadata": {},
   "source": [
    "## Guiding-centers in a Tokamak equilibrium\n",
    "\n",
    "Let us run a similar test for guiding-centers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75",
   "metadata": {},
   "outputs": [],
   "source": [
    "from struphy.models import GuidingCenter\n",
    "\n",
    "# light-weight model instance\n",
    "model = GuidingCenter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76",
   "metadata": {},
   "outputs": [],
   "source": [
    "time_opts = Time(dt=0.1, Tend=100, split_algo=\"Strang\")\n",
    "\n",
    "sim_gc = sim_asdex.spawn_sister(model=model, time_opts=time_opts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial phase space coordinates\n",
    "initial = (\n",
    "    (0.501, 0.001, 0.001, -1.935, 1.72),  # co-passing particle\n",
    "    (0.501, 0.001, 0.001, 1.935, 1.72),  # couner-passing particle\n",
    "    (0.501, 0.001, 0.001, -0.6665, 1.72),  # co-trapped particle\n",
    "    (0.501, 0.001, 0.001, 0.4515, 1.72),\n",
    ")  # counter-trapped particl\n",
    "\n",
    "loading_params = LoadingParameters(Np=4, seed=1608, specific_markers=initial)\n",
    "boundary_params = BoundaryParameters(bc=(\"remove\", \"periodic\", \"periodic\"))\n",
    "model.kinetic_ions.set_markers(\n",
    "    loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params, bufsize=2.0\n",
    ")\n",
    "model.kinetic_ions.set_sorting_boxes()\n",
    "model.kinetic_ions.set_save_data(n_markers=1.0)\n",
    "\n",
    "# propagator options\n",
    "model.propagators.push_bxe.options = model.propagators.push_bxe.Options(tol=1e-5)\n",
    "model.propagators.push_parallel.options = model.propagators.push_parallel.Options(tol=1e-5)\n",
    "\n",
    "# initial conditions (background + perturbation)\n",
    "perturbation = None\n",
    "background = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil)\n",
    "\n",
    "model.kinetic_ions.var.add_background(background)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate two axes\n",
    "fig, axs = plt.subplots(2, 1, figsize=(8, 16))\n",
    "ax = axs[0]\n",
    "ax_top = axs[1]\n",
    "\n",
    "# min/max of field strength\n",
    "equil.domain = domain\n",
    "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n",
    "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n",
    "levels = np.linspace(Bmin, Bmax, 51)\n",
    "\n",
    "# absolute magnetic field at phi = 0\n",
    "im = ax.contourf(x_pol, z_pol, equil.absB0(*eta_poloidal, squeeze_out=True), levels=levels)\n",
    "\n",
    "# absolute magnetic field at Z = 0\n",
    "im_top = ax_top.contourf(x_top1, y_top1, equil.absB0(*eta_topview_1, squeeze_out=True), levels=levels)\n",
    "ax_top.contourf(x_top2, y_top2, equil.absB0(*eta_topview_2, squeeze_out=True), levels=levels)\n",
    "\n",
    "# last closed flux surface, poloidal\n",
    "ax.plot(x_pol[-1], z_pol[-1], color=\"k\")\n",
    "\n",
    "# last closed flux surface, toroidal\n",
    "ax_top.plot(x_top1[-1], y_top1[-1], color=\"k\")\n",
    "ax_top.plot(x_top2[-1], y_top2[-1], color=\"k\")\n",
    "\n",
    "# limiter, poloidal\n",
    "ax.plot(equil.limiter_pts_R, equil.limiter_pts_Z, \"tab:orange\")\n",
    "ax.axis(\"equal\")\n",
    "ax.set_xlabel(\"R\")\n",
    "ax.set_ylabel(\"Z\")\n",
    "ax.set_title(\"abs(B) at $\\phi=0$\")\n",
    "fig.colorbar(im)\n",
    "# limiter, toroidal\n",
    "limiter_Rmax = np.max(equil.limiter_pts_R)\n",
    "limiter_Rmin = np.min(equil.limiter_pts_R)\n",
    "\n",
    "thetas = 2 * np.pi * e2\n",
    "limiter_x_max = limiter_Rmax * np.cos(thetas)\n",
    "limiter_y_max = -limiter_Rmax * np.sin(thetas)\n",
    "limiter_x_min = limiter_Rmin * np.cos(thetas)\n",
    "limiter_y_min = -limiter_Rmin * np.sin(thetas)\n",
    "\n",
    "ax_top.plot(limiter_x_max, limiter_y_max, \"tab:orange\")\n",
    "ax_top.plot(limiter_x_min, limiter_y_min, \"tab:orange\")\n",
    "ax_top.axis(\"equal\")\n",
    "ax_top.set_xlabel(\"x\")\n",
    "ax_top.set_ylabel(\"y\")\n",
    "ax_top.set_title(\"abs(B) at $Z=0$\")\n",
    "fig.colorbar(im_top);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_gc.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_gc.pproc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_gc.load_plotting_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82",
   "metadata": {},
   "outputs": [],
   "source": [
    "orbits = sim_gc.orbits.kinetic_ions\n",
    "\n",
    "Nt = orbits.shape[0]\n",
    "Np = orbits.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "colors = [\"tab:blue\", \"tab:orange\", \"tab:green\", \"tab:red\"]\n",
    "\n",
    "dt = time_opts.dt\n",
    "Tend = time_opts.Tend\n",
    "\n",
    "for i in range(Np):\n",
    "    r = np.sqrt(orbits[:, i, 0] ** 2 + orbits[:, i, 1] ** 2)\n",
    "    # poloidal\n",
    "    ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n",
    "    # top view\n",
    "    ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n",
    "\n",
    "ax.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n",
    "ax_top.set_title(f\"{math.ceil(Tend / dt)} time steps\")\n",
    "fig"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env (3.12.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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}